Port track feature IDs for authentication related to SSO credentials

awscli/botocore/credentials.py

@@ -2298,6 +2298,7 @@ def _get_credentials(self):
             'accessToken': token,
         }
         try:
+            register_feature_ids(self.feature_ids)
             response = client.get_role_credentials(**kwargs)
         except client.exceptions.UnauthorizedException:
             raise UnauthorizedSSOTokenError()
@@ -2353,6 +2354,7 @@ def __init__(
         self._load_config = load_config
         self._client_creator = client_creator
         self._profile_name = profile_name
+        self._feature_ids = set()
 
     def _load_sso_config(self):
         loaded_config = self._load_config()
@@ -2427,11 +2429,23 @@ def load(self):
             'token_loader': SSOTokenLoader(cache=self._token_cache),
             'cache': self.cache,
         }
-        if 'sso_session' in sso_config:
+        sso_session_in_config = 'sso_session' in sso_config
+        if sso_session_in_config:
             fetcher_kwargs['sso_session_name'] = sso_config['sso_session']
             fetcher_kwargs['token_provider'] = self._token_provider
+            self._feature_ids.add('CREDENTIALS_PROFILE_SSO')
+        else:
+            self._feature_ids.add('CREDENTIALS_PROFILE_SSO_LEGACY')
 
         sso_fetcher = SSOCredentialFetcher(**fetcher_kwargs)
