feat: Add framework for BYOID metrics headers (#1332)

* Add framework for BYOID metrics headers

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* responding to PR comments

* fix: changing try catch to if statement

* Fix lint and test coverage issue

* fix comment

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
Co-authored-by: Leo <39062083+lsirac@users.noreply.github.com>

Author: 3 people

packages/google-auth/google/auth/aws.py

@@ -742,6 +742,11 @@ def _should_use_metadata_server(self):
 
         return False
 
+    def _create_default_metrics_options(self):
+        metrics_options = super(Credentials, self)._create_default_metrics_options()
+        metrics_options["source"] = "aws"
+        return metrics_options
+
     @classmethod
     def from_info(cls, info, **kwargs):
         """Creates an AWS Credentials instance from parsed external account info.

packages/google-auth/google/auth/external_account.py

@@ -40,6 +40,7 @@
 from google.auth import credentials
 from google.auth import exceptions
 from google.auth import impersonated_credentials
+from google.auth import metrics
 from google.oauth2 import sts
 from google.oauth2 import utils
 
@@ -140,6 +141,8 @@ def __init__(
             self._client_auth = None
         self._sts_client = sts.Client(self._token_url, self._client_auth)
 
+        self._metrics_options = self._create_default_metrics_options()
+
         if self._service_account_impersonation_url:
             self._impersonated_credentials = self._initialize_impersonated_credentials()
         else:
@@ -284,7 +287,9 @@ def token_info_url(self):
     def with_scopes(self, scopes, default_scopes=None):
         kwargs = self._constructor_args()
         kwargs.update(scopes=scopes, default_scopes=default_scopes)
-        return self.__class__(**kwargs)
+        scoped = self.__class__(**kwargs)
+        scoped._metrics_options = self._metrics_options
+        return scoped
 
     @abc.abstractmethod
     def retrieve_subject_token(self, request):
@@ -362,6 +367,11 @@ def refresh(self, request):
             # is used. The client ID is sufficient for determining the user project.
             if self._workforce_pool_user_project and not self._client_id:
                 additional_options = {"userProject": self._workforce_pool_user_project}
+            additional_headers = {
+                metrics.API_CLIENT_HEADER: metrics.byoid_metrics_header(
+                    self._metrics_options
+                )
+            }
             response_data = self._sts_client.exchange_token(
                 request=request,
                 grant_type=_STS_GRANT_TYPE,
@@ -371,6 +381,7 @@ def refresh(self, request):
                 scopes=scopes,
                 requested_token_type=_STS_REQUESTED_TOKEN_TYPE,
                 additional_options=additional_options,
+                additional_headers=additional_headers,
             )
             self.token = response_data.get("access_token")
             lifetime = datetime.timedelta(seconds=response_data.get("expires_in"))
@@ -381,13 +392,17 @@ def with_quota_project(self, quota_project_id):
         # Return copy of instance with the provided quota project ID.
         kwargs = self._constructor_args()
         kwargs.update(quota_project_id=quota_project_id)
-        return self.__class__(**kwargs)
+        new_cred = self.__class__(**kwargs)
+        new_cred._metrics_options = self._metrics_options
+        return new_cred
 
     @_helpers.copy_docstring(credentials.CredentialsWithTokenUri)
     def with_token_uri(self, token_uri):
         kwargs = self._constructor_args()
         kwargs.update(token_url=token_uri)
-        return self.__class__(**kwargs)
+        new_cred = self.__class__(**kwargs)
+        new_cred._metrics_options = self._metrics_options
+        return new_cred
 
     def _initialize_impersonated_credentials(self):
         """Generates an impersonated credentials.
@@ -411,6 +426,7 @@ def _initialize_impersonated_credentials(self):
             service_account_impersonation_options={},
         )
         source_credentials = self.__class__(**kwargs)
+        source_credentials._metrics_options = self._metrics_options
 
         # Determine target_principal.
         target_principal = self.service_account_email
@@ -432,6 +448,19 @@ def _initialize_impersonated_credentials(self):
             ),
         )
 
