feat(auth): Implement python mtls helpers

packages/google-auth/google/auth/environment_vars.py

@@ -133,3 +133,6 @@
     "GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES"
 )
 """Environment variable to prevent agent token sharing for GCP services."""
+
+GOOGLE_API_USE_MTLS_ENDPOINT = "GOOGLE_API_USE_MTLS_ENDPOINT"
+"""Environment variable controlling whether to use mTLS endpoint or not."""

packages/google-auth/google/auth/transport/mtls.py

@@ -14,12 +14,19 @@
 
 """Utilites for mutual TLS."""
 
+import logging
 from os import getenv
+import ssl
+from typing import Optional
 
+from google.auth import environment_vars
 from google.auth import exceptions
 from google.auth.transport import _mtls_helper
 
 
+_LOGGER = logging.getLogger(__name__)
+
+
 def has_default_client_cert_source(include_context_aware=True):
     """Check if default client SSL credentials exists on the device.
 
@@ -60,7 +67,7 @@ def default_client_cert_source():
             client certificate bytes and private key bytes, both in PEM format.
 
     Raises:
-        google.auth.exceptions.DefaultClientCertSourceError: If the default
+        google.auth.exceptions.MutualTLSChannelError: If the default
             client SSL credentials don't exist or are malformed.
     """
     if not has_default_client_cert_source(include_context_aware=True):
@@ -71,7 +78,12 @@ def default_client_cert_source():
     def callback():
         try:
             _, cert_bytes, key_bytes = _mtls_helper.get_client_cert_and_key()
-        except (OSError, RuntimeError, ValueError) as caught_exc:
+        except (
+            exceptions.ClientCertError,
+            OSError,
+            RuntimeError,
+            ValueError,
+        ) as caught_exc:
             new_exc = exceptions.MutualTLSChannelError(caught_exc)
             raise new_exc from caught_exc
 
@@ -96,7 +108,7 @@ def default_client_encrypted_cert_source(cert_path, key_path):
             returns the cert_path, key_path and passphrase bytes.
 
     Raises:
-        google.auth.exceptions.DefaultClientCertSourceError: If any problem
+        google.auth.exceptions.MutualTLSChannelError: If any problem
             occurs when loading or saving the client certificate and key.
     """
     if not has_default_client_cert_source(include_context_aware=True):
@@ -140,3 +152,158 @@ def should_use_client_cert():
        bool: indicating whether the client certificate should be used for mTLS.
     """
     return _mtls_helper.check_use_client_cert()
