Add configurable redirect port for SSO and login OAuth flows

.changes/next-release/feature-login-34536.json

@@ -0,0 +1,5 @@
+{
+  "type": "enhancement",
+  "category": "``login``",
+  "description": "Add a ``--redirect-port`` option to set the localhost port that the callback server listens on, instead of choosing an available port at random."
+}

.changes/next-release/feature-sso-34535.json

@@ -0,0 +1,5 @@
+{
+  "type": "enhancement",
+  "category": "``sso``",
+  "description": "Add a ``--redirect-port`` option to ``aws sso login`` and ``aws configure sso`` to set the localhost port that the callback server listens on, instead of choosing an available port at random."
+}

awscli/customizations/configure/sso_commands.py

@@ -49,6 +49,7 @@
     BaseSSOCommand,
     PrintOnlyHandler,
     do_sso_login,
+    validate_redirect_port,
 )
 from awscli.customizations.utils import uni_print
 from awscli.formatter import CLI_OUTPUT_FORMATS
@@ -315,6 +316,7 @@ def _unset_session_profile(self):
 
     def _run_main(self, parsed_args, parsed_globals):
         super()._run_main(parsed_args, parsed_globals)
+        validate_redirect_port(parsed_args.redirect_port)
         self._unset_session_profile()
         on_pending_authorization = None
         if parsed_args.no_browser:
@@ -331,6 +333,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

@@ -24,10 +24,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
 
@@ -60,7 +62,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__(
@@ -85,6 +88,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)
@@ -115,7 +119,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
 
@@ -49,6 +50,7 @@ class LoginCommand(BaseSSOCommand):
     ]
 
     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,29 @@
 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 port on localhost that the local callback server listens on, '
+        'which is used to build the ``redirect_uri`` for the authorization '
+        'request. By default an available port is chosen at random. Set this '
+        'when the port must be known in advance, such as when mapping it from '
+        'a Docker container to the host. Has no effect when a flow that does '
+        'not start a callback server is used.'
+    ),
+}
+
 LOGIN_ARGS = [
     {
         'name': 'no-browser',
@@ -61,9 +77,21 @@
             'instead of the Authorization Code flow.'
         ),
     },
+    REDIRECT_PORT_ARG,
 ]
 
 
+def validate_redirect_port(redirect_port):
+    """Validates a ``--redirect-port`` value before starting a login flow."""
+    if redirect_port is None:
+        return
+    if not 1 <= redirect_port <= 65535:
+        raise ParamValidationError(
+            f'Invalid value for --redirect-port: {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 +114,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 +130,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,16 +259,20 @@ 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
 
+        # Binding to port 0 lets the OS pick any available port, which is what
+        # we want unless the caller needs the port to be known in advance.
+        port = 0 if redirect_port is None else redirect_port
+
         # We do this so that the request handler can have a reference to this
         # AuthCodeFetcher so that it can pass back the state and auth code
         try:
             handler = partial(OAuthCallbackHandler, self)
-            self.http_server = HTTPServer(('', 0), handler)
+            self.http_server = HTTPServer(('', 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
 )
@@ -104,6 +107,50 @@ def test_run_main_same_device_flow(
     )
 
 
+@mock.patch('awscli.customizations.login.utils.get_base_sign_in_uri')
+@mock.patch('awscli.customizations.login.login.AuthCodeFetcher')
+@mock.patch(
+    'awscli.customizations.login.utils.SameDeviceLoginTokenFetcher.fetch_token'
+)
+@pytest.mark.parametrize('redirect_port', [None, 34535])
+def test_run_main_passes_redirect_port_to_auth_code_fetcher(
+    mock_token_fetcher,
+    mock_auth_code_fetcher,
+    mock_base_sign_in_uri,
+    mock_login_command,
+    redirect_port,
+):
+    mock_base_sign_in_uri.return_value = 'https://foo'
+    mock_token_fetcher.return_value = (
+        {
+            'accessToken': 'access_token',
+            'idToken': SAMPLE_ID_TOKEN,
+            'expiresIn': 3600,
+        },
+        'arn:aws:iam::0123456789012:user/Admin',
+    )
+    args = Namespace(**vars(DEFAULT_ARGS))
+    args.redirect_port = redirect_port
+
+    mock_login_command._run_main(args, DEFAULT_GLOBAL_ARGS)
+
+    mock_auth_code_fetcher.assert_called_once_with(redirect_port=redirect_port)
+
+
+@mock.patch('awscli.customizations.login.login.AuthCodeFetcher')
+def test_run_main_rejects_out_of_range_redirect_port(
+    mock_auth_code_fetcher, mock_login_command
+):
+    args = Namespace(**vars(DEFAULT_ARGS))
+    args.redirect_port = 65536
+
+    with pytest.raises(ParamValidationError) as excinfo:
+        mock_login_command._run_main(args, DEFAULT_GLOBAL_ARGS)
+
+    assert 'must be between 1 and 65535' in str(excinfo.value)
+    mock_auth_code_fetcher.assert_not_called()
+
+
 @mock.patch('awscli.customizations.login.utils.get_base_sign_in_uri')
 @mock.patch(
     'awscli.customizations.login.utils.CrossDeviceLoginTokenFetcher.fetch_token'

tests/functional/sso/test_login.py

@@ -443,6 +443,29 @@ def test_login_auth_sso_state_mismatch_throws_error(self):
         _, stderr, _ = self.run_cmd('sso login', expected_rc=255)
         self.assertIn('State parameter does not match expected value.', stderr)
 
+    def test_login_auth_default_redirect_port(self):
+        content = self.get_sso_session_config('test-session')
+        self.set_config_file_content(content=content)
+        self.add_oidc_auth_code_responses(self.access_token)
+        self.run_cmd('sso login')
+        self.fetcher_mock.assert_called_once_with(redirect_port=None)
+
+    def test_login_auth_explicit_redirect_port(self):
+        content = self.get_sso_session_config('test-session')
+        self.set_config_file_content(content=content)
+        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)
+
+    def test_login_rejects_out_of_range_redirect_port(self):
+        content = self.get_sso_session_config('test-session')
+        self.set_config_file_content(content=content)
+        _, stderr, _ = self.run_cmd(
+            'sso login --redirect-port 65536', expected_rc=252
+        )
+        self.assertIn('must be between 1 and 65535', stderr)
+        self.fetcher_mock.assert_not_called()
+
     def test_login_device_no_extra_user_agent(self):
         self.add_oidc_device_responses(self.access_token)
         self.run_cmd('sso login --use-device-code')

tests/unit/customizations/configure/test_sso.py

@@ -41,6 +41,7 @@
     display_account,
     get_account_sorting_key,
 )
+from awscli.customizations.exceptions import ParamValidationError
 from awscli.customizations.sso.utils import (
     PrintOnlyHandler,
     do_sso_login,
@@ -881,6 +882,7 @@ def assert_do_sso_login_call(
         expected_scopes=None,
         expected_auth_handler_cls=None,
         expected_force_refresh=None,
+        expected_redirect_port=None,
     ):
         expected_kwargs = {
             "sso_region": expected_sso_region,
@@ -889,6 +891,7 @@ def assert_do_sso_login_call(
             "on_pending_authorization": None,
             "token_cache": None,
             "use_device_code": expected_use_device_code,
+            "redirect_port": expected_redirect_port,
         }
         if expected_session_name is not None:
             expected_kwargs["session_name"] = expected_session_name
@@ -957,6 +960,44 @@ def test_legacy_configure_sso_flow(
             in stdout
         )
 
+    def test_configure_sso_passes_redirect_port(
+        self,
+        sso_cmd,
+        ptk_stubber,
+        aws_config,
+        stub_simple_single_item_sso_responses,
+        mock_do_sso_login,
+        botocore_session,
+        parsed_globals,
+        configure_sso_legacy_inputs,
+        account_id,
+        role_name,
+    ):
+        inputs = configure_sso_legacy_inputs
+        inputs.skip_account_and_role_selection()
+        ptk_stubber.user_inputs = inputs
+        stub_simple_single_item_sso_responses(account_id, role_name)
+
+        sso_cmd(["--redirect-port", "34535"], parsed_globals)
+        self.assert_do_sso_login_call(
+            mock_do_sso_login,
+            botocore_session,
+            expected_sso_region=inputs.sso_region_prompt.answer,
+            expected_start_url=inputs.start_url_prompt.answer,
+            expected_redirect_port=34535,
+        )
+
+    def test_configure_sso_rejects_out_of_range_redirect_port(
+        self,
+        sso_cmd,
+        parsed_globals,
+        mock_do_sso_login,
+    ):
+        with pytest.raises(ParamValidationError) as excinfo:
+            sso_cmd(["--redirect-port", "65536"], parsed_globals)
+        assert "must be between 1 and 65535" in str(excinfo.value)
+        mock_do_sso_login.assert_not_called()
+
     def test_single_account_single_role_flow_no_browser(
         self,
         sso_cmd,