+    def _create_default_metrics_options(self):
+        metrics_options = {}
+        if self._service_account_impersonation_url:
+            metrics_options["sa-impersonation"] = "true"
+        else:
+            metrics_options["sa-impersonation"] = "false"
+        if self._service_account_impersonation_options.get("token_lifetime_seconds"):
+            metrics_options["config-lifetime"] = "true"
+        else:
+            metrics_options["config-lifetime"] = "false"
+
+        return metrics_options
+
     @classmethod
     def from_info(cls, info, **kwargs):
         """Creates a Credentials instance from parsed external account info.

packages/google-auth/google/auth/identity_pool.py

@@ -216,6 +216,18 @@ def _parse_token_data(
             )
         return token
 
+    def _create_default_metrics_options(self):
+        metrics_options = super(Credentials, self)._create_default_metrics_options()
+        # Check that credential source is a dict before checking for file vs url. This check needs to be done
+        # here because the external_account credential constructor needs to pass the metrics options to the
+        # impersonated credential object before the identity_pool credentials are validated.
+        if isinstance(self._credential_source, Mapping):
+            if self._credential_source.get("file"):
+                metrics_options["source"] = "file"
+            else:
+                metrics_options["source"] = "url"
+        return metrics_options
+
     @classmethod
     def from_info(cls, info, **kwargs):
         """Creates an Identity Pool Credentials instance from parsed external account info.

packages/google-auth/google/auth/metrics.py

@@ -23,6 +23,9 @@
 
 API_CLIENT_HEADER = "x-goog-api-client"
 
+# BYOID Specific consts
+BYOID_HEADER_SECTION = "google-byoid-sdk"
+
 # Auth request type
 REQUEST_TYPE_ACCESS_TOKEN = "auth-request-type/at"
 REQUEST_TYPE_ID_TOKEN = "auth-request-type/it"
@@ -123,6 +126,15 @@ def reauth_continue():
     return "{} {}".format(python_and_auth_lib_version(), REQUEST_TYPE_REAUTH_CONTINUE)
 
 