+
+
+def _load_client_cert_into_context(
+    ctx: ssl.SSLContext,
+    cert_bytes: bytes,
+    key_bytes: bytes,
+    passphrase: Optional[bytes] = None,
+) -> None:
+    """Load a client certificate and key into an SSL context.
+
+    Args:
+        ctx (ssl.SSLContext): The SSL context to load the certificate and key into.
+        cert_bytes (bytes): The client certificate bytes in PEM format.
+        key_bytes (bytes): The client private key bytes in PEM format.
+        passphrase (Optional[bytes]): The passphrase for the client private key.
+
+    Raises:
+        google.auth.exceptions.MutualTLSChannelError: If the SSL context is invalid,
+            or if loading the certificate and key fails.
+    """
+    if not isinstance(ctx, ssl.SSLContext):
+        raise exceptions.MutualTLSChannelError(
+            "Failed to load client certificate and key for mTLS. The provided context "
+            "object is invalid or does not support loading certificate chains."
+        )
+
+    try:
+        with _mtls_helper.secure_cert_key_paths(
+            cert_bytes, key_bytes, passphrase=passphrase
+        ) as (
+            cert_path,
+            key_path,
+            passphrase_val,
+        ):
+            password_val = (
+                passphrase_val.decode("utf-8")
+                if isinstance(passphrase_val, bytes)
+                else passphrase_val
+            )
+            if cert_path is None or key_path is None:
+                raise exceptions.MutualTLSChannelError(
+                    "Failed to generate temporary file paths for the client certificate and key."
+                )
+            ctx.load_cert_chain(
+                certfile=cert_path, keyfile=key_path, password=password_val
+            )
+    except (
+        ssl.SSLError,
+        OSError,
+        ValueError,
+        RuntimeError,
+        TypeError,
+    ) as caught_exc:
+        new_exc = exceptions.MutualTLSChannelError(caught_exc)
+        raise new_exc from caught_exc
+
+
+def load_default_client_cert(ctx: ssl.SSLContext) -> bool:
+    """Load the default client certificate and key into an SSL context if configured.
+
+    If client certificates are enabled and a default client certificate source is
+    found, the certificate and key are loaded into the SSL context.
+
+    Args:
+        ctx (ssl.SSLContext): The SSL context to load the default client certificate
+            and key into.
+
+    Returns:
+        bool: True if client certificates are enabled and the default client
+            certificate was successfully loaded. False if client certificates
+            are disabled or if no default certificate source is configured.
+
+    Raises:
+        google.auth.exceptions.ClientCertError: If the default client certificate
+            source exists but cannot be loaded or parsed.
+        google.auth.exceptions.MutualTLSChannelError: If the default client certificate
+            or key is malformed.
+    """
+    if not should_use_client_cert() or not has_default_client_cert_source():
+        return False
+    try:
+        (
+            has_cert,
+            cert_bytes,
+            key_bytes,
+            passphrase,
+        ) = _mtls_helper.get_client_ssl_credentials()
+    except (
+        exceptions.ClientCertError,
+        OSError,
+        RuntimeError,
+        ValueError,
+    ) as caught_exc:
+        new_exc = exceptions.MutualTLSChannelError(caught_exc)
+        raise new_exc from caught_exc
+    if not has_cert:
+        return False
+    _load_client_cert_into_context(ctx, cert_bytes, key_bytes, passphrase)
+    return True
+
+
+def get_default_ssl_context() -> Optional[ssl.SSLContext]:
+    """Get a default SSL context loaded with the default client certificate.
+
+    Returns:
+        ssl.SSLContext: An SSL context loaded with the default client
+            certificate, or None if client certificates are not configured
+            or available.
+
+    Raises:
+        google.auth.exceptions.ClientCertError: If the default client certificate
+            source exists but cannot be loaded or parsed.
+        google.auth.exceptions.MutualTLSChannelError: If the default client certificate
+            or key is malformed.
+    """
+    if not should_use_client_cert() or not has_default_client_cert_source():
+        return None
+
+    ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
+    return ctx if load_default_client_cert(ctx) else None
+
+
+def should_use_mtls_endpoint(
+    client_cert_available: Optional[bool] = None,
+) -> bool:
+    """Determine whether to use an mTLS endpoint.
+
+    This relies on the GOOGLE_API_USE_MTLS_ENDPOINT environment variable. If set to
+    "always", returns True. If set to "never", returns False. If set to "auto"
+    or unset, returns whether a client certificate is available.
+
+    Args:
+        client_cert_available (Optional[bool]): indicating if a client certificate
+            is available. If None, this is determined by checking if client
+            certificates are enabled using :func:`should_use_client_cert`.
+
+    Returns:
+        bool: indicating if an mTLS endpoint should be used.
+    """
+    if client_cert_available is None:
+        client_cert_available = should_use_client_cert()
+
+    use_mtls_endpoint = getenv(environment_vars.GOOGLE_API_USE_MTLS_ENDPOINT, "auto")
+    use_mtls_endpoint = (use_mtls_endpoint or "auto").strip().lower()
+    if use_mtls_endpoint == "always":
+        return True
+    if use_mtls_endpoint == "never":
+        return False
+    if use_mtls_endpoint == "auto":
+        return client_cert_available
+
+    raise exceptions.MutualTLSChannelError(
+        f"Unsupported {environment_vars.GOOGLE_API_USE_MTLS_ENDPOINT} value "
+        f"'{use_mtls_endpoint}'. Accepted values: never, auto, always."
+    )

packages/google-auth/tests/transport/test__mtls_helper.py

