Skip to content

Commit 8454f0b

Browse files
Port (#9754)
1 parent 087c6e3 commit 8454f0b

File tree

4 files changed

+144
-1
lines changed

4 files changed

+144
-1
lines changed

awscli/botocore/credentials.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2298,6 +2298,7 @@ def _get_credentials(self):
22982298
'accessToken': token,
22992299
}
23002300
try:
2301+
register_feature_ids(self.feature_ids)
23012302
response = client.get_role_credentials(**kwargs)
23022303
except client.exceptions.UnauthorizedException:
23032304
raise UnauthorizedSSOTokenError()
@@ -2353,6 +2354,7 @@ def __init__(
23532354
self._load_config = load_config
23542355
self._client_creator = client_creator
23552356
self._profile_name = profile_name
2357+
self._feature_ids = set()
23562358

23572359
def _load_sso_config(self):
23582360
loaded_config = self._load_config()
@@ -2427,11 +2429,23 @@ def load(self):
24272429
'token_loader': SSOTokenLoader(cache=self._token_cache),
24282430
'cache': self.cache,
24292431
}
2430-
if 'sso_session' in sso_config:
2432+
sso_session_in_config = 'sso_session' in sso_config
2433+
if sso_session_in_config:
24312434
fetcher_kwargs['sso_session_name'] = sso_config['sso_session']
24322435
fetcher_kwargs['token_provider'] = self._token_provider
2436+
self._feature_ids.add('CREDENTIALS_PROFILE_SSO')
2437+
else:
2438+
self._feature_ids.add('CREDENTIALS_PROFILE_SSO_LEGACY')
24332439

24342440
sso_fetcher = SSOCredentialFetcher(**fetcher_kwargs)
2441+
sso_fetcher.feature_ids = self._feature_ids.copy()
2442+
2443+
if sso_session_in_config:
2444+
self._feature_ids.add('CREDENTIALS_SSO')
2445+
else:
2446+
self._feature_ids.add('CREDENTIALS_SSO_LEGACY')
2447+
2448+
register_feature_ids(self._feature_ids)
24352449

