# -*- coding: utf-8 -*-

import os
import boto.sts
import shlex
from subprocess import Popen

from aws_role_credentials.models import SamlAssertion, AwsCredentialsFile


class Actions:

    def __init__(self,
                 credentials_filename,
                 profile,
                 region,
                 quiet,
                 **kwargs):

        self.credentials_filename = credentials_filename
        self.profile = profile
        self.region = region
        self.quiet = quiet

    @staticmethod
    def print_credentials(credentials_filename, profile, credentials):
        print('\n\n----------------------------------------------------------------')
        print('Your credentials have been stored in the AWS configuration file {0} under the {1} profile.'.format(credentials_filename, profile))
        print('Note that they will expire at {0}.'.format(credentials.expiration))
        print('You may safely rerun this script at any time to refresh your credentials.')
        print('To use this credential, call the AWS CLI with the --profile option (e.g. aws --profile {0} ec2 describe-instances).'.format(profile))
        print('----------------------------------------------------------------\n\n')

    @staticmethod
    def persist_credentials(credentials_filename,
                            profile, region, token, quiet, **kwargs):
        AwsCredentialsFile(credentials_filename).add_profile(profile,
                                                             region,
                                                             token.credentials)
        if not quiet:
            Actions.print_credentials(credentials_filename,
                                      profile,
                                      token.credentials)

    @staticmethod
    def credentials_handler(credentials_filename,
                            profile,
                            region, quiet, **kwargs):

        return lambda token: Actions.persist_credentials(credentials_filename,
                                                         profile,
                                                         region,
                                                         token,
                                                         quiet)

    @staticmethod
    def exec_with_credentials(region, command, token):
        env = os.environ.copy()

        env["AWS_ACCESS_KEY_ID"] = token.credentials.access_key
        env["AWS_SECRET_ACCESS_KEY"] = token.credentials.secret_key
        env["AWS_SESSION_TOKEN"] = token.credentials.session_token
        env["AWS_DEFAULT_REGION"] = region

        Popen(shlex.split(command),
              env=env,
              shell=False).wait()

    @staticmethod
    def exec_handler(region, exec_command, **kwargs):
        return lambda token: Actions.exec_with_credentials(region, exec_command, token)

    @staticmethod
    def saml_token(region, assertion, **kwargs):
        assertion = SamlAssertion(assertion)
        role = assertion.roles()[0]

        conn = boto.sts.connect_to_region(region, anon=True)
        return conn.assume_role_with_saml(role['role'], role['principle'],
                                          assertion.encode())

    @staticmethod
    def user_token(region, role_arn, session_name,
                   mfa_serial_number=None, mfa_token=None,
                   **kwargs):
        conn = boto.sts.connect_to_region(region)

        return conn.assume_role(role_arn, session_name,
                                mfa_serial_number=mfa_serial_number,
                                mfa_token=mfa_token)

@@ -25,17 +25,23 @@
 %0Aimport 
-o
+boto.st
 s%0Aimport
@@ -33,39 +33,33 @@
 boto.sts%0Aimport 
-boto.st
+o
 s%0Aimport shlex%0Af
@@ -57,16 +57,44 @@
 t shlex%0A
+from getpass import getpass%0A
 from sub
@@ -3020,16 +3020,17 @@
     role
+s
  = asser
@@ -3041,16 +3041,704 @@
 .roles()
+%0A        if len(roles) %3E 1:%0A            print('Please select the role you would like to assume:')%0A            for i, role in enumerate(roles):%0A                print('%5B%7B%7D%5D - %7B%7D'.format(i, role%5B'role'%5D))%0A            while True:%0A                # We use getpass() instead of input() here because we are already listening for stdin as part of%0A                # read_stdin()%0A                selectedroleindex = getpass('Selection: ')%0A                try:%0A                    role = roles%5Bint(selectedroleindex)%5D%0A                    break%0A                except (IndexError, ValueError):%0A                    print('Invalid selection, please try again...')%0A        else:%0A            role = roles
 %5B0%5D%0A%0A   

Allow for role selection Rather than assume the first role, this will prompt the user to choose a role that they would like to assume.