@@ -1204,6 +1204,31 @@ def test_falls_back_to_tempfile_when_memfd_fails(
             assert key_path == "/tmp/key"
             assert passphrase == b"new_pass"
 
+    @mock.patch.object(sys, "platform", "linux")
+    @mock.patch.object(os, "memfd_create", create=True)
+    @mock.patch.object(_mtls_helper, "_memfd_cert_key_paths", autospec=True)
+    @mock.patch.object(_mtls_helper, "_tempfile_cert_key_paths", autospec=True)
+    def test_tier3_fallback_failure_propagation(
+        self, mock_tempfile_cm, mock_memfd_cm, mock_memfd_create
+    ):
+        mock_memfd_ctx = mock.MagicMock()
+        mock_memfd_ctx.__enter__.side_effect = _mtls_helper._MemfdCreationError(
+            "memfd failed"
+        )
+        mock_memfd_cm.return_value = mock_memfd_ctx
+
+        mock_tempfile_ctx = mock.MagicMock()
+        mock_tempfile_ctx.__enter__.side_effect = OSError("tempfile failed")
+        mock_tempfile_cm.return_value = mock_tempfile_ctx
+
+        with pytest.raises(OSError) as exc_info:
+            with _mtls_helper.secure_cert_key_paths(
+                pytest.public_cert_bytes, pytest.private_key_bytes, b"passphrase"
+            ):
+                pass
+
+        assert "tempfile failed" in str(exc_info.value)
+
     @mock.patch.object(sys, "platform", "darwin")
     @mock.patch.object(_mtls_helper, "_tempfile_cert_key_paths", autospec=True)
     def test_uses_tempfile_directly_on_unsupported_os(self, mock_tempfile_cm):
@@ -1378,6 +1403,43 @@ def _redirect_mkstemp(dir=None):
             # Verify cert file is still cleaned up even if key cleanup raised KeyboardInterrupt
             assert not os.path.exists(cert_path)
 
+    @mock.patch.object(os, "access", return_value=True)
+    @mock.patch.object(os.path, "isdir", return_value=True)
+    @mock.patch.object(_mtls_helper, "_encrypt_key_if_plaintext", autospec=True)
+    def test_cleanup_on_key_write_error(
+        self, mock_encrypt, mock_isdir, mock_access, tmpdir
+    ):
+        original_mkstemp = tempfile.mkstemp
+
+        cert_path_ref = []
+
+        def _redirect_mkstemp(dir=None):
+            fd, path = original_mkstemp(dir=str(tmpdir))
+            cert_path_ref.append(path)
+            return fd, path
+
+        with mock.patch.object(tempfile, "mkstemp", side_effect=_redirect_mkstemp):
+            mock_encrypt.return_value = (b"encrypted_key", b"new_pass")
+
+            original_fdopen = os.fdopen
+            call_count = [0]
+
+            def _failing_fdopen(fd, *args, **kwargs):
+                call_count[0] += 1
+                if call_count[0] == 2:
+                    raise OSError("key write failed")
+                return original_fdopen(fd, *args, **kwargs)
+
+            with pytest.raises(OSError):
+                with mock.patch("os.fdopen", side_effect=_failing_fdopen):
+                    with _mtls_helper._tempfile_cert_key_paths(
+                        b"cert", b"key", b"pass"
+                    ):
+                        pass
+
+            assert len(cert_path_ref) > 0
+            assert not os.path.exists(cert_path_ref[0])
+
     @mock.patch.object(os, "access", return_value=True)
     @mock.patch.object(os.path, "isdir", return_value=True)
     @mock.patch.object(_mtls_helper, "_encrypt_key_if_plaintext", autospec=True)
@@ -1472,6 +1534,21 @@ def test_encrypts_plaintext_key(self):
         )
         assert decrypted
 
+        original_key = serialization.load_pem_private_key(
+            pytest.private_key_bytes, password=None
+        )
+        decrypted_bytes = decrypted.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+        original_bytes = original_key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+        assert decrypted_bytes == original_bytes
+
     @mock.patch("secrets.token_hex", return_value="0123456789abcdef0123456789abcdef")
     def test_default_passphrase_generation(self, mock_secrets):
         encrypted_bytes, passphrase = _mtls_helper._encrypt_key_if_plaintext(
@@ -1495,6 +1572,24 @@ def test_encrypts_plaintext_key_string_passphrase(self):
         assert encrypted_bytes != pytest.private_key_bytes
         assert b"ENCRYPTED PRIVATE KEY" in encrypted_bytes
 
+        decrypted = serialization.load_pem_private_key(
+            encrypted_bytes, password=b"my_passphrase_str"
+        )
+        original_key = serialization.load_pem_private_key(
+            pytest.private_key_bytes, password=None
+        )
+        decrypted_bytes = decrypted.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+        original_bytes = original_key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+        assert decrypted_bytes == original_bytes
+
     def test_returns_unsupported_algorithm_asis(self):
         import cryptography.exceptions
 
@@ -1537,6 +1632,7 @@ def test_success(
 
         mock_open.assert_called_once_with("/path/to/secret", "r+b")
         mock_fh.write.assert_called_once_with(b"\0" * 10)
+        mock_fh.flush.assert_called_once()
         mock_fsync.assert_called_once()
         mock_remove.assert_called_once_with("/path/to/secret")
 
@@ -1583,5 +1679,6 @@ def test_remove_oserror_ignored(
         _mtls_helper._secure_wipe_and_remove("/path/to/secret")
 
         mock_open.assert_called_once_with("/path/to/secret", "r+b")
+        mock_fh.flush.assert_called_once()
         mock_fsync.assert_called_once()
         mock_remove.assert_called_once_with("/path/to/secret")