Merge branch 'main' into remove_rsa_3

Author: daniel-sanche

google/auth/_agent_identity_utils.py

@@ -24,7 +24,6 @@
 
 from google.auth import environment_vars
 from google.auth import exceptions
-from google.auth.transport import _mtls_helper
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -263,14 +262,6 @@ def should_request_bound_token(cert):
     return is_agent_cert and is_opted_in
 
 
-def call_client_cert_callback():
-    """Calls the client cert callback and returns the certificate and key."""
-    _, cert_bytes, key_bytes, passphrase = _mtls_helper.get_client_ssl_credentials(
-        generate_encrypted_key=True
-    )
-    return cert_bytes, key_bytes
-
-
 def get_cached_cert_fingerprint(cached_cert):
     """Returns the fingerprint of the cached certificate."""
     if cached_cert:

google/auth/credentials.py

@@ -17,20 +17,22 @@
 
 import abc
 from enum import Enum
+import logging
 import os
 from typing import List
 
 from google.auth import _helpers, environment_vars
 from google.auth import exceptions
 from google.auth import metrics
 from google.auth._credentials_base import _BaseCredentials
-from google.auth._default import _LOGGER
 from google.auth._refresh_worker import RefreshThreadManager
 
 DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
 NO_OP_TRUST_BOUNDARY_LOCATIONS: List[str] = []
 NO_OP_TRUST_BOUNDARY_ENCODED_LOCATIONS = "0x0"
 
+_LOGGER = logging.getLogger("google.auth._default")
+
 
 class Credentials(_BaseCredentials):
     """Base class for all credentials.

google/auth/iam.py

@@ -22,12 +22,14 @@
 import base64
 import http.client as http_client
 import json
+import os
 
 from google.auth import _exponential_backoff
 from google.auth import _helpers
 from google.auth import credentials
 from google.auth import crypt
 from google.auth import exceptions
+from google.auth.transport import mtls
 
 IAM_RETRY_CODES = {
     http_client.INTERNAL_SERVER_ERROR,
@@ -38,25 +40,33 @@
 
 _IAM_SCOPE = ["https://www.googleapis.com/auth/iam"]
 
-_IAM_ENDPOINT = (
-    "https://iamcredentials.googleapis.com/v1/projects/-"
-    + "/serviceAccounts/{}:generateAccessToken"
-)
-
-_IAM_SIGN_ENDPOINT = (
-    "https://iamcredentials.googleapis.com/v1/projects/-"
-    + "/serviceAccounts/{}:signBlob"
-)
-
-_IAM_SIGNJWT_ENDPOINT = (
-    "https://iamcredentials.googleapis.com/v1/projects/-"
-    + "/serviceAccounts/{}:signJwt"
-)
-
-_IAM_IDTOKEN_ENDPOINT = (
-    "https://iamcredentials.googleapis.com/v1/"
-    + "projects/-/serviceAccounts/{}:generateIdToken"
-)
+# 1. Determine if we should use mTLS.
+# Note: We only support automatic mTLS on the default googleapis.com universe.
+if hasattr(mtls, "should_use_client_cert"):
+    use_client_cert = mtls.should_use_client_cert()
+else:  # pragma: NO COVER
+    # if unsupported, fallback to reading from env var
+    use_client_cert = (
+        os.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false").lower() == "true"
+    )
+
+# 2. Construct the template domain using the library's DEFAULT_UNIVERSE_DOMAIN constant.
+# This ensures that the .replace() calls in the classes will work correctly.
+if use_client_cert:
+    # We use the .mtls. prefix only for the default universe template
+    _IAM_DOMAIN = f"iamcredentials.mtls.{credentials.DEFAULT_UNIVERSE_DOMAIN}"
+else:
+    _IAM_DOMAIN = f"iamcredentials.{credentials.DEFAULT_UNIVERSE_DOMAIN}"
+
+# 3. Create the common base URL template
+# We use double brackets {{}} so .format() can be called later for the email.
+_IAM_BASE_URL = f"https://{_IAM_DOMAIN}/v1/projects/-/serviceAccounts/{{}}"
+
+# 4. Define the endpoints as templates
+_IAM_ENDPOINT = _IAM_BASE_URL + ":generateAccessToken"
+_IAM_SIGN_ENDPOINT = _IAM_BASE_URL + ":signBlob"
+_IAM_SIGNJWT_ENDPOINT = _IAM_BASE_URL + ":signJwt"
+_IAM_IDTOKEN_ENDPOINT = _IAM_BASE_URL + ":generateIdToken"
 
 
 class Signer(crypt.Signer):

google/auth/transport/_mtls_helper.py

@@ -489,7 +489,7 @@ def check_parameters_for_unauthorized_response(cached_cert):
         str: The base64-encoded SHA256 cached fingerprint.
         str: The base64-encoded SHA256 current cert fingerprint.
     """
