import argparse

from awscli.customizations.commands import BasicCommand

import os

def awscli_initialize(event_hooks):
    event_hooks.register('building-command-table.main', inject_commands)

def inject_commands(command_table, session, **kwargs):
    command_table['roleshell'] = RoleShell(session)

def get_exec_args(input_command):
    if len(input_command) == 0:
        input_command = (os.environ['SHELL'],)

    return (input_command[0], input_command)


class RoleShell(BasicCommand):
    NAME = 'roleshell'
    DESCRIPTION = ('Executes a shell with temporary AWS credentials provided as environment variables')
    ARG_TABLE = [
        dict(name='command', nargs=argparse.REMAINDER, positional_arg=True),
    ]

    def _run_main(self, args, parsed_globals):
        a = self._session.get_credentials()
        os.environ['AWS_ACCESS_KEY_ID'] = a.access_key
        os.environ['AWS_SECRET_ACCESS_KEY'] = a.secret_key
        os.environ['AWS_SESSION_TOKEN'] = a.token

        os.execvp(*get_exec_args(args.command))

@@ -8,16 +8,55 @@
 argparse
+%0Aimport os%0Aimport shlex%0Aimport textwrap
 %0A%0Afrom a
@@ -106,26 +106,16 @@
 ommand%0A%0A
-import os%0A
 %0Adef aws
@@ -217,16 +217,17 @@
 mands)%0A%0A
+%0A
 def inje
@@ -330,165 +330,729 @@
 n)%0A%0A
-def get_exec_args(input_command):%0A    if len(input_command) == 0:%0A        input_command = (os.environ%5B'SHELL'%5D,)%0A%0A    return (input_command%5B0%5D, input_command
+%0Adef print_creds(creds):%0A    quoted_vars = map(shlex.quote, (creds.access_key,%0A                                    creds.secret_key, creds.token))%0A%0A    print(textwrap.dedent(%22%22%22%5C%0A        export AWS_ACCESS_KEY_ID=%7B%7D%0A        export AWS_SECRET_ACCESS_KEY=%7B%7D%0A        export AWS_SESSION_TOKEN=%7B%7D%5C%0A    %22%22%22.format(*quoted_vars)))%0A%0A%0Adef get_exec_args(input_command):%0A    if len(input_command) == 0:%0A        input_command = (os.environ%5B'SHELL'%5D,)%0A%0A    return (input_command%5B0%5D, input_command)%0A%0A%0Adef run_command(creds, command):%0A    os.environ%5B'AWS_ACCESS_KEY_ID'%5D = creds.access_key%0A    os.environ%5B'AWS_SECRET_ACCESS_KEY'%5D = creds.secret_key%0A    os.environ%5B'AWS_SESSION_TOKEN'%5D = creds.token%0A%0A    os.execvp(*get_exec_args(command)
 )%0A%0A%0A
@@ -1124,16 +1124,25 @@
 TION = (
+%0A        
 'Execute
@@ -1375,17 +1375,17 @@
         
-a
+c
  = self.
@@ -1411,24 +1411,25 @@
 tials()%0A
+%0A
         
 os.envir
@@ -1424,198 +1424,103 @@
     
-os.environ%5B'AWS_ACCESS_KEY_ID'%5D = a.access_key%0A        os.environ%5B'AWS_SECRET_ACCESS_KEY'%5D = a.secret_key%0A        os.environ%5B'AWS_SESSION_TOKEN'%5D = a.token%0A%0A        os.execvp(*get_exec_args(
+if len(args.command) == 0:%0A            print_creds(c)%0A        else:%0A            run_command(c, 
 args
@@ -1520,18 +1520,17 @@
 c, args.command)
-)
 %0A

Make it possible to eval() output instead of execing another shell.