[v2] Add redirect port option to SSO and login commands

.changes/next-release/enhancement-login-redirect-port.json

@@ -0,0 +1,5 @@
+{
+  "type": "enhancement",
+  "category": "login",
+  "description": "Support ``--redirect-port`` in ``aws login`` to select a fixed port for the Authorization Code callback server."
+}

.changes/next-release/enhancement-sso-redirect-port.json

@@ -0,0 +1,5 @@
+{
+  "type": "enhancement",
+  "category": "sso",
+  "description": "Support ``--redirect-port`` in ``aws sso login`` and ``aws configure sso`` to select a fixed port for the Authorization Code callback server."
+}

awscli/customizations/configure/sso_commands.py

@@ -52,6 +52,7 @@
     BaseSSOCommand,
     PrintOnlyHandler,
     do_sso_login,
+    validate_redirect_port,
 )
 from awscli.customizations.utils import uni_print
 from awscli.formatter import CLI_OUTPUT_FORMATS
@@ -321,6 +322,7 @@ def _unset_session_profile(self):
         config_store.set_config_provider('profile', ConstantProvider(None))
 
     def _run_main(self, parsed_args, parsed_globals):
+        validate_redirect_port(parsed_args.redirect_port)
         super()._run_main(parsed_args, parsed_globals)
         self._unset_session_profile()
         on_pending_authorization = None
@@ -338,6 +340,7 @@ def _run_main(self, parsed_args, parsed_globals):
             token_cache=self._sso_token_cache,
             on_pending_authorization=on_pending_authorization,
             use_device_code=parsed_args.use_device_code,
+            redirect_port=parsed_args.redirect_port,
             **sso_registration_args,
         )
 

awscli/customizations/login/login.py

@@ -27,10 +27,12 @@
 )
 from awscli.customizations.prompts import yes_no_choice
 from awscli.customizations.sso.utils import (
+    REDIRECT_PORT_ARG,
     AuthCodeFetcher,
     OpenBrowserHandler,
     PrintOnlyHandler,
     open_browser_with_original_ld_path,
+    validate_redirect_port,
 )
 from awscli.customizations.utils import uni_print
 
@@ -63,7 +65,8 @@ class LoginCommand(BasicCommand):
                 'intended when running the CLI on remote hosts via SSH '
                 'where a local browser is not available.'
             ),
-        }
+        },
+        REDIRECT_PORT_ARG,
     ]
 
     def __init__(
@@ -88,6 +91,7 @@ def __init__(
         self._config_file_writer = config_file_writer
 
     def _run_main(self, parsed_args, parsed_globals):
+        validate_redirect_port(parsed_args.redirect_port)
         region = self._resolve_region(parsed_globals)
         profile_name = self.resolve_profile_name()
         sign_in_type = self.resolve_sign_in_type(parsed_args)
@@ -119,7 +123,9 @@ def _run_main(self, parsed_args, parsed_globals):
         if sign_in_type is LoginType.SAME_DEVICE:
             token_fetcher = SameDeviceLoginTokenFetcher(
                 client=client,
-                auth_code_fetcher=AuthCodeFetcher(),
+                auth_code_fetcher=AuthCodeFetcher(
+                    redirect_port=parsed_args.redirect_port,
+                ),
                 on_pending_authorization=OpenBrowserHandler(
                     open_browser=open_browser_with_original_ld_path
                 ),

awscli/customizations/sso/login.py

@@ -21,6 +21,7 @@
     BaseSSOCommand,
     PrintOnlyHandler,
     do_sso_login,
+    validate_redirect_port,
 )
 from awscli.customizations.utils import uni_print
 
@@ -45,10 +46,11 @@ class LoginCommand(BaseSSOCommand):
                 'of the requested profile and generally does not require this '
                 'argument to be set.'
             ),
-        }
+        },
     ]
 
     def _run_main(self, parsed_args, parsed_globals):
+        validate_redirect_port(parsed_args.redirect_port)
         sso_config = self._get_sso_config(sso_session=parsed_args.sso_session)
         start_url = sso_config['sso_start_url']
         configured_region = sso_config.get('sso_region')