-    call_cert_bytes, call_key_bytes = _agent_identity_utils.call_client_cert_callback()
+    call_cert_bytes, call_key_bytes = call_client_cert_callback()
     cert_obj = _agent_identity_utils.parse_certificate(call_cert_bytes)
     current_cert_fingerprint = _agent_identity_utils.calculate_certificate_fingerprint(
         cert_obj
@@ -501,3 +501,11 @@ def check_parameters_for_unauthorized_response(cached_cert):
     else:
         cached_fingerprint = current_cert_fingerprint
     return call_cert_bytes, call_key_bytes, cached_fingerprint, current_cert_fingerprint
+
+
+def call_client_cert_callback():
+    """Calls the client cert callback and returns the certificate and key."""
+    _, cert_bytes, key_bytes, passphrase = get_client_ssl_credentials(
+        generate_encrypted_key=True
+    )
+    return cert_bytes, key_bytes

tests/compute_engine/test__metadata.py

@@ -201,8 +201,9 @@ def test_ping_success_custom_root(mock_metrics_header_value):
     )
 
 
+@mock.patch("time.sleep", return_value=None)
 @mock.patch("google.auth.metrics.mds_ping", return_value=MDS_PING_METRICS_HEADER_VALUE)
-def test_ping_failure_custom_retry(mock_metrics_header_value):
+def test_ping_failure_custom_retry(mock_metrics_header_value, _mock_sleep):
     request = make_request("")
     request.side_effect = exceptions.TransportError()
 
@@ -450,7 +451,8 @@ def test_get_failure_connection_failed(mock_sleep):
     assert request.call_count == 5
 
 
-def test_get_too_many_requests_retryable_error_failure():
+@mock.patch("time.sleep", return_value=None)
+def test_get_too_many_requests_retryable_error_failure(_mock_sleep):
     request = make_request("too many requests", status=http_client.TOO_MANY_REQUESTS)
 
     with pytest.raises(exceptions.TransportError) as excinfo:
@@ -546,7 +548,8 @@ def test_get_universe_domain_not_found():
     assert universe_domain == "googleapis.com"
 
 
-def test_get_universe_domain_retryable_error_failure():
+@mock.patch("time.sleep", return_value=None)
+def test_get_universe_domain_retryable_error_failure(_mock_sleep):
     # Test that if the universe domain endpoint returns a retryable error
     # we should retry.
     #

tests/test_agent_identity_utils.py

