# -*- coding: utf-8 -*-
"""
envassume
"""

from __future__ import print_function
import sys
import os
from socket import gethostname
import boto3


class EnvAssumeException(Exception):
    pass


def parse_arguments(arg_list):
    role_arn = os.environ.get('AWS_ASSUME_ROLE')
    if role_arn:
        if len(arg_list) <= 1:
            raise EnvAssumeException('No arguments supplied')

        arg_list = arg_list[1:]

    else:
        if len(arg_list) <= 2:
            raise EnvAssumeException('Not enough arguments')

        role_arn = arg_list[1]
        arg_list = arg_list[2:]

    if arg_list[0] == '--':
        arg_list = arg_list[1:]

    return role_arn, arg_list


def assume_role(role_arn, external_id = None, session_name = None):
    if not external_id:
        external_id = os.environ.get('AWS_EXTERNAL_ID')

    if not session_name:
        session_name = 'envassume-' + gethostname()

    request = {
        'RoleArn': role_arn,
        'RoleSessionName': session_name
    }

    if external_id:
        request['ExternalId'] = external_id

    boto3_session = boto3.Session()
    sts_client = boto3_session.client('sts')
    response = sts_client.assume_role(**request)

    return response.get('Credentials')


def update_env(credentials_lookup):
    credentials_list = (
        ('AccessKeyId', 'AWS_ACCESS_KEY_ID'),
        ('SecretAccessKey', 'AWS_SECRET_ACCESS_KEY'),
        ('SessionToken', 'AWS_SESSION_TOKEN'),
    )
    for credential_key, env_var_name in credentials_list:
        os.environ[env_var_name] = credentials_lookup[credential_key]


def run_command(arg_list):
    try:
        os.execvpe(arg_list[0], arg_list, os.environ)

    except OSError as e:
        raise EnvAssumeException('Unable to run command {}: {}'.format(arg_list[0], e))

    raise EnvAssumeException('Unable to run command {}'.format(arg_list[0]))


def run():
    role_arn, arg_list = parse_arguments(sys.argv)

    credentials_lookup = assume_role(role_arn)

    update_env(credentials_lookup)

    run_command(arg_list)


def run_script():
    try:
        run()

    except Exception as e:
        print('Error({}) {}'.format(e.__class__.__name__, e), file = sys.stderr)
        sys.exit(1)

    except KeyboardInterrupt:
        pass

@@ -194,20 +194,18 @@
 %0A%0A%0Adef p
-arse
+op
 _argumen
@@ -205,17 +205,16 @@
 argument
-s
 (arg_lis
@@ -225,68 +225,11 @@
     
-role_arn = os.environ.get('AWS_ASSUME_ROLE')%0A    if role_arn
+try
 :%0A  
@@ -230,39 +230,39 @@
 ry:%0A        
-if len(
+return 
 arg_list
 ) %3C= 1:%0A    
@@ -253,20 +253,40 @@
 list
-) %3C= 1:%0A    
+.pop(0)%0A%0A    except IndexError:%0A
     
@@ -317,16 +317,24 @@
 tion('No
+t enough
  argumen
@@ -339,78 +339,170 @@
 ents
- supplied')%0A%0A        arg_list = arg_list%5B1:%5D%0A%0A    else:%0A        if l
+')%0A%0A%0Adef parse_arguments(arg_list):%0A    _ = pop_argument(arg_list)%0A%0A    external_id = None%0A    if arg_list%5B0%5D in ('-i', '--external-id'):%0A        _ = pop_argum
 en
+t
 (arg
@@ -511,14 +511,8 @@
 ist)
- %3C= 2:
 %0A   
@@ -520,120 +520,156 @@
     
-    raise EnvAssumeException('Not enough arguments')%0A%0A        role_arn = arg_list%5B1%5D%0A        arg_list = 
+external_id = pop_argument(arg_list)%0A%0A    role_arn = os.environ.get('AWS_ASSUME_ROLE')%0A    if not role_arn:%0A        role_arn = pop_argument(
 arg_list
 %5B2:%5D
@@ -668,12 +668,9 @@
 list
-%5B2:%5D
+)
 %0A%0A  
@@ -707,27 +707,33 @@
     
-arg_list = 
+_ = pop_argument(
 arg_list
 %5B1:%5D
@@ -732,12 +732,9 @@
 list
-%5B1:%5D
+)
 %0A%0A  
@@ -743,32 +743,45 @@
 return role_arn,
+ external_id,
  arg_list%0A%0A%0Adef 
@@ -1329,16 +1329,22 @@
 ntials')
+ or %7B%7D
 %0A%0A%0Adef u
@@ -1990,16 +1990,29 @@
 ole_arn,
+ external_id,
  arg_lis
@@ -2087,16 +2087,29 @@
 role_arn
+, external_id
 )%0A%0A    u

Tidy up arg parsing, add external_id option.