diff --git a/src/awscli_login/exceptions.py b/src/awscli_login/exceptions.py index 31824509..8ec98c45 100644 --- a/src/awscli_login/exceptions.py +++ b/src/awscli_login/exceptions.py @@ -86,3 +86,11 @@ class RoleParseFail(SAML): def __init__(self, role: str) -> None: mesg = "Bad SAML Response! Failed to parse role: %s!" super().__init__(mesg % role) + + +class InvalidSelection(ConfigError): + code = 11 + + def __init__(self) -> None: + mesg = "Invalid selection!\a" + super().__init__(mesg) diff --git a/src/awscli_login/util.py b/src/awscli_login/util.py index 8739dd00..465c7215 100644 --- a/src/awscli_login/util.py +++ b/src/awscli_login/util.py @@ -10,7 +10,10 @@ from botocore.session import Session from .const import ERROR_INVALID_PROFILE_ROLE -from .exceptions import SAML +from .exceptions import ( + InvalidSelection, + SAML, +) from .typing import Role awsconfigfile = path.join('.aws', 'credentials') @@ -69,8 +72,10 @@ def get_selection(role_arns: List[Role], profile_role: str = None) -> Role: i += 1 print("Selection:\a ", end='') -# TODO need error checking - return role_arns[select[int(input())]] + try: + return role_arns[select[int(input())]] + except (ValueError, KeyError): + raise InvalidSelection elif n == 1: return role_arns[0] else: diff --git a/src/tests/test_util.py b/src/tests/test_util.py index b8be542a..b09bd363 100755 --- a/src/tests/test_util.py +++ b/src/tests/test_util.py @@ -11,7 +11,10 @@ from botocore.session import Session from awscli_login.const import ERROR_INVALID_PROFILE_ROLE -from awscli_login.exceptions import SAML +from awscli_login.exceptions import ( + InvalidSelection, + SAML, +) from awscli_login.util import ( get_selection, remove_credentials, @@ -130,6 +133,30 @@ def test_get_2of2_selections(self, *args): self.assertEqual(get_selection(roles), roles[1]) + @patch('builtins.input', return_value=3) + @patch('sys.stdout', new=StringIO()) + def test_get_bad_numeric_selection(self, *args): + """ Invalid numeric selection of two roles """ + roles = [ + ('idp1', 'arn:aws:iam::224588347132:role/KalturaAdmin'), + ('idp2', 'arn:aws:iam::617683844790:role/BoxAdmin'), + ] + + with self.assertRaises(InvalidSelection): + get_selection(roles) + + @patch('builtins.input', return_value="foo") + @patch('sys.stdout', new=StringIO()) + def test_get_bad_type_selection(self, *args): + """ Invalid string selection of two roles """ + roles = [ + ('idp1', 'arn:aws:iam::224588347132:role/KalturaAdmin'), + ('idp2', 'arn:aws:iam::617683844790:role/BoxAdmin'), + ] + + with self.assertRaises(InvalidSelection): + get_selection(roles) + @patch('builtins.input', return_value=1) def test_selections_profile_role(self, *args): """ Profile role is selected when valid and present """