diff --git a/aws_google_auth/__init__.py b/aws_google_auth/__init__.py index 318ec92..9ec2bdb 100644 --- a/aws_google_auth/__init__.py +++ b/aws_google_auth/__init__.py @@ -167,7 +167,8 @@ def cli(cli_args): if config.role_arn in roles and not config.ask_role: config.provider = roles[config.role_arn] else: - config.role_arn, config.provider = util.Util.pick_a_role(roles) + aliases = amazon_client.resolve_aws_aliases(roles) + config.role_arn, config.provider = util.Util.pick_a_role(roles, aliases) print("Assuming " + config.role_arn) print("Credentials Expiration: " + format(amazon_client.expiration.astimezone(get_localzone()))) diff --git a/aws_google_auth/amazon.py b/aws_google_auth/amazon.py index 270648a..11bf680 100644 --- a/aws_google_auth/amazon.py +++ b/aws_google_auth/amazon.py @@ -1,9 +1,11 @@ #!/usr/bin/env python -import boto3 import base64 -from lxml import etree from datetime import datetime +from threading import Thread + +import boto3 +from lxml import etree class Amazon: @@ -68,6 +70,35 @@ def roles(self): roles[res[0]] = res[1] return roles + def resolve_aws_aliases(self, roles): + def resolve_aws_alias(role, principal, aws_dict): + saml = self.sts_client.assume_role_with_saml(RoleArn=role, + PrincipalArn=principal, + SAMLAssertion=self.base64_encoded_saml) + iam = boto3.client('iam', + aws_access_key_id=saml['Credentials']['AccessKeyId'], + aws_secret_access_key=saml['Credentials']['SecretAccessKey'], + aws_session_token=saml['Credentials']['SessionToken'], + region_name=self.config.region) + try: + response = iam.list_account_aliases() + account_alias = response['AccountAliases'][0] + aws_dict[role.split(':')[4]] = account_alias + except: + aws_dict[role.split(':')[4]] = "AliasNotAvailable" + + threads = [] + aws_id_alias = {} + for number, (role, principal) in enumerate(roles.items()): + t = Thread(target=resolve_aws_alias, args=(role, principal, aws_id_alias)) + t.start() + threads.append(t) + + for t in threads: + t.join() + + return aws_id_alias + @staticmethod def is_valid_saml_assertion(saml_xml): if saml_xml is None: diff --git a/aws_google_auth/util.py b/aws_google_auth/util.py index f80b24a..c341550 100644 --- a/aws_google_auth/util.py +++ b/aws_google_auth/util.py @@ -1,6 +1,7 @@ #!/usr/bin/env python import os +from collections import OrderedDict class Util: @@ -13,16 +14,25 @@ def get_input(prompt): return input(prompt) @staticmethod - def pick_a_role(roles): + def pick_a_role(roles, aliases): + enriched_roles = {} + for role, principal in roles.items(): + enriched_roles['{} {}'.format(aliases[role.split(':')[4]], role)] = principal + enriched_roles = OrderedDict(sorted(enriched_roles.items(), key=lambda t: t[0])) + + ordered_roles = OrderedDict() + for role, principal in enriched_roles.items(): + ordered_roles[role.split(' ')[1]] = principal + while True: - for i, role in enumerate(roles): + for i, role in enumerate(enriched_roles): print("[{:>3d}] {}".format(i + 1, role)) - prompt = 'Type the number (1 - {:d}) of the role to assume: '.format(len(roles)) + prompt = 'Type the number (1 - {:d}) of the role to assume: '.format(len(enriched_roles)) choice = Util.get_input(prompt) try: - return list(roles.items())[int(choice) - 1] + return list(ordered_roles.items())[int(choice) - 1] except IndexError: print("Invalid choice, try again.")