Resolve AWS account aliases if possible and print them to the users

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())))

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:

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.")