diff --git a/gimme_aws_creds/main.py b/gimme_aws_creds/main.py index b110602..7040295 100644 --- a/gimme_aws_creds/main.py +++ b/gimme_aws_creds/main.py @@ -420,25 +420,30 @@ def _get_user_int_selections_many(self, min_int, max_int, max_retries=5): for _ in range(max_retries): selections = set() error = False + input_values = self.ui.input('Selections (comma separated; specify "A" for all profiles): ').lower().split(',') - for value in self.ui.input('Selections (comma separated): ').split(','): - value = value.strip() - - if not value: - continue - - try: - selection = int(value) - except ValueError: - self.ui.warning('Invalid selection {}, must be an integer value.'.format(repr(value))) - error = True - continue - - if min_int <= selection <= max_int: - selections.add(value) - else: - self.ui.warning( - 'Selection {} out of range <{}, {}>'.format(repr(selection), min_int, max_int)) + if 'a' in input_values or 'all' in input_values: + selections = set(range(min_int, max_int+1)) + self.ui.message("Generating credentials for profile indices {} to {} ".format(min_int, max_int)) + else: + for value in input_values: + value = value.strip() + + if not value: + continue + + try: + selection = int(value) + except ValueError: + self.ui.warning('Invalid selection {}, must be an integer value, or "A" for all profiles'.format(repr(value))) + error = True + continue + + if min_int <= selection <= max_int: + selections.add(value) + else: + self.ui.warning( + 'Selection {} out of range <{}, {}>'.format(repr(selection), min_int, max_int)) if error: continue diff --git a/tests/test_main.py b/tests/test_main.py index f6469b2..ca4883f 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -47,12 +47,26 @@ def test_choose_roles_app_2(self, mock): self.assertRaises(errors.GimmeAWSCredsExitBase, creds._choose_roles, self.APP_INFO) self.assertRaises(errors.GimmeAWSCredsExitBase, creds._choose_app, self.AWS_INFO) - @patch('builtins.input', return_value='a') - def test_choose_roles_app_a(self, mock): + @patch('builtins.input', return_value='b') + def test_choose_roles_app_b(self, mock): creds = GimmeAWSCreds() self.assertRaises(errors.GimmeAWSCredsExitBase, creds._choose_roles, self.APP_INFO) self.assertRaises(errors.GimmeAWSCredsExitBase, creds._choose_app, self.AWS_INFO) + @patch('builtins.input', return_value='all') + def test_choose_roles_app_select_all(self, mock): + creds = GimmeAWSCreds() + selections = creds._choose_roles(self.APP_INFO) + expected_selections = {self.APP_INFO[0].role, self.APP_INFO[1].role} + self.assertEqual(selections, expected_selections) + + @patch('builtins.input', return_value='a') + def test_choose_roles_app_select_all_a(self, mock): + creds = GimmeAWSCreds() + selections = creds._choose_roles(self.APP_INFO) + expected_selections = {self.APP_INFO[0].role, self.APP_INFO[1].role} + self.assertEqual(selections, expected_selections) + def test_get_selected_app_from_config_0(self): creds = GimmeAWSCreds()