Track assume-role and profile related credential feature ids (#3554)

Author: arandito

botocore/credentials.py

@@ -44,7 +44,7 @@
     UnknownCredentialError,
 )
 from botocore.tokens import SSOTokenProvider
-from botocore.useragent import register_feature_id
+from botocore.useragent import register_feature_id, register_feature_ids
 from botocore.utils import (
     ArnParser,
     ContainerMetadataFetcher,
@@ -705,6 +705,7 @@ def __init__(self, cache=None, expiry_window_seconds=None):
         if expiry_window_seconds is None:
             expiry_window_seconds = self.DEFAULT_EXPIRY_WINDOW_SECONDS
         self._expiry_window_seconds = expiry_window_seconds
+        self.feature_ids = set()
 
     def _create_cache_key(self):
         raise NotImplementedError('_create_cache_key()')
@@ -885,6 +886,7 @@ def __init__(
 
     def _get_credentials(self):
         """Get credentials by calling assume role."""
+        register_feature_ids(self.feature_ids)
         kwargs = self._assume_role_kwargs()
         client = self._create_client()
         response = client.assume_role(**kwargs)
@@ -971,6 +973,7 @@ def __init__(
 
     def _get_credentials(self):
         """Get credentials by calling assume role."""
+        register_feature_ids(self.feature_ids)
         kwargs = self._assume_role_kwargs()
         # Assume role with web identity does not require credentials other than
         # the token, explicitly configure the client to not sign requests.
@@ -1367,6 +1370,7 @@ def load(self):
                 )
                 token = self._get_session_token(config)
                 account_id = self._get_account_id(config)
+                register_feature_id('CREDENTIALS_PROFILE')
                 return Credentials(
                     access_key,
                     secret_key,
@@ -1434,6 +1438,7 @@ def load(self):
                 )
                 token = self._get_session_token(profile_config)
                 account_id = self._get_account_id(profile_config)
+                register_feature_id('CREDENTIALS_PROFILE')
                 return Credentials(
                     access_key,
                     secret_key,
@@ -1513,6 +1518,11 @@ class AssumeRoleProvider(CredentialProvider):
     # remaining time left until the credentials expires is less than the
     # EXPIRY_WINDOW.
     EXPIRY_WINDOW_SECONDS = 60 * 15
+    NAMED_PROVIDER_FEATURE_MAP = {
+        'Ec2InstanceMetadata': 'CREDENTIALS_IMDS',
+        'Environment': 'CREDENTIALS_ENV_VARS',
+        'EcsContainer': 'CREDENTIALS_HTTP',
+    }
 
     def __init__(
         self,
@@ -1575,6 +1585,7 @@ def __init__(
         self._credential_sourcer = credential_sourcer
         self._profile_provider_builder = profile_provider_builder
         self._visited_profiles = [self._profile_name]
+        self._feature_ids = set()
 
     def load(self):
         self._loaded_config = self._load_config()
@@ -1625,10 +1636,13 @@ def _load_creds_via_assume_role(self, profile_name):
             mfa_prompter=self._prompter,
             cache=self.cache,
         )
+        fetcher.feature_ids = self._feature_ids.copy()
         refresher = fetcher.fetch_credentials
         if mfa_serial is not None:
             refresher = create_mfa_serial_refresher(refresher)
 
+        self._feature_ids.add('CREDENTIALS_STS_ASSUME_ROLE')
+        register_feature_ids(self._feature_ids)
         # The initial credentials are empty and the expiration time is set
         # to now so that we can delay the call to assume role until it is
         # strictly needed.
@@ -1757,18 +1771,20 @@ def _has_static_credentials(self, profile):
     def _resolve_source_credentials(self, role_config, profile_name):
         credential_source = role_config.get('credential_source')
         if credential_source is not None:
+            self._feature_ids.add('CREDENTIALS_PROFILE_NAMED_PROVIDER')
             return self._resolve_credentials_from_source(
                 credential_source, profile_name
             )
 
         source_profile = role_config['source_profile']
         self._visited_profiles.append(source_profile)
+        self._feature_ids.add('CREDENTIALS_PROFILE_SOURCE_PROFILE')
         return self._resolve_credentials_from_profile(source_profile)
 
     def _resolve_credentials_from_profile(self, profile_name):
         profiles = self._loaded_config.get('profiles', {})
         profile = profiles[profile_name]
-
+        self._feature_ids.add('CREDENTIALS_PROFILE')
         if (
             self._has_static_credentials(profile)
             and not self._profile_provider_builder
@@ -1824,6 +1840,11 @@ def _resolve_credentials_from_source(
                     f'in profile {profile_name}'
                 ),
             )
+        named_provider_feature_id = self.NAMED_PROVIDER_FEATURE_MAP.get(
+            credential_source
+        )
+        if named_provider_feature_id:
+            self._feature_ids.add(named_provider_feature_id)
         return credentials
 
 
@@ -1854,6 +1875,7 @@ def __init__(
         if token_loader_cls is None:
             token_loader_cls = FileWebIdentityTokenLoader
         self._token_loader_cls = token_loader_cls
+        self._feature_ids = set()
 
     def load(self):
         return self._assume_role_with_web_identity()
@@ -1876,8 +1898,15 @@ def _get_env_config(self, key):
     def _get_config(self, key):
         env_value = self._get_env_config(key)
         if env_value is not None:
+            self._feature_ids.add('CREDENTIALS_ENV_VARS_STS_WEB_ID_TOKEN')
             return env_value
-        return self._get_profile_config(key)
+
+        config_value = self._get_profile_config(key)
+        if config_value is not None:
+            self._feature_ids.add('CREDENTIALS_PROFILE_STS_WEB_ID_TOKEN')
+            return config_value
+
+        return None
 
     def _assume_role_with_web_identity(self):
         token_path = self._get_config('web_identity_token_file')
@@ -1907,6 +1936,10 @@ def _assume_role_with_web_identity(self):
             extra_args=extra_args,
             cache=self.cache,
         )
+        fetcher.feature_ids = self._feature_ids.copy()
+
+        self._feature_ids.add('CREDENTIALS_STS_ASSUME_ROLE_WEB_ID')
+        register_feature_ids(self._feature_ids)
         # The initial credentials are empty and the expiration time is set
         # to now so that we can delay the call to assume role until it is
         # strictly needed.

botocore/useragent.py

@@ -80,6 +80,13 @@
     'FLEXIBLE_CHECKSUMS_RES_WHEN_REQUIRED': 'c',
     'CREDENTIALS_CODE': 'e',
     'CREDENTIALS_ENV_VARS': 'g',
+    'CREDENTIALS_ENV_VARS_STS_WEB_ID_TOKEN': 'h',
+    'CREDENTIALS_STS_ASSUME_ROLE': 'i',
+    'CREDENTIALS_STS_ASSUME_ROLE_WEB_ID': 'k',
+    'CREDENTIALS_PROFILE': 'n',
+    'CREDENTIALS_PROFILE_SOURCE_PROFILE': 'o',
+    'CREDENTIALS_PROFILE_NAMED_PROVIDER': 'p',
+    'CREDENTIALS_PROFILE_STS_WEB_ID_TOKEN': 'q',
     'CREDENTIALS_PROFILE_PROCESS': 'v',
     'CREDENTIALS_PROCESS': 'w',
     'CREDENTIALS_BOTO2_CONFIG_FILE': 'x',
@@ -108,6 +115,17 @@ def register_feature_id(feature_id):
         ctx.features.add(val)
 
 
+def register_feature_ids(feature_ids):
+    """Adds multiple feature IDs to the current context object's ``features`` set.
+
+    :type feature_ids: iterable of str
+    :param feature_ids: An iterable of feature ID strings to register. Each
+        value must be a key in the ``_USERAGENT_FEATURE_MAPPINGS`` dict.
+    """
+    for feature_id in feature_ids:
+        register_feature_id(feature_id)
+
+
 def sanitize_user_agent_string_component(raw_str, allow_hash):
     """Replaces all not allowed characters in the string with a dash ("-").