@@ -81,6 +83,7 @@ def _run_main(self, parsed_args, parsed_globals):
             session_name=sso_config.get('session_name'),
             registration_scopes=sso_config.get('registration_scopes'),
             use_device_code=parsed_args.use_device_code,
+            redirect_port=parsed_args.redirect_port,
         )
 
         # Only rewrite sso_region after successful login.

awscli/customizations/sso/utils.py

@@ -35,13 +35,28 @@
 from awscli import __version__ as awscli_version
 from awscli.customizations.assumerole import CACHE_DIR as AWS_CREDS_CACHE_DIR
 from awscli.customizations.commands import BasicCommand
-from awscli.customizations.exceptions import ConfigurationError
+from awscli.customizations.exceptions import (
+    ConfigurationError,
+    ParamValidationError,
+)
 from awscli.customizations.utils import uni_print
 
 LOG = logging.getLogger(__name__)
 
 SSO_TOKEN_DIR = os.path.expanduser(os.path.join('~', '.aws', 'sso', 'cache'))
 
+REDIRECT_PORT_ARG = {
+    'name': 'redirect-port',
+    'cli_type_name': 'integer',
+    'help_text': (
+        'The localhost port to use for the Authorization Code '
+        'callback server. When omitted, a random available port is '
+        'selected. This option is ignored when the login flow does not '
+        'use a callback server.'
+    ),
+    'required': False,
+}
+
 LOGIN_ARGS = [
     {
         'name': 'no-browser',
@@ -61,9 +76,20 @@
             'instead of the Authorization Code flow.'
         ),
     },
+    REDIRECT_PORT_ARG,
 ]
 
 
+def validate_redirect_port(redirect_port):
+    if redirect_port is None:
+        return
+    if redirect_port < 1 or redirect_port > 65535:
+        raise ParamValidationError(
+            'Invalid value for --redirect-port. '
+            'Value must be between 1 and 65535.'
+        )
+
+
 def _serialize_utc_timestamp(obj):
     if isinstance(obj, datetime.datetime):
         return obj.strftime('%Y-%m-%dT%H:%M:%SZ')
@@ -86,6 +112,7 @@ def do_sso_login(
     session_name=None,
     use_device_code=False,
     resolved_start_url=None,
+    redirect_port=None,
 ):
     if token_cache is None:
         token_cache = JSONFileCache(SSO_TOKEN_DIR, dumps_func=_sso_json_dumps)
@@ -101,7 +128,7 @@ def do_sso_login(
             sso_region=sso_region,
             client_creator=session.create_client,
             parsed_globals=parsed_globals,
-            auth_code_fetcher=AuthCodeFetcher(),
+            auth_code_fetcher=AuthCodeFetcher(redirect_port=redirect_port),
             cache=token_cache,
             on_pending_authorization=on_pending_authorization,
         )
@@ -230,7 +257,7 @@ class AuthCodeFetcher:
     # How long we wait overall for the callback
     _OVERALL_TIMEOUT = 60 * 10
 
-    def __init__(self):
+    def __init__(self, redirect_port=None):
         self._auth_code = None
         self._state = None
         self._is_done = False
@@ -239,7 +266,8 @@ def __init__(self):
         # AuthCodeFetcher so that it can pass back the state and auth code
         try:
             handler = partial(OAuthCallbackHandler, self)
-            self.http_server = HTTPServer(('', 0), handler)
+            server_port = 0 if redirect_port is None else redirect_port
+            self.http_server = HTTPServer(('', server_port), handler)
             self.http_server.timeout = self._REQUEST_TIMEOUT
         except OSError as e:
             raise AuthCodeFetcherError(error_msg=e)

tests/functional/login/test_login.py

@@ -6,10 +6,13 @@
 
 import pytest
 
-from awscli.customizations.exceptions import ConfigurationError
+from awscli.customizations.exceptions import (
+    ConfigurationError,
+    ParamValidationError,
+)
 from awscli.customizations.login.login import LoginCommand
 
