From 3f911d31195491afd5e0167844a030f89f8b225b Mon Sep 17 00:00:00 2001 From: Adam Vondersaar Date: Wed, 27 Nov 2019 14:00:09 -0800 Subject: [PATCH 1/2] Add ability to filter for an account on the CLI with flag -A --- aws_google_auth/__init__.py | 14 ++++++++++- aws_google_auth/configuration.py | 8 +++++++ aws_google_auth/tests/test_args_parser.py | 10 +++++++- .../tests/test_backwards_compatibility.py | 1 + aws_google_auth/tests/test_config_parser.py | 24 +++++++++++++++++++ .../tests/test_configuration_persistence.py | 1 + aws_google_auth/tests/test_init.py | 12 ++++++++-- aws_google_auth/util.py | 15 ++++++++---- 8 files changed, 76 insertions(+), 9 deletions(-) diff --git a/aws_google_auth/__init__.py b/aws_google_auth/__init__.py index 71a1dfb..35d0c0c 100644 --- a/aws_google_auth/__init__.py +++ b/aws_google_auth/__init__.py @@ -29,6 +29,7 @@ def parse_args(args): parser.add_argument('-R', '--region', help='AWS region endpoint ($AWS_DEFAULT_REGION)') parser.add_argument('-d', '--duration', type=int, help='Credential duration ($DURATION)') parser.add_argument('-p', '--profile', help='AWS profile (defaults to value of $AWS_PROFILE, then falls back to \'sts\')') + parser.add_argument('-A', '--account', help='Filter for specific AWS account.') parser.add_argument('-D', '--disable-u2f', action='store_true', help='Disable U2F functionality.') parser.add_argument('-q', '--quiet', action='store_true', help='Quiet output') parser.add_argument('--bg-response', help='Override default bgresponse challenge token.') @@ -153,6 +154,12 @@ def resolve_config(args): os.getenv('GOOGLE_USERNAME'), config.username) + # Account (Option priority = ARGS, ENV_VAR, DEFAULT) + config.account = coalesce( + args.account, + os.getenv('AWS_ACCOUNT'), + config.account) + config.keyring = coalesce( args.keyring, config.keyring) @@ -239,7 +246,12 @@ def process_auth(args, config): if config.role_arn in roles and not config.ask_role: config.provider = roles[config.role_arn] else: - if config.resolve_aliases: + if config.account and config.resolve_aliases: + aliases = amazon_client.resolve_aws_aliases(roles) + config.role_arn, config.provider = util.Util.pick_a_role(roles, aliases, config.account) + elif config.account: + config.role_arn, config.provider = util.Util.pick_a_role(roles, account=config.account) + elif config.resolve_aliases: aliases = amazon_client.resolve_aws_aliases(roles) config.role_arn, config.provider = util.Util.pick_a_role(roles, aliases) else: diff --git a/aws_google_auth/configuration.py b/aws_google_auth/configuration.py index f1643b2..3eef927 100644 --- a/aws_google_auth/configuration.py +++ b/aws_google_auth/configuration.py @@ -36,6 +36,7 @@ def __init__(self, **kwargs): self.print_creds = False self.quiet = False self.bg_response = None + self.account = "" # For the "~/.aws/config" file, we use the format "[profile testing]" # for the 'testing' profile. The credential file will just be "[testing]" @@ -137,6 +138,9 @@ def raise_if_invalid(self): # quiet assert (self.quiet.__class__ is bool), "Expected quiet to be a boolean. Got {}.".format(self.quiet.__class__) + # account + assert (self.account.__class__ is str), "Expected account to be string. Got {}".format(self.account.__class__) + # Write the configuration (and credentials) out to disk. This allows for # regular AWS tooling (aws cli and boto) to use the credentials in the # profile the user specified. @@ -241,6 +245,10 @@ def read(self, profile): read_bg_response = unicode_to_string(config_parser[profile_string].get('google_config.bg_response', None)) self.bg_response = coalesce(read_bg_response, self.bg_response) + # Account + read_account= unicode_to_string(config_parser[profile_string].get('account', None)) + self.account = coalesce(read_account, self.account) + # SAML Cache try: with open(self.saml_cache_file, 'r') as f: diff --git a/aws_google_auth/tests/test_args_parser.py b/aws_google_auth/tests/test_args_parser.py index d083453..59837b0 100644 --- a/aws_google_auth/tests/test_args_parser.py +++ b/aws_google_auth/tests/test_args_parser.py @@ -31,12 +31,14 @@ def test_no_arguments(self): self.assertEqual(parser.username, None) self.assertEqual(parser.quiet, False) self.assertEqual(parser.bg_response, None) + self.assertEqual(parser.account, None) + self.assertFalse(parser.save_failure_html) # Assert the size of the parameter so that new parameters trigger a review of this function # and the appropriate defaults are added here to track backwards compatibility in the future. - self.assertEqual(len(vars(parser)), 17) + self.assertEqual(len(vars(parser)), 18) def test_username(self): @@ -53,6 +55,8 @@ def test_username(self): self.assertEqual(parser.region, None) self.assertEqual(parser.role_arn, None) self.assertEqual(parser.username, 'username@gmail.com') + self.assertEqual(parser.account, None) + def test_nocache(self): @@ -69,6 +73,8 @@ def test_nocache(self): self.assertEqual(parser.region, None) self.assertEqual(parser.role_arn, None) self.assertEqual(parser.username, None) + self.assertEqual(parser.account, None) + def test_resolvealiases(self): @@ -85,6 +91,8 @@ def test_resolvealiases(self): self.assertEqual(parser.region, None) self.assertEqual(parser.role_arn, None) self.assertEqual(parser.username, None) + self.assertEqual(parser.account, None) + def test_ask_and_supply_role(self): diff --git a/aws_google_auth/tests/test_backwards_compatibility.py b/aws_google_auth/tests/test_backwards_compatibility.py index b3ff051..f8962f6 100644 --- a/aws_google_auth/tests/test_backwards_compatibility.py +++ b/aws_google_auth/tests/test_backwards_compatibility.py @@ -29,6 +29,7 @@ def setUp(self): self.c.sp_id = "sample_sp_id" self.c.u2f_disabled = False self.c.username = "sample_username" + self.c.account = "123456789012" self.c.raise_if_invalid() self.c.write(None) diff --git a/aws_google_auth/tests/test_config_parser.py b/aws_google_auth/tests/test_config_parser.py index 87dd9ab..beccb4a 100644 --- a/aws_google_auth/tests/test_config_parser.py +++ b/aws_google_auth/tests/test_config_parser.py @@ -249,3 +249,27 @@ def test_with_environment(self): args = parse_args([]) config = resolve_config(args) self.assertEqual(config.bg_response, 'foo') + + +class TestAccountProcessing(unittest.TestCase): + + @nottest + def test_default(self): + args = parse_args([]) + config = resolve_config(args) + self.assertEqual(None, config.account) + + def test_cli_param_supplied(self): + args = parse_args(['--account', "123456789012"]) + config = resolve_config(args) + self.assertEqual("123456789012", config.account) + + @mock.patch.dict(os.environ, {'AWS_ACCOUNT': '123456789012'}) + def test_with_environment(self): + args = parse_args([]) + config = resolve_config(args) + self.assertEqual("123456789012", config.account) + + args = parse_args(['--region', "123456789012"]) + config = resolve_config(args) + self.assertEqual("123456789012", config.account) diff --git a/aws_google_auth/tests/test_configuration_persistence.py b/aws_google_auth/tests/test_configuration_persistence.py index da07150..d05135b 100644 --- a/aws_google_auth/tests/test_configuration_persistence.py +++ b/aws_google_auth/tests/test_configuration_persistence.py @@ -33,6 +33,7 @@ def setUp(self): self.c.bg_response = "foo" self.c.raise_if_invalid() self.c.write(None) + self.c.account = "123456789012" self.config_parser = configparser.RawConfigParser() self.config_parser.read(self.c.config_file) diff --git a/aws_google_auth/tests/test_init.py b/aws_google_auth/tests/test_init.py index 3450e4e..99aac13 100644 --- a/aws_google_auth/tests/test_init.py +++ b/aws_google_auth/tests/test_init.py @@ -60,7 +60,8 @@ def test_main_method_chaining(self, process_auth, resolve_config, exit_if_unsupp print_creds=False, username=None, quiet=False, - bg_response=None)) + bg_response=None, + account=None)) ], resolve_config.mock_calls) @@ -80,7 +81,8 @@ def test_main_method_chaining(self, process_auth, resolve_config, exit_if_unsupp print_creds=False, username=None, quiet=False, - bg_response=None), + bg_response=None, + account=None), mock_config) ], process_auth.mock_calls) @@ -98,6 +100,7 @@ def test_process_auth_standard(self, mock_google, mock_amazon, mock_util): mock_config.idp_id = None mock_config.sp_id = None mock_config.return_value = None + mock_config.account = None mock_amazon_client = Mock() mock_google_client = Mock() @@ -169,6 +172,7 @@ def test_process_auth_print_creds(self, mock_google, mock_amazon, mock_util): mock_config.sp_id = None mock_config.return_value = None mock_config.print_creds = True + mock_config.account = None mock_amazon_client = Mock() mock_google_client = Mock() @@ -316,6 +320,7 @@ def test_process_auth_dont_resolve_alias(self, mock_google, mock_amazon, mock_ut mock_config.sp_id = None mock_config.return_value = None mock_config.keyring = False + mock_config.account = None mock_amazon_client = Mock() mock_google_client = Mock() @@ -349,6 +354,7 @@ def test_process_auth_dont_resolve_alias(self, mock_google, mock_amazon, mock_ut self.assertEqual(mock_config.password, "pass") self.assertEqual(mock_config.provider, "da_provider") self.assertEqual(mock_config.role_arn, "da_role") + self.assertEqual(mock_config.account, None) # Assert calls occur self.assertEqual([call.Util.get_input('Google username: '), @@ -387,6 +393,7 @@ def test_process_auth_with_profile(self, mock_google, mock_amazon, mock_util): mock_config.profile = "blart" mock_config.return_value = None mock_config.role_arn = 'arn:aws:iam::123456789012:role/admin' + mock_config.account = None mock_amazon_client = Mock() mock_google_client = Mock() @@ -459,6 +466,7 @@ def test_process_auth_with_saml_cache(self, mock_google, mock_amazon, mock_util) mock_config.password = None mock_config.return_value = None mock_config.role_arn = 'arn:aws:iam::123456789012:role/admin' + mock_config.account = None mock_amazon_client = Mock() mock_google_client = Mock() diff --git a/aws_google_auth/util.py b/aws_google_auth/util.py index 0e42ea9..4aacac6 100644 --- a/aws_google_auth/util.py +++ b/aws_google_auth/util.py @@ -18,10 +18,15 @@ def get_input(prompt): return input(prompt) @staticmethod - def pick_a_role(roles, aliases=None): + def pick_a_role(roles, aliases=None, account=None): + if account: + filtered_roles = {role: principal for role, principal in roles.items() if(account in role)} + else: + filtered_roles = roles + if aliases: enriched_roles = {} - for role, principal in roles.items(): + for role, principal in filtered_roles.items(): enriched_roles[role] = [ aliases[role.split(':')[4]], role.split('role/')[1], @@ -48,14 +53,14 @@ def pick_a_role(roles, aliases=None): print("Invalid choice, try again.") else: while True: - for i, role in enumerate(roles): + for i, role in enumerate(filtered_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(filtered_roles)) choice = Util.get_input(prompt) try: - return list(roles.items())[int(choice) - 1] + return list(filtered_roles.items())[int(choice) - 1] except (IndexError, ValueError): print("Invalid choice, try again.") From d5c10deb24bc181486dbc7a6c21d2e33e2035c20 Mon Sep 17 00:00:00 2001 From: Adam Vondersaar Date: Wed, 27 Nov 2019 14:21:25 -0800 Subject: [PATCH 2/2] fix flake8 --- aws_google_auth/configuration.py | 2 +- aws_google_auth/tests/test_args_parser.py | 4 ---- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/aws_google_auth/configuration.py b/aws_google_auth/configuration.py index 3eef927..2dc26b1 100644 --- a/aws_google_auth/configuration.py +++ b/aws_google_auth/configuration.py @@ -246,7 +246,7 @@ def read(self, profile): self.bg_response = coalesce(read_bg_response, self.bg_response) # Account - read_account= unicode_to_string(config_parser[profile_string].get('account', None)) + read_account = unicode_to_string(config_parser[profile_string].get('account', None)) self.account = coalesce(read_account, self.account) # SAML Cache diff --git a/aws_google_auth/tests/test_args_parser.py b/aws_google_auth/tests/test_args_parser.py index 59837b0..8abadb2 100644 --- a/aws_google_auth/tests/test_args_parser.py +++ b/aws_google_auth/tests/test_args_parser.py @@ -33,7 +33,6 @@ def test_no_arguments(self): self.assertEqual(parser.bg_response, None) self.assertEqual(parser.account, None) - self.assertFalse(parser.save_failure_html) # Assert the size of the parameter so that new parameters trigger a review of this function @@ -57,7 +56,6 @@ def test_username(self): self.assertEqual(parser.username, 'username@gmail.com') self.assertEqual(parser.account, None) - def test_nocache(self): parser = parse_args(['--no-cache']) @@ -75,7 +73,6 @@ def test_nocache(self): self.assertEqual(parser.username, None) self.assertEqual(parser.account, None) - def test_resolvealiases(self): parser = parse_args(['--resolve-aliases']) @@ -93,7 +90,6 @@ def test_resolvealiases(self): self.assertEqual(parser.username, None) self.assertEqual(parser.account, None) - def test_ask_and_supply_role(self): with self.assertRaises(SystemExit):