24362450
return DeferredRefreshableCredentials(
24372451
method=self.METHOD,

awscli/botocore/useragent.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,10 @@
8585
'CREDENTIALS_PROFILE_SOURCE_PROFILE': 'o',
8686
'CREDENTIALS_PROFILE_NAMED_PROVIDER': 'p',
8787
'CREDENTIALS_PROFILE_STS_WEB_ID_TOKEN': 'q',
88+
'CREDENTIALS_PROFILE_SSO': 'r',
89+
'CREDENTIALS_SSO': 's',
90+
'CREDENTIALS_PROFILE_SSO_LEGACY': 't',
91+
'CREDENTIALS_SSO_LEGACY': 'u',
8892
'CREDENTIALS_HTTP': 'z',
8993
'CREDENTIALS_IMDS': '0',
9094
'BEARER_SERVICE_ENV_VARS': '3',

tests/functional/botocore/test_credentials.py

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1285,6 +1285,89 @@ def _assert_feature_ids_in_ua(client, expected_feature_ids):
12851285
assert expected_id in feature_list
12861286

12871287

1288+
@patch("botocore.credentials.CachedCredentialFetcher._load_from_cache")
1289+
@patch("botocore.credentials.SSOProvider._load_sso_config")
1290+
@patch(
1291+
"botocore.credentials.AssumeRoleWithWebIdentityProvider.load",
1292+
return_value=None,
1293+
)
1294+
@patch("botocore.credentials.AssumeRoleProvider.load", return_value=None)
1295+
@patch("botocore.credentials.EnvProvider.load", return_value=None)
1296+
def test_user_agent_has_sso_legacy_credentials_feature_id(
1297+
_unused_mock_env_load,
1298+
_unused_mock_shared_load,
1299+
_unused_mock_config_load,
1300+
mock_load_sso_config,
1301+
mock_load_sso_credentials,
1302+
monkeypatch,
1303+
patched_session,
1304+
):
1305+
fake_fetcher_kwargs = {
1306+
'sso_start_url': "https://test.awsapps.com/start",
1307+
'sso_region': "us-east-1",
1308+
'sso_role_name': "Administrator",
1309+
'sso_account_id': "1234567890",
1310+
}
1311+
fake_response = {
1312+
"ProviderType": "sso",
1313+
"Credentials": {
1314+
"role_name": "FAKEROLE",
1315+
"AccessKeyId": "FAKEACCESSKEY",
1316+
"SecretAccessKey": "FAKESECRET",
1317+
"SessionToken": "FAKETOKEN",
1318+
"Expiration": "2099-01-01T00:00:00Z",
1319+
},
1320+
}
1321+
1322+
mock_load_sso_config.return_value = fake_fetcher_kwargs
1323+
client_one = patched_session.create_client("s3", region_name="us-east-1")
1324+
mock_load_sso_credentials.return_value = fake_response
1325+
1326+
_assert_feature_ids_in_ua(client_one, ['t', 'u'])
1327+
1328+
1329+
@patch("botocore.credentials.CachedCredentialFetcher._load_from_cache")
1330+
@patch("botocore.credentials.SSOProvider._load_sso_config")
1331+
@patch(
1332+
"botocore.credentials.AssumeRoleWithWebIdentityProvider.load",
1333+
return_value=None,
1334+
)
1335+
@patch("botocore.credentials.AssumeRoleProvider.load", return_value=None)
1336+
@patch("botocore.credentials.EnvProvider.load", return_value=None)
1337+
def test_user_agent_has_sso_credentials_feature_id(
1338+
_unused_mock_env_load,
1339+
_unused_mock_shared_load,
1340+
_unused_mock_config_load,
1341+
mock_load_sso_config,
1342+
mock_load_sso_credentials,
1343+
monkeypatch,
1344+
patched_session,
1345+
):
1346+
fake_fetcher_kwargs = {
1347+
'sso_session': 'sample_test',
1348+
'sso_start_url': "https://test.awsapps.com/start",
1349+
'sso_region': "us-east-1",
1350+
'sso_role_name': "Administrator",
1351+
'sso_account_id': "1234567890",
1352+
}
1353+
fake_response = {
1354+
"ProviderType": "sso",
1355+
"Credentials": {
1356+
"role_name": "FAKEROLE",
1357+
"AccessKeyId": "FAKEACCESSKEY",
1358+
"SecretAccessKey": "FAKESECRET",
1359+
"SessionToken": "FAKETOKEN",
1360+
"Expiration": "2099-01-01T00:00:00Z",
1361+
},
1362+
}
1363+
1364+
mock_load_sso_config.return_value = fake_fetcher_kwargs
1365+
client_one = patched_session.create_client("s3", region_name="us-east-1")
1366+
mock_load_sso_credentials.return_value = fake_response
1367+
1368+
_assert_feature_ids_in_ua(client_one, ['r', 's'])
1369+
1370+
12881371
@pytest.mark.parametrize(
12891372
"config_content,env_vars,expected_source_features,expected_provider_feature",
12901373
[

tests/unit/botocore/test_credentials.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3868,6 +3868,48 @@ def test_expired_legacy_token_has_expected_behavior(self):
38683868
fetcher.fetch_credentials()
38693869
self.assertFalse(mock_client.get_role_credentials.called)
38703870

3871+
@mock.patch('botocore.credentials.register_feature_ids')
3872+
def test_feature_ids_registered_during_get_credentials(
3873+
self, mock_register
3874+
):
3875+
response = {
3876+
'roleCredentials': {
3877+
'accessKeyId': 'foo',
3878+
'secretAccessKey': 'bar',
3879+
'sessionToken': 'baz',
3880+
'expiration': self.now_timestamp + 1000000,
3881+
}
3882+
}
3883+
params = {
3884+
'roleName': self.role_name,
3885+
'accountId': self.account_id,
3886+
'accessToken': self.access_token['accessToken'],
3887+
}
3888+
self.stubber.add_response(
3889+
'get_role_credentials',
3890+
response,
3891+
expected_params=params,
3892+
)
3893+
3894+
self.stubber.activate()
3895+
try:
3896+
fetcher = SSOCredentialFetcher(
3897+
self.start_url,
3898+
self.sso_region,
3899+
self.role_name,
3900+
self.account_id,
3901+
self.mock_session.create_client,
3902+
token_loader=self.loader,
3903+
cache=self.cache,
3904+
time_fetcher=self.mock_time_fetcher,
3905+
)
3906+
test_feature_ids = {'test_feature_1', 'test_feature_2'}
3907+
fetcher.feature_ids = test_feature_ids
3908+
fetcher.fetch_credentials()
3909+
mock_register.assert_called_once_with(test_feature_ids)
3910+
finally:
3911+
self.stubber.deactivate()
3912+
38713913

38723914
class TestSSOProvider(unittest.TestCase):
38733915
def setUp(self):

0 commit comments

Comments
 (0)