+# x-goog-api-client header value for BYOID calls to the Security Token Service exchange token endpoint.
+# Example: "gl-python/3.7 auth/1.1 google-byoid-sdk source/aws sa-impersonation/true sa-impersonation/true"
+def byoid_metrics_header(metrics_options):
+    header = "{} {}".format(python_and_auth_lib_version(), BYOID_HEADER_SECTION)
+    for key, value in metrics_options.items():
+        header = "{} {}/{}".format(header, key, value)
+    return header
+
+
 def add_metric_header(headers, metric_header_value):
     """Add x-goog-api-client header with the given value.
 

packages/google-auth/google/auth/pluggable.py

@@ -422,3 +422,8 @@ def _validate_running_mode(self):
             raise exceptions.InvalidValue(
                 "Interactive mode is only enabled for workforce pool."
             )
+
+    def _create_default_metrics_options(self):
+        metrics_options = super(Credentials, self)._create_default_metrics_options()
+        metrics_options["source"] = "executable"
+        return metrics_options

packages/google-auth/tests/test_aws.py

@@ -32,6 +32,8 @@
     "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/imp"
 )
 
+LANG_LIBRARY_METRICS_HEADER_VALUE = "gl-python/3.7 auth/1.1"
+
 CLIENT_ID = "username"
 CLIENT_SECRET = "password"
 # Base64 encoding of "username:password".
@@ -1793,8 +1795,14 @@ def test_retrieve_subject_token_error_determining_aws_security_creds(self):
 
         assert excinfo.match(r"Unable to retrieve AWS security credentials")
 
+    @mock.patch(
+        "google.auth.metrics.python_and_auth_lib_version",
+        return_value=LANG_LIBRARY_METRICS_HEADER_VALUE,
+    )
     @mock.patch("google.auth._helpers.utcnow")
-    def test_refresh_success_without_impersonation_ignore_default_scopes(self, utcnow):
+    def test_refresh_success_without_impersonation_ignore_default_scopes(
+        self, utcnow, mock_auth_lib_value
+    ):
         utcnow.return_value = datetime.datetime.strptime(
             self.AWS_SIGNATURE_TIME, "%Y-%m-%dT%H:%M:%SZ"
         )
@@ -1808,6 +1816,7 @@ def test_refresh_success_without_impersonation_ignore_default_scopes(self, utcno
         token_headers = {
             "Content-Type": "application/x-www-form-urlencoded",
             "Authorization": "Basic " + BASIC_AUTH_ENCODING,
+            "x-goog-api-client": "gl-python/3.7 auth/1.1 google-byoid-sdk sa-impersonation/false config-lifetime/false source/aws",
         }
         token_request_data = {
             "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
@@ -1849,8 +1858,14 @@ def test_refresh_success_without_impersonation_ignore_default_scopes(self, utcno
         assert credentials.scopes == SCOPES
         assert credentials.default_scopes == ["ignored"]
 
+    @mock.patch(
+        "google.auth.metrics.python_and_auth_lib_version",
+        return_value=LANG_LIBRARY_METRICS_HEADER_VALUE,
+    )
     @mock.patch("google.auth._helpers.utcnow")
-    def test_refresh_success_without_impersonation_use_default_scopes(self, utcnow):
+    def test_refresh_success_without_impersonation_use_default_scopes(
+        self, utcnow, mock_auth_lib_value
+    ):
         utcnow.return_value = datetime.datetime.strptime(
             self.AWS_SIGNATURE_TIME, "%Y-%m-%dT%H:%M:%SZ"
         )
@@ -1864,6 +1879,7 @@ def test_refresh_success_without_impersonation_use_default_scopes(self, utcnow):
         token_headers = {
             "Content-Type": "application/x-www-form-urlencoded",
             "Authorization": "Basic " + BASIC_AUTH_ENCODING,
+            "x-goog-api-client": "gl-python/3.7 auth/1.1 google-byoid-sdk sa-impersonation/false config-lifetime/false source/aws",
         }
         token_request_data = {
             "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
@@ -1909,9 +1925,13 @@ def test_refresh_success_without_impersonation_use_default_scopes(self, utcnow):
         "google.auth.metrics.token_request_access_token_impersonate",
         return_value=IMPERSONATE_ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE,
     )
+    @mock.patch(
+        "google.auth.metrics.python_and_auth_lib_version",
+        return_value=LANG_LIBRARY_METRICS_HEADER_VALUE,
+    )
     @mock.patch("google.auth._helpers.utcnow")
     def test_refresh_success_with_impersonation_ignore_default_scopes(
-        self, utcnow, mock_metrics_header_value
+        self, utcnow, mock_metrics_header_value, mock_auth_lib_value
     ):
         utcnow.return_value = datetime.datetime.strptime(
             self.AWS_SIGNATURE_TIME, "%Y-%m-%dT%H:%M:%SZ"
@@ -1929,6 +1949,7 @@ def test_refresh_success_with_impersonation_ignore_default_scopes(
         token_headers = {
             "Content-Type": "application/x-www-form-urlencoded",
             "Authorization": "Basic " + BASIC_AUTH_ENCODING,
+            "x-goog-api-client": "gl-python/3.7 auth/1.1 google-byoid-sdk sa-impersonation/true config-lifetime/false source/aws",
         }
         token_request_data = {
             "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
@@ -2000,9 +2021,13 @@ def test_refresh_success_with_impersonation_ignore_default_scopes(
         "google.auth.metrics.token_request_access_token_impersonate",
         return_value=IMPERSONATE_ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE,
     )
+    @mock.patch(
+        "google.auth.metrics.python_and_auth_lib_version",
+        return_value=LANG_LIBRARY_METRICS_HEADER_VALUE,
+    )
     @mock.patch("google.auth._helpers.utcnow")
     def test_refresh_success_with_impersonation_use_default_scopes(
-        self, utcnow, mock_metrics_header_value
+        self, utcnow, mock_metrics_header_value, mock_auth_lib_value
     ):
         utcnow.return_value = datetime.datetime.strptime(
             self.AWS_SIGNATURE_TIME, "%Y-%m-%dT%H:%M:%SZ"
@@ -2020,6 +2045,7 @@ def test_refresh_success_with_impersonation_use_default_scopes(
         token_headers = {
             "Content-Type": "application/x-www-form-urlencoded",
             "Authorization": "Basic " + BASIC_AUTH_ENCODING,
+            "x-goog-api-client": "gl-python/3.7 auth/1.1 google-byoid-sdk sa-impersonation/true config-lifetime/false source/aws",
         }
         token_request_data = {
             "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",