@@ -282,23 +282,6 @@ def test_get_agent_identity_certificate_path_fallback_to_well_known_path(
         mock_sleep.assert_called_once()
         assert mock_is_ready.call_count == 2
 
-    @mock.patch("google.auth.transport._mtls_helper.get_client_ssl_credentials")
-    def test_call_client_cert_callback(self, mock_get_client_ssl_credentials):
-        mock_get_client_ssl_credentials.return_value = (
-            True,
-            b"cert_bytes",
-            b"key_bytes",
-            b"passphrase",
-        )
-
-        cert, key = _agent_identity_utils.call_client_cert_callback()
-
-        assert cert == b"cert_bytes"
-        assert key == b"key_bytes"
-        mock_get_client_ssl_credentials.assert_called_once_with(
-            generate_encrypted_key=True
-        )
-
     def test_get_cached_cert_fingerprint_no_cert(self):
         with pytest.raises(ValueError, match="mTLS connection is not configured."):
             _agent_identity_utils.get_cached_cert_fingerprint(None)

tests/transport/test__mtls_helper.py

@@ -813,11 +813,12 @@ def test_no_env_vars_set(self):
 
 
 class TestMtlsHelper:
+    @mock.patch.object(_mtls_helper, "call_client_cert_callback")
     @mock.patch("google.auth.transport._mtls_helper._agent_identity_utils")
     def test_check_parameters_for_unauthorized_response_with_cached_cert(
-        self, mock_agent_identity_utils
+        self, mock_agent_identity_utils, mock_call_client_cert_callback
     ):
-        mock_agent_identity_utils.call_client_cert_callback.return_value = (
+        mock_call_client_cert_callback.return_value = (
             CERT_MOCK_VAL,
             KEY_MOCK_VAL,
         )
@@ -841,16 +842,17 @@ def test_check_parameters_for_unauthorized_response_with_cached_cert(
         assert key == KEY_MOCK_VAL
         assert cached_fingerprint == "cached_fingerprint"
         assert current_fingerprint == "current_fingerprint"
-        mock_agent_identity_utils.call_client_cert_callback.assert_called_once()
+        mock_call_client_cert_callback.assert_called_once()
         mock_agent_identity_utils.get_cached_cert_fingerprint.assert_called_once_with(
             b"cached_cert_bytes"
         )
 
-    @mock.patch("google.auth.transport._mtls_helper._agent_identity_utils")
+    @mock.patch.object(_mtls_helper, "call_client_cert_callback")
+    @mock.patch.object(_mtls_helper, "_agent_identity_utils")
     def test_check_parameters_for_unauthorized_response_without_cached_cert(
-        self, mock_agent_identity_utils
+        self, mock_agent_identity_utils, mock_call_client_cert_callback
     ):
-        mock_agent_identity_utils.call_client_cert_callback.return_value = (
+        mock_call_client_cert_callback.return_value = (
             CERT_MOCK_VAL,
             KEY_MOCK_VAL,
         )
@@ -869,5 +871,22 @@ def test_check_parameters_for_unauthorized_response_without_cached_cert(
         assert key == KEY_MOCK_VAL
         assert cached_fingerprint == "current_fingerprint"
         assert current_fingerprint == "current_fingerprint"
-        mock_agent_identity_utils.call_client_cert_callback.assert_called_once()
+        mock_call_client_cert_callback.assert_called_once()
         mock_agent_identity_utils.get_cached_cert_fingerprint.assert_not_called()
+
+    @mock.patch("google.auth.transport._mtls_helper.get_client_ssl_credentials")
+    def test_call_client_cert_callback(self, mock_get_client_ssl_credentials):
+        mock_get_client_ssl_credentials.return_value = (
+            True,
+            b"cert_bytes",
+            b"key_bytes",
+            b"passphrase",
+        )
+
+        cert, key = _mtls_helper.call_client_cert_callback()
+
+        assert cert == b"cert_bytes"
+        assert key == b"key_bytes"
+        mock_get_client_ssl_credentials.assert_called_once_with(
+            generate_encrypted_key=True
+        )

tests/transport/test_requests.py

@@ -566,7 +566,7 @@ def test_cert_rotation_when_cert_mismatch_and_mtls_enabled(self):
 
         # Mock call_client_cert_callback to return the new certificate.
         with mock.patch.object(
-            google.auth.transport._mtls_helper._agent_identity_utils,
+            google.auth.transport._mtls_helper,
             "call_client_cert_callback",
             return_value=(new_cert, new_key),
         ) as mock_callback:
@@ -605,7 +605,7 @@ def test_no_cert_rotation_when_cert_match_and_mTLS_enabled(self):
 
         # Mock call_client_cert_callback to return the new certificate.
         with mock.patch.object(
-            google.auth.transport._mtls_helper._agent_identity_utils,
+            google.auth.transport._mtls_helper,
             "call_client_cert_callback",
             return_value=(new_cert, new_key),
         ):
@@ -638,7 +638,7 @@ def test_no_cert_match_check_when_mtls_disabled(self):
 
         # Mock call_client_cert_callback to return the new certificate.
         with mock.patch.object(
-            google.auth.transport._mtls_helper._agent_identity_utils,
+            google.auth.transport._mtls_helper,
             "call_client_cert_callback",
             return_value=(new_cert, new_key),
         ) as mock_callback:
@@ -688,7 +688,7 @@ def test_cert_rotation_failure_raises_error(self):
         authed_session._is_mtls = True
 
         with mock.patch.object(
-            google.auth.transport._mtls_helper._agent_identity_utils,
+            google.auth.transport._mtls_helper,
             "call_client_cert_callback",
             return_value=(new_cert, new_key),
         ):

tests/transport/test_urllib3.py

@@ -343,7 +343,7 @@ def test_cert_rotation_when_cert_mismatch_and_mtls_endpoint_used(self):
         authed_http._is_mtls = True
         # Mock call_client_cert_callback to return the new certificate.
         with mock.patch.object(
-            google.auth._agent_identity_utils,
+            google.auth.transport._mtls_helper,
             "call_client_cert_callback",
             return_value=(new_cert, new_key),
         ) as mock_callback:
@@ -378,7 +378,7 @@ def test_no_cert_rotation_when_cert_match_and_mtls_endpoint_used(self):
         authed_http._is_mtls = True
         # Mock call_client_cert_callback to return the certificate.
         with mock.patch.object(
-            google.auth._agent_identity_utils,
+            google.auth.transport._mtls_helper,
             "call_client_cert_callback",
             return_value=(new_cert, new_key),
         ):
@@ -408,7 +408,7 @@ def test_no_cert_match_check_when_mtls_endpoint_not_used(self):
 
         # Mock call_client_cert_callback to return the certificate.
         with mock.patch.object(
-            google.auth._agent_identity_utils,
+            google.auth.transport._mtls_helper,
             "call_client_cert_callback",
             return_value=(new_cert, new_key),
         ) as mock_callback: