From 736c8e6315f3508e095433cf1be0ae8da1183e4b Mon Sep 17 00:00:00 2001
From: "David D. Riddle" <ddriddle@illinois.edu>
Date: Tue, 19 Aug 2025 11:20:54 -0500
Subject: [PATCH] Prompt user again on invalid selection

Closes #94
---
 src/awscli_login/exceptions.py | 12 ++++++++---
 src/awscli_login/util.py       | 38 +++++++++++++++++++++++-----------
 src/tests/test_util.py         | 19 ++++++++++++++---
 3 files changed, 51 insertions(+), 18 deletions(-)

diff --git a/src/awscli_login/exceptions.py b/src/awscli_login/exceptions.py
index fd7aba72..fac8b49d 100644
--- a/src/awscli_login/exceptions.py
+++ b/src/awscli_login/exceptions.py
@@ -10,6 +10,13 @@ class ConfigError(AWSCLILogin):
     pass
 
 
+class UserExit(ConfigError):
+    code = 0
+
+    def __init__(self) -> None:
+        super().__init__("No role selected. Good bye!")
+
+
 class AlreadyLoggedIn(ConfigError):
     code = 2
 
@@ -89,12 +96,11 @@ def __init__(self, role: str) -> None:
         super().__init__(mesg % role)
 
 
-class InvalidSelection(ConfigError):
+class TooManyInvalidSelections(ConfigError):
     code = 11
 
     def __init__(self) -> None:
-        mesg = "Invalid selection!\a"
-        super().__init__(mesg)
+        super().__init__("Too many invalid selections!")
 
 
 class TooManyHttpTrafficFlags(ConfigError):
diff --git a/src/awscli_login/util.py b/src/awscli_login/util.py
index b6fd2a18..a4c84583 100644
--- a/src/awscli_login/util.py
+++ b/src/awscli_login/util.py
@@ -41,9 +41,10 @@ class Session:  # type: ignore
     ConfigError,
     CredentialProcessMisconfigured,
     CredentialProcessNotSet,
-    InvalidSelection,
     SAML,
     TooManyHttpTrafficFlags,
+    TooManyInvalidSelections,
+    UserExit,
 )
 from ._typing import Role
 
@@ -75,9 +76,7 @@ def get_selection(role_arns: List[Role], profile_role: Optional[str] = None,
                   interactive: bool = True, aliases: Dict[str, str] = {}
                   ) -> Role:
     """ Interactively prompts the user for a role selection. """
-    i = 0
     n = len(role_arns)
-    select: Dict[int, int] = {}
 
     # Return profile_role if valid and set
     if profile_role is not None:
@@ -90,9 +89,22 @@ def get_selection(role_arns: List[Role], profile_role: Optional[str] = None,
             logger.error(ERROR_INVALID_PROFILE_ROLE % profile_role)
 
     if n > 1:
-        print("Please choose the role you would like to assume:")
+        return prompt_for_role_arn(role_arns)
+    elif n == 1:
+        return role_arns[0]
+    else:
+        raise SAML("No roles returned!")
+
+
+def prompt_for_role_arn(role_arns: List[Role], aliases: Dict[str, str] = {}):
+    """ Prompts user to select a role from the given list of roles. """
+    accounts = sort_roles(role_arns)
 
-        accounts = sort_roles(role_arns)
+    for _ in range(3):
+        i = 0
+        select: Dict[int, int] = {}
+
+        print("Please choose the role you would like to assume:")
         for acct, roles in accounts:
             name = f"{aliases[acct]} ({acct})" if acct in aliases else acct
             print(' ' * 4, "Account:", name)
@@ -102,15 +114,17 @@ def get_selection(role_arns: List[Role], profile_role: Optional[str] = None,
                 select[i] = index
                 i += 1
 
-        print("Selection:\a ", end='')
+        print("Select a role or enter 'q' to quit:\a ", end='')
         try:
-            return role_arns[select[int(input())]]
+            user_input = input()
+            return role_arns[select[int(user_input)]]
         except (ValueError, KeyError):
-            raise InvalidSelection
-    elif n == 1:
-        return role_arns[0]
-    else:
-        raise SAML("No roles returned!")
+            if user_input == 'q':
+                raise UserExit
+            print("Invalid value. Try again.")
+            continue
+
+    raise TooManyInvalidSelections
 
 
 def file2bytes(filename: str) -> bytes:
diff --git a/src/tests/test_util.py b/src/tests/test_util.py
index 21d68ef1..c3c27aea 100755
--- a/src/tests/test_util.py
+++ b/src/tests/test_util.py
@@ -14,9 +14,10 @@
 from awscli_login.exceptions import (
     CredentialProcessMisconfigured,
     CredentialProcessNotSet,
-    InvalidSelection,
     SAML,
     TooManyHttpTrafficFlags,
+    TooManyInvalidSelections,
+    UserExit,
 )
 from awscli_login.util import (
     config_vcr,
@@ -128,6 +129,18 @@ def test_get_2of2_selections(self, *args):
 
         self.assertEqual(get_selection(roles), roles[1])
 
+    @patch('builtins.input', return_value='q')
+    @patch('sys.stdout', new=StringIO())
+    def test_user_exit(self, *args):
+        """ User exits by typing 'q' """
+        roles = [
+            ('idp1', 'arn:aws:iam::123577191723:role/KalturaAdmin'),
+            ('idp2', 'arn:aws:iam::271867855970:role/BoxAdmin'),
+        ]
+
+        with self.assertRaises(UserExit):
+            get_selection(roles)
+
     @patch('builtins.input', return_value=3)
     @patch('sys.stdout', new=StringIO())
     def test_get_bad_numeric_selection(self, *args):
@@ -137,7 +150,7 @@ def test_get_bad_numeric_selection(self, *args):
             ('idp2', 'arn:aws:iam::271867855970:role/BoxAdmin'),
         ]
 
-        with self.assertRaises(InvalidSelection):
+        with self.assertRaises(TooManyInvalidSelections):
             get_selection(roles)
 
     @patch('builtins.input', return_value="foo")
@@ -149,7 +162,7 @@ def test_get_bad_type_selection(self, *args):
             ('idp2', 'arn:aws:iam::271867855970:role/BoxAdmin'),
         ]
 
-        with self.assertRaises(InvalidSelection):
+        with self.assertRaises(TooManyInvalidSelections):
             get_selection(roles)
 
     @patch('builtins.input', return_value=1)