-DEFAULT_ARGS = Namespace(remote=False)
+DEFAULT_ARGS = Namespace(remote=False, redirect_port=None)
 DEFAULT_GLOBAL_ARGS = Namespace(
     region='us-east-1', endpoint_url=None, verify_ssl=None
 )
@@ -46,6 +49,8 @@ def config_variables(key):
     mock_session.get_config_variable.side_effect = config_variables
     mock_session.full_config = {'profiles': {'profile-name': {}}}
     mock_session._profile_map = {'profile-name': {}}
+    mock_session.user_agent_extra = ''
+    mock_session.emit_first_non_none_response.return_value = None
 
     return mock_session
 
@@ -61,16 +66,20 @@ def mock_login_command(
     )
 
 
+@pytest.mark.parametrize('redirect_port', [None, 34535])
+@mock.patch('awscli.customizations.login.login.AuthCodeFetcher')
 @mock.patch('awscli.customizations.login.utils.get_base_sign_in_uri')
 @mock.patch(
     'awscli.customizations.login.utils.SameDeviceLoginTokenFetcher.fetch_token'
 )
 def test_run_main_same_device_flow(
     mock_token_fetcher,
     mock_base_sign_in_uri,
+    mock_auth_code_fetcher,
     mock_login_command,
     mock_token_loader,
     mock_config_file_writer,
+    redirect_port,
 ):
     mock_base_sign_in_uri.return_value = 'https://foo'
     mock_token_fetcher.return_value = (
@@ -82,9 +91,13 @@ def test_run_main_same_device_flow(
         'arn:aws:iam::0123456789012:user/Admin',
     )
 
-    mock_login_command._run_main(DEFAULT_ARGS, DEFAULT_GLOBAL_ARGS)
+    args = []
+    if redirect_port is not None:
+        args = ['--redirect-port', str(redirect_port)]
+    mock_login_command(args, DEFAULT_GLOBAL_ARGS)
 
     mock_token_fetcher.assert_called_once()
+    mock_auth_code_fetcher.assert_called_once_with(redirect_port=redirect_port)
 
     mock_token_loader.save_token.assert_called_once_with(
         'arn:aws:iam::0123456789012:user/Admin',
@@ -104,20 +117,35 @@ def test_run_main_same_device_flow(
     )
 
 
+@pytest.mark.parametrize('redirect_port', [-1, 0, 65536])
+def test_invalid_redirect_port(mock_login_command, redirect_port):
+    with mock.patch.object(mock_login_command, '_resolve_region') as region:
+        with pytest.raises(ParamValidationError, match='--redirect-port'):
+            mock_login_command(
+                ['--redirect-port', str(redirect_port)], DEFAULT_GLOBAL_ARGS
+            )
+    region.assert_not_called()
+    mock_login_command._session.create_client.assert_not_called()
+
+
+@pytest.mark.parametrize('redirect_port', [None, 34535])
+@mock.patch('awscli.customizations.login.login.AuthCodeFetcher')
 @mock.patch('awscli.customizations.login.utils.get_base_sign_in_uri')
 @mock.patch(
     'awscli.customizations.login.utils.CrossDeviceLoginTokenFetcher.fetch_token'
 )
 def test_run_main_cross_device_flow(
     mock_token_fetcher,
     mock_base_sign_in_uri,
+    mock_auth_code_fetcher,
     mock_login_command,
     mock_token_loader,
     mock_config_file_writer,
+    redirect_port,
 ):
-    # Set the --remote argument
-    args = Namespace(**vars(DEFAULT_ARGS))
-    args.remote = True
+    args = ['--remote']
+    if redirect_port is not None:
+        args += ['--redirect-port', str(redirect_port)]
 
     mock_base_sign_in_uri.return_value = 'https://foo'
     mock_token_fetcher.return_value = (
@@ -129,9 +157,10 @@ def test_run_main_cross_device_flow(
         'arn:aws:iam::0123456789012:user/Admin',
     )
 
-    mock_login_command._run_main(args, DEFAULT_GLOBAL_ARGS)
+    mock_login_command(args, DEFAULT_GLOBAL_ARGS)
 
     mock_token_fetcher.assert_called_once()
+    mock_auth_code_fetcher.assert_not_called()
 
     mock_token_loader.save_token.assert_called_once_with(
         'arn:aws:iam::0123456789012:user/Admin',

tests/functional/sso/__init__.py

@@ -168,14 +168,17 @@ def assert_device_browser_handler_called_with(
             verificationUriComplete, kwargs['verificationUriComplete']
         )
 
-    def assert_auth_browser_handler_called_with(self, expected_scopes):
+    def assert_auth_browser_handler_called_with(
+        self, expected_scopes, expected_redirect_port=55555
+    ):
         # The endpoint is subject to the endpoint rules, and the
         # code_challenge is not fixed so assert against the rest of the url
         expected_url = (
             'authorize?'
             'response_type=code'
             '&client_id=auth-client-id'
-            '&redirect_uri=http%3A%2F%2F127.0.0.1%3A55555%2Foauth%2Fcallback'
+            '&redirect_uri=http%3A%2F%2F127.0.0.1%3A'
+            f'{expected_redirect_port}%2Foauth%2Fcallback'
             '&state=00000000-0000-0000-0000-000000000000'
             '&code_challenge_method=S256'
             '&scopes=' + expected_scopes

tests/functional/sso/test_login.py

@@ -297,6 +297,7 @@ def test_login_auth_sso_session(self):
         self.run_cmd('sso login')
         self.assert_used_expected_sso_region(expected_region=self.sso_region)
         self.assert_auth_browser_handler_called_with('sso%3Aaccount%3Aaccess')
+        self.fetcher_mock.assert_called_once_with(redirect_port=None)
         self.assert_cache_contains_registration(
             start_url=self.start_url,
             session_name='test-session',
@@ -309,6 +310,49 @@ def test_login_auth_sso_session(self):
             expected_token=self.access_token,
         )
 
+    def test_login_auth_sso_session_with_redirect_port(self):
+        content = self.get_sso_session_config('test-session')
+        self.set_config_file_content(content=content)
+        self.fetcher_mock.return_value.redirect_uri_with_port.return_value = (
+            'http://127.0.0.1:34535/oauth/callback'
+        )
+        self.add_oidc_auth_code_responses(self.access_token)
+        self.run_cmd('sso login --redirect-port 34535')
+        self.fetcher_mock.assert_called_once_with(redirect_port=34535)
+        self.assert_auth_browser_handler_called_with(
+            'sso%3Aaccount%3Aaccess',
+            expected_redirect_port=34535,
+        )
+
+    def test_login_invalid_redirect_port(self):
+        for redirect_port in (-1, 0, 65536):
+            with self.subTest(redirect_port=redirect_port):
+                _, stderr, _ = self.run_cmd(
+                    f'sso login --redirect-port {redirect_port}',
+                    expected_rc=252,
+                )
+                self.assertIn('Invalid value for --redirect-port', stderr)
+                self.fetcher_mock.assert_not_called()
+
+    def test_login_non_integer_redirect_port(self):
+        _, stderr, _ = self.run_cmd(
+            'sso login --redirect-port invalid', expected_rc=255
+        )
+        self.assertIn("invalid literal for int()", stderr)
+        self.fetcher_mock.assert_not_called()
+
+    def test_login_device_code_ignores_redirect_port(self):
+        content = self.get_sso_session_config('test-session')
+        self.set_config_file_content(content=content)
+        self.add_oidc_device_responses(self.access_token)
+        self.run_cmd('sso login --use-device-code --redirect-port 34535')
+        self.fetcher_mock.assert_not_called()
+        self.assert_cache_contains_token(
+            start_url=self.start_url,
+            session_name='test-session',
+            expected_token=self.access_token,
+        )
+
     def test_login_device_sso_with_explicit_sso_session_arg(self):
         content = self.get_sso_session_config(
             'test-session', include_profile=False