Merge pull request #39 from s-maj/master

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

aws_google_auth/__init__.py

@@ -27,6 +27,7 @@ def parse_args(args):
     parser.add_argument('-p', '--profile', help='AWS profile (defaults to value of $AWS_PROFILE, then falls back to \'sts\')')
     parser.add_argument('-D', '--disable-u2f', action='store_true', help='Disable U2F functionality.')
     parser.add_argument('--no-cache', dest="saml_cache", action='store_false', help='Do not cache the SAML Assertion.')
+    parser.add_argument('--resolve-aliases', action='store_true', help='Resolve AWS account aliases.')
 
     role_group = parser.add_mutually_exclusive_group()
     role_group.add_argument('-a', '--ask-role', action='store_true', help='Set true to always pick the role')
@@ -119,6 +120,12 @@ def cli(cli_args):
         os.getenv('U2F_DISABLED'),
         config.u2f_disabled)
 
+    # Resolve AWS aliases enabled (Option priority = ARGS, ENV_VAR, DEFAULT)
+    config.resolve_aliases = coalesce(
+        args.resolve_aliases,
+        os.getenv('RESOLVE_AWS_ALIASES'),
+        config.resolve_aliases)
+
     # Username (Option priority = ARGS, ENV_VAR, DEFAULT)
     config.username = coalesce(
         args.username,
@@ -167,7 +174,11 @@ 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)
+        if config.resolve_aliases:
+            aliases = amazon_client.resolve_aws_aliases(roles)
+            config.role_arn, config.provider = util.Util.pick_a_role(roles, aliases)
+        else:
+            config.role_arn, config.provider = util.Util.pick_a_role(roles)
 
     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,42 @@ 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:
+                sts = boto3.client('sts',
+                                   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)
+
+                account_id = sts.get_caller_identity().get('Account')
+                aws_dict[role.split(':')[4]] = '{}'.format(account_id)
+
+        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/configuration.py

@@ -25,6 +25,7 @@ def __init__(self, **kwargs):
         self.__saml_cache = None
         self.sp_id = None
         self.u2f_disabled = False
+        self.resolve_aliases = False
         self.username = None
 
     # For the "~/.aws/config" file, we use the format "[profile testing]"

aws_google_auth/util.py

@@ -1,6 +1,8 @@
 #!/usr/bin/env python
 
 import os
+from collections import OrderedDict
+from tabulate import tabulate
 
 
 class Util:
@@ -13,18 +15,46 @@ def get_input(prompt):
             return input(prompt)
 
     @staticmethod
-    def pick_a_role(roles):
-        while True:
-            for i, role in enumerate(roles):
-                print("[{:>3d}] {}".format(i + 1, role))
+    def pick_a_role(roles, aliases=None):
+        if aliases:
+            enriched_roles = {}
+            for role, principal in roles.items():
+                enriched_roles[role] = [
+                    aliases[role.split(':')[4]],
+                    role.split('role/')[1],
+                    principal
+                ]
+            enriched_roles = OrderedDict(sorted(enriched_roles.items(), key=lambda t: (t[1][0], t[1][1])))
 
-            prompt = 'Type the number (1 - {:d}) of the role to assume: '.format(len(roles))
-            choice = Util.get_input(prompt)
+            ordered_roles = OrderedDict()
+            for role, role_property in enriched_roles.items():
+                ordered_roles[role] = role_property[2]
 
-            try:
-                return list(roles.items())[int(choice) - 1]
-            except IndexError:
-                print("Invalid choice, try again.")
+            enriched_roles_tab = []
+            for i, (role, role_property) in enumerate(enriched_roles.items()):
+                enriched_roles_tab.append([i + 1, role_property[0], role_property[1]])
+
+            while True:
+                print(tabulate(enriched_roles_tab, headers=['No', 'AWS account', 'Role'], ))
+                prompt = 'Type the number (1 - {:d}) of the role to assume: '.format(len(enriched_roles))
+                choice = Util.get_input(prompt)
+
+                try:
+                    return list(ordered_roles.items())[int(choice) - 1]
+                except IndexError:
+                    print("Invalid choice, try again.")
+        else:
+            while True:
+                for i, role in enumerate(roles):
+                    print("[{:>3d}] {}".format(i + 1, role))
+
+                prompt = 'Type the number (1 - {:d}) of the role to assume: '.format(len(roles))
+                choice = Util.get_input(prompt)
+
+                try:
+                    return list(roles.items())[int(choice) - 1]
+                except IndexError:
+                    print("Invalid choice, try again.")
 
     @staticmethod
     def touch(file_name, mode=0o600):

requirements.txt

@@ -3,3 +3,4 @@ boto3
 configparser
 lxml
 requests
+tabulate

setup.py

@@ -85,7 +85,7 @@
     # https://packaging.python.org/en/latest/requirements.html
     # install_requires=['peppercorn'],
     install_requires=['boto3', 'lxml', 'requests', 'beautifulsoup4',
-        'configparser', 'tzlocal'],
+        'configparser', 'tzlocal', 'tabulate'],
 
     # List additional groups of dependencies here (e.g. development
     # dependencies). You can install these using the following syntax,