diff --git a/packages/google-auth/google/auth/_agent_identity_utils.py b/packages/google-auth/google/auth/_agent_identity_utils.py index 299c56ccbfde..8a1eddbe1cd3 100644 --- a/packages/google-auth/google/auth/_agent_identity_utils.py +++ b/packages/google-auth/google/auth/_agent_identity_utils.py @@ -22,9 +22,7 @@ import time from urllib.parse import quote, urlparse -from google.auth import environment_vars -from google.auth import exceptions - +from google.auth import environment_vars, exceptions _LOGGER = logging.getLogger(__name__) @@ -37,6 +35,8 @@ _AGENT_IDENTITY_SPIFFE_TRUST_DOMAIN_PATTERNS = [ r"^agents\.global\.org-\d+\.system\.id\.goog$", r"^agents\.global\.proj-\d+\.system\.id\.goog$", + r"^agents-nonprod\.global\.org-\d+\.system\.id\.goog$", + r"^agents-nonprod\.global\.proj-\d+\.system\.id\.goog$", ] _WELL_KNOWN_CERT_PATH = "/var/run/secrets/workload-spiffe-credentials/certificates.pem" diff --git a/packages/google-auth/tests/test_agent_identity_utils.py b/packages/google-auth/tests/test_agent_identity_utils.py index b2e690c7465a..f74bdad9e475 100644 --- a/packages/google-auth/tests/test_agent_identity_utils.py +++ b/packages/google-auth/tests/test_agent_identity_utils.py @@ -21,9 +21,7 @@ from cryptography import x509 import pytest -from google.auth import _agent_identity_utils -from google.auth import environment_vars -from google.auth import exceptions +from google.auth import _agent_identity_utils, environment_vars, exceptions # A mock PEM-encoded certificate without an Agent Identity SPIFFE ID. NON_AGENT_IDENTITY_CERT_BYTES = ( @@ -60,15 +58,22 @@ def test__is_agent_identity_certificate_invalid(self): cert = _agent_identity_utils.parse_certificate(NON_AGENT_IDENTITY_CERT_BYTES) assert not _agent_identity_utils._is_agent_identity_certificate(cert) - def test__is_agent_identity_certificate_valid_spiffe(self): + @pytest.mark.parametrize( + "spiffe_id", + [ + "spiffe://agents.global.proj-12345.system.id.goog/workload", + "spiffe://agents.global.org-12345.system.id.goog/workload", + "spiffe://agents-nonprod.global.proj-12345.system.id.goog/workload", + "spiffe://agents-nonprod.global.org-12345.system.id.goog/workload", + ], + ) + def test__is_agent_identity_certificate_valid_spiffe(self, spiffe_id): mock_cert = mock.MagicMock() mock_ext = mock.MagicMock() mock_san_value = mock.MagicMock() mock_cert.extensions.get_extension_for_oid.return_value = mock_ext mock_ext.value = mock_san_value - mock_san_value.get_values_for_type.return_value = [ - "spiffe://agents.global.proj-12345.system.id.goog/workload" - ] + mock_san_value.get_values_for_type.return_value = [spiffe_id] assert _agent_identity_utils._is_agent_identity_certificate(mock_cert) def test__is_agent_identity_certificate_non_matching_spiffe(self):