+        sso_fetcher.feature_ids = self._feature_ids.copy()
+
+        if sso_session_in_config:
+            self._feature_ids.add('CREDENTIALS_SSO')
+        else:
+            self._feature_ids.add('CREDENTIALS_SSO_LEGACY')
+
+        register_feature_ids(self._feature_ids)
 
         return DeferredRefreshableCredentials(
             method=self.METHOD,

awscli/botocore/useragent.py

@@ -85,6 +85,10 @@
     'CREDENTIALS_PROFILE_SOURCE_PROFILE': 'o',
     'CREDENTIALS_PROFILE_NAMED_PROVIDER': 'p',
     'CREDENTIALS_PROFILE_STS_WEB_ID_TOKEN': 'q',
+    'CREDENTIALS_PROFILE_SSO': 'r',
+    'CREDENTIALS_SSO': 's',
+    'CREDENTIALS_PROFILE_SSO_LEGACY': 't',
+    'CREDENTIALS_SSO_LEGACY': 'u',
     'CREDENTIALS_HTTP': 'z',
     'CREDENTIALS_IMDS': '0',
     'BEARER_SERVICE_ENV_VARS': '3',

tests/functional/botocore/test_credentials.py

@@ -1285,6 +1285,89 @@ def _assert_feature_ids_in_ua(client, expected_feature_ids):
             assert expected_id in feature_list
 
 
+@patch("botocore.credentials.CachedCredentialFetcher._load_from_cache")
+@patch("botocore.credentials.SSOProvider._load_sso_config")
+@patch(
+    "botocore.credentials.AssumeRoleWithWebIdentityProvider.load",
+    return_value=None,
+)
+@patch("botocore.credentials.AssumeRoleProvider.load", return_value=None)
+@patch("botocore.credentials.EnvProvider.load", return_value=None)
+def test_user_agent_has_sso_legacy_credentials_feature_id(
+    _unused_mock_env_load,
+    _unused_mock_shared_load,
+    _unused_mock_config_load,
+    mock_load_sso_config,
+    mock_load_sso_credentials,
+    monkeypatch,
+    patched_session,
+):
+    fake_fetcher_kwargs = {
+        'sso_start_url': "https://test.awsapps.com/start",
+        'sso_region': "us-east-1",
+        'sso_role_name': "Administrator",
+        'sso_account_id': "1234567890",
+    }
+    fake_response = {
+        "ProviderType": "sso",
+        "Credentials": {
+            "role_name": "FAKEROLE",
+            "AccessKeyId": "FAKEACCESSKEY",
+            "SecretAccessKey": "FAKESECRET",
+            "SessionToken": "FAKETOKEN",
+            "Expiration": "2099-01-01T00:00:00Z",
+        },
+    }
+
+    mock_load_sso_config.return_value = fake_fetcher_kwargs
+    client_one = patched_session.create_client("s3", region_name="us-east-1")
+    mock_load_sso_credentials.return_value = fake_response
+
+    _assert_feature_ids_in_ua(client_one, ['t', 'u'])
+
+
+@patch("botocore.credentials.CachedCredentialFetcher._load_from_cache")
+@patch("botocore.credentials.SSOProvider._load_sso_config")
+@patch(
+    "botocore.credentials.AssumeRoleWithWebIdentityProvider.load",
+    return_value=None,
+)
+@patch("botocore.credentials.AssumeRoleProvider.load", return_value=None)
+@patch("botocore.credentials.EnvProvider.load", return_value=None)
+def test_user_agent_has_sso_credentials_feature_id(
+    _unused_mock_env_load,
+    _unused_mock_shared_load,
+    _unused_mock_config_load,
+    mock_load_sso_config,
+    mock_load_sso_credentials,
+    monkeypatch,
+    patched_session,
+):
+    fake_fetcher_kwargs = {
+        'sso_session': 'sample_test',
+        'sso_start_url': "https://test.awsapps.com/start",
+        'sso_region': "us-east-1",
+        'sso_role_name': "Administrator",
+        'sso_account_id': "1234567890",
+    }
+    fake_response = {
+        "ProviderType": "sso",
+        "Credentials": {
+            "role_name": "FAKEROLE",
+            "AccessKeyId": "FAKEACCESSKEY",
+            "SecretAccessKey": "FAKESECRET",
+            "SessionToken": "FAKETOKEN",
+            "Expiration": "2099-01-01T00:00:00Z",
+        },
+    }
+
+    mock_load_sso_config.return_value = fake_fetcher_kwargs
+    client_one = patched_session.create_client("s3", region_name="us-east-1")
+    mock_load_sso_credentials.return_value = fake_response
+
+    _assert_feature_ids_in_ua(client_one, ['r', 's'])
+
+
 @pytest.mark.parametrize(
     "config_content,env_vars,expected_source_features,expected_provider_feature",
     [

tests/unit/botocore/test_credentials.py

@@ -3868,6 +3868,48 @@ def test_expired_legacy_token_has_expected_behavior(self):
             fetcher.fetch_credentials()
         self.assertFalse(mock_client.get_role_credentials.called)
 
+    @mock.patch('botocore.credentials.register_feature_ids')
+    def test_feature_ids_registered_during_get_credentials(
+        self, mock_register
+    ):
+        response = {
+            'roleCredentials': {
+                'accessKeyId': 'foo',
+                'secretAccessKey': 'bar',
+                'sessionToken': 'baz',
+                'expiration': self.now_timestamp + 1000000,
+            }
+        }
+        params = {
+            'roleName': self.role_name,
+            'accountId': self.account_id,
+            'accessToken': self.access_token['accessToken'],
+        }
+        self.stubber.add_response(
+            'get_role_credentials',
+            response,
+            expected_params=params,
+        )
+
+        self.stubber.activate()
+        try:
+            fetcher = SSOCredentialFetcher(
+                self.start_url,
+                self.sso_region,
+                self.role_name,
+                self.account_id,
+                self.mock_session.create_client,
+                token_loader=self.loader,
+                cache=self.cache,
+                time_fetcher=self.mock_time_fetcher,
+            )
+            test_feature_ids = {'test_feature_1', 'test_feature_2'}
+            fetcher.feature_ids = test_feature_ids
+            fetcher.fetch_credentials()
+            mock_register.assert_called_once_with(test_feature_ids)
+        finally:
+            self.stubber.deactivate()
+
 
 class TestSSOProvider(unittest.TestCase):
     def setUp(self):