Update SSO token provider configuration and caching:

* Migrate SSO token provider configuration to dedicated section
* Cache SSO tokens based on the sso-session name

Author: joguSD

botocore/configloader.py

@@ -253,6 +253,7 @@ def build_profile_map(parsed_ini_config):
     """
     parsed_config = copy.deepcopy(parsed_ini_config)
     profiles = {}
+    sso_sessions = {}
     final_config = {}
     for key, values in parsed_config.items():
         if key.startswith("profile"):
@@ -262,6 +263,13 @@ def build_profile_map(parsed_ini_config):
                 continue
             if len(parts) == 2:
                 profiles[parts[1]] = values
+        elif key.startswith("sso-session"):
+            try:
+                parts = shlex.split(key)
+            except ValueError:
+                continue
+            if len(parts) == 2:
+                sso_sessions[parts[1]] = values
         elif key == 'default':
             # default section is special and is considered a profile
             # name but we don't require you use 'profile "default"'
@@ -270,4 +278,5 @@ def build_profile_map(parsed_ini_config):
         else:
             final_config[key] = values
     final_config['profiles'] = profiles
+    final_config['sso_sessions'] = sso_sessions
     return final_config

botocore/tokens.py

@@ -198,24 +198,43 @@ def __init__(self, session, cache=None, time_fetcher=_utc_now):
     def _load_sso_config(self):
         loaded_config = self._session.full_config
         profiles = loaded_config.get("profiles", {})
+        sso_sessions = loaded_config.get("sso_sessions", {})
         profile_name = self._session.get_config_variable("profile")
         if not profile_name:
             profile_name = "default"
         profile_config = profiles.get(profile_name, {})
 
-        if "sso_start_url" not in profile_config:
-            return None
+        if "sso_session" not in profile_config:
+            return
+
+        sso_session_name = profile_config["sso_session"]
+        sso_config = sso_sessions.get(sso_session_name, None)
+
+        if not sso_config:
+            error_msg = (
+                f'The profile "{profile_name}" is configured to use the SSO '
+                f'token provider but the "{sso_session_name}" sso_session '
+                f"configuration does not exist."
+            )
+            raise InvalidConfigError(error_msg=error_msg)
+
+        missing_configs = []
+        for var in self._SSO_CONFIG_VARS:
+            if var not in sso_config:
+                missing_configs.append(var)
 
-        if "sso_region" not in profile_config:
+        if missing_configs:
             error_msg = (
                 f'The profile "{profile_name}" is configured to use the SSO '
-                f"token provider but is missing the sso_region configuration"
+                f"token provider but is missing the following configuration: "
+                f"{missing_configs}."
             )
             raise InvalidConfigError(error_msg=error_msg)
 
         return {
-            "sso_region": profile_config["sso_region"],
-            "sso_start_url": profile_config["sso_start_url"],
+            "session_name": sso_session_name,
+            "sso_region": sso_config["sso_region"],
+            "sso_start_url": sso_config["sso_start_url"],
         }
 
     @CachedProperty
@@ -250,9 +269,6 @@ def _attempt_create_token(self, token):
         }
         if "refreshToken" in response:
             new_token["refreshToken"] = response["refreshToken"]
-        elif "refreshToken" in token:
-            # TODO: Verify if we should preserve the old refresh token
-            new_token["refreshToken"] = token["refreshToken"]
         logger.info("SSO Token refresh succeeded")
         return new_token
 
@@ -282,8 +298,9 @@ def _refresh_access_token(self, token):
 
     def _refresher(self):
         start_url = self._sso_config["sso_start_url"]
-        logger.info(f"Loading cached SSO token for {start_url}")
-        token_dict = self._token_loader(start_url)
+        session_name = self._sso_config["session_name"]
+        logger.info(f"Loading cached SSO token for {session_name}")
+        token_dict = self._token_loader(start_url, session_name=session_name)
         expiration = dateutil.parser.parse(token_dict["expiresAt"])
         logger.debug(f"Cached SSO token expires at {expiration}")
 
@@ -293,7 +310,9 @@ def _refresher(self):
             if new_token_dict is not None:
                 token_dict = new_token_dict
                 expiration = token_dict["expiresAt"]
-                self._token_loader.save_token(start_url, token_dict)
+                self._token_loader.save_token(
+                    start_url, token_dict, session_name=session_name
+                )
 
         return FrozenAuthToken(
             token_dict["accessToken"], expiration=expiration

botocore/utils.py

@@ -2764,17 +2764,24 @@ def __init__(self, cache=None):
             cache = {}
         self._cache = cache
 
-    def _generate_cache_key(self, start_url):
-        return hashlib.sha1(start_url.encode('utf-8')).hexdigest()
-
-    def save_token(self, start_url, token):
-        cache_key = self._generate_cache_key(start_url)
+    def _generate_cache_key(self, start_url, session_name):
+        input_str = start_url
+        if session_name is not None:
+            input_str = session_name
+        return hashlib.sha1(input_str.encode('utf-8')).hexdigest()
+
+    def save_token(self, start_url, token, session_name=None):
+        cache_key = self._generate_cache_key(start_url, session_name)
         self._cache[cache_key] = token
 
-    def __call__(self, start_url):
-        cache_key = self._generate_cache_key(start_url)
+    def __call__(self, start_url, session_name=None):
+        cache_key = self._generate_cache_key(start_url, session_name)
+        logger.debug(f'Checking for cached token at: {cache_key}')
         if cache_key not in self._cache:
-            error_msg = f'Token for {start_url} does not exist'
+            name = start_url
+            if session_name is not None:
+                name = session_name
+            error_msg = f'Token for {name} does not exist'
             raise SSOTokenLoadError(error_msg=error_msg)
 
         token = self._cache[cache_key]

tests/unit/cfg/aws_sso_session_config

@@ -0,0 +1,6 @@
+[default]
+sso_session = sso
+
+[sso-session sso]
+sso_start_url = https://example.com
+sso_region = us-east-1

tests/unit/test_configloader.py

@@ -178,6 +178,16 @@ def test_unicode_bytes_path(self):
         self.assertIn('default', loaded_config['profiles'])
         self.assertIn('personal', loaded_config['profiles'])
 
+    def test_sso_session_config(self):
+        filename = path('aws_sso_session_config')
+        loaded_config = load_config(filename)
+        self.assertIn('profiles', loaded_config)
+        self.assertIn('default', loaded_config['profiles'])
+        self.assertIn('sso_sessions', loaded_config)
+        self.assertIn('sso', loaded_config['sso_sessions'])
+        sso_config = loaded_config['sso_sessions']['sso']
+        self.assertEqual(sso_config['sso_region'], 'us-east-1')
+        self.assertEqual(sso_config['sso_start_url'], 'https://example.com')
 
 if __name__ == "__main__":
     unittest.main()

tests/unit/test_tokens.py

@@ -33,42 +33,79 @@ def parametrize(cases):
 sso_provider_resolution_cases = [
     {
         "documentation": "Full valid profile",
-        "profile": {
-            "sso_region": "us-east-1",
-            "sso_start_url": "https://d-abc123.awsapps.com/start",
+        "config": {
+            "profiles": {"test": {"sso_session": "admin"}},
+            "sso_sessions": {
+                "admin": {
+                    "sso_region": "us-east-1",
+                    "sso_start_url": "https://d-abc123.awsapps.com/start",
+                }
+            },
         },
         "resolves": True,
     },
     {
         "documentation": "Non-SSO profiles are skipped",
-        "profile": {"region": "us-west-2"},
+        "config": {"profiles": {"test": {"region": "us-west-2"}}},
         "resolves": False,
     },
     {
         "documentation": "Only start URL is invalid",
-        "profile": {"sso_start_url": "https://d-abc123.awsapps.com/start"},
+        "config": {
+            "profiles": {"test": {"sso_session": "admin"}},
+            "sso_sessions": {
+                "admin": {
+                    "sso_start_url": "https://d-abc123.awsapps.com/start"
+                }
+            },
+        },
+        "resolves": False,
+        "expectedException": InvalidConfigError,
+    },
+    {
+        "documentation": "Only sso_region is invalid",
+        "config": {
+            "profiles": {"test": {"sso_session": "admin"}},
+            "sso_sessions": {"admin": {"sso_region": "us-east-1"}},
+        },
         "resolves": False,
         "expectedException": InvalidConfigError,
     },
     {
-        "documentation": "SSO Region only is skipped",
-        "profile": {"sso_region": "us-east-1"},
+        "documentation": "Specified sso-session must exist",
+        "config": {
+            "profiles": {"test": {"sso_session": "dev"}},
+            "sso_sessions": {"admin": {"sso_region": "us-east-1"}},
+        },
+        "resolves": False,
+        "expectedException": InvalidConfigError,
+    },
+    {
+        "documentation": "The sso_session must be specified",
+        "config": {
+            "profiles": {"test": {"region": "us-west-2"}},
+            "sso_sessions": {
+                "admin": {
+                    "sso_region": "us-east-1",
+                    "sso_start_url": "https://d-abc123.awsapps.com/start",
+                }
+            },
+        },
         "resolves": False,
     },
 ]
 
 
 def _create_mock_session(config):
     mock_session = mock.Mock(spec=Session)
-    mock_session.get_config_variable.return_value = "default"
-    mock_session.full_config = {"profiles": {"default": config}}
+    mock_session.get_config_variable.return_value = "test"
+    mock_session.full_config = config
     return mock_session
 
 
 @parametrize(sso_provider_resolution_cases)
 def test_sso_token_provider_resolution(test_case):
-    config = test_case["profile"]
-    mock_session = _create_mock_session(config)
+    mock_session = _create_mock_session(test_case["config"])
     resolver = SSOTokenProvider(mock_session)
 
     expected_exception = test_case.get("expectedException")
@@ -196,8 +233,6 @@ def test_sso_token_provider_resolution(test_case):
             "clientId": "clientid",
             "clientSecret": "YSBzZWNyZXQ=",
             "registrationExpiresAt": "2022-12-25T13:30:00Z",
-            # TODO: Verify if we should preserve old refresh token
-            "refreshToken": "cachedrefreshtoken",
         },
         "expectedToken": {
             "token": "newtoken",
@@ -225,10 +260,15 @@ def test_sso_token_provider_resolution(test_case):
 @parametrize(sso_provider_refresh_cases)
 def test_sso_token_provider_refresh(test_case):
     config = {
-        "sso_region": "us-west-2",
-        "sso_start_url": "https://d-123.awsapps.com/start",
+        "profiles": {"test": {"sso_session": "admin"}},
+        "sso_sessions": {
+            "admin": {
+                "sso_region": "us-west-2",
+                "sso_start_url": "https://d-123.awsapps.com/start",
+            }
+        },
     }
-    cache_key = "2b829a45f04c9828cb45b7d092d8e4aa30818393"
+    cache_key = "d033e22ae348aeb5660fc2140aec35850c4da997"
     token_cache = {}
 
     # Prepopulate the token cache

tests/unit/test_utils.py

@@ -3148,8 +3148,10 @@ def test_imds_service_endpoint_overrides_ipv6_endpoint(self, send):
 class TestSSOTokenLoader(unittest.TestCase):
     def setUp(self):
         super().setUp()
+        self.session_name = 'admin'
         self.start_url = 'https://d-abc123.awsapps.com/start'
         self.cache_key = '40a89917e3175433e361b710a9d43528d7f1890a'
+        self.session_cache_key = 'd033e22ae348aeb5660fc2140aec35850c4da997'
         self.access_token = 'totally.a.token'
         self.cached_token = {
             'accessToken': self.access_token,
@@ -3172,6 +3174,27 @@ def test_can_handle_invalid_cache(self):
         with self.assertRaises(SSOTokenLoadError):
             self.loader(self.start_url)
 
+    def test_can_save_token(self):
+        self.loader.save_token(self.start_url, self.cached_token)
+        access_token = self.loader(self.start_url)
+        self.assertEqual(self.cached_token, access_token)
+
+    def test_can_save_token_sso_session(self):
+        self.loader.save_token(
+            self.start_url, self.cached_token, session_name=self.session_name,
+        )
+        access_token = self.loader(
+            self.start_url, session_name=self.session_name,
+        )
+        self.assertEqual(self.cached_token, access_token)
+
+    def test_can_load_token_exists_sso_session_name(self):
+        self.cache[self.session_cache_key] = self.cached_token
+        access_token = self.loader(
+            self.start_url, session_name=self.session_name,
+        )
+        self.assertEqual(self.cached_token, access_token)
+
 
 @pytest.mark.parametrize(
     'header_name, headers, expected',