From a4d432a9e62a5e759e94168463795cd0e7d7ab72 Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Sat, 14 Jun 2025 13:39:14 -0700 Subject: [PATCH] chore: Add Service Account Credential Exchanger (Experimental) PiperOrigin-RevId: 771507089 --- .../service_account_credential_exchanger.py | 92 +++++ ...st_service_account_credential_exchanger.py | 341 ++++++++++++++++++ 2 files changed, 433 insertions(+) create mode 100644 src/google/adk/auth/service_account_credential_exchanger.py create mode 100644 tests/unittests/auth/test_service_account_credential_exchanger.py diff --git a/src/google/adk/auth/service_account_credential_exchanger.py b/src/google/adk/auth/service_account_credential_exchanger.py new file mode 100644 index 00000000..644501ee --- /dev/null +++ b/src/google/adk/auth/service_account_credential_exchanger.py @@ -0,0 +1,92 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Credential fetcher for Google Service Account.""" + +from __future__ import annotations + +import google.auth +from google.auth.transport.requests import Request +from google.oauth2 import service_account + +from ..utils.feature_decorator import experimental +from .auth_credential import AuthCredential +from .auth_credential import AuthCredentialTypes +from .auth_credential import HttpAuth +from .auth_credential import HttpCredentials + + +@experimental +class ServiceAccountCredentialExchanger: + """Exchanges Google Service Account credentials for an access token. + + Uses the default service credential if `use_default_credential = True`. + Otherwise, uses the service account credential provided in the auth + credential. + """ + + def __init__(self, credential: AuthCredential): + if credential.auth_type != AuthCredentialTypes.SERVICE_ACCOUNT: + raise ValueError("Credential is not a service account credential.") + self._credential = credential + + def exchange(self) -> AuthCredential: + """Exchanges the service account auth credential for an access token. + + If the AuthCredential contains a service account credential, it will be used + to exchange for an access token. Otherwise, if use_default_credential is True, + the default application credential will be used for exchanging an access token. + + Returns: + An AuthCredential in HTTP Bearer format, containing the access token. + + Raises: + ValueError: If service account credentials are missing or invalid. + Exception: If credential exchange or refresh fails. + """ + if ( + self._credential is None + or self._credential.service_account is None + or ( + self._credential.service_account.service_account_credential is None + and not self._credential.service_account.use_default_credential + ) + ): + raise ValueError( + "Service account credentials are missing. Please provide them, or set" + " `use_default_credential = True` to use application default" + " credential in a hosted service like Google Cloud Run." + ) + + try: + if self._credential.service_account.use_default_credential: + credentials, _ = google.auth.default() + else: + config = self._credential.service_account + credentials = service_account.Credentials.from_service_account_info( + config.service_account_credential.model_dump(), scopes=config.scopes + ) + + # Refresh credentials to ensure we have a valid access token + credentials.refresh(Request()) + + return AuthCredential( + auth_type=AuthCredentialTypes.HTTP, + http=HttpAuth( + scheme="bearer", + credentials=HttpCredentials(token=credentials.token), + ), + ) + except Exception as e: + raise ValueError(f"Failed to exchange service account token: {e}") from e diff --git a/tests/unittests/auth/test_service_account_credential_exchanger.py b/tests/unittests/auth/test_service_account_credential_exchanger.py new file mode 100644 index 00000000..a5c66843 --- /dev/null +++ b/tests/unittests/auth/test_service_account_credential_exchanger.py @@ -0,0 +1,341 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the ServiceAccountCredentialExchanger.""" + +from unittest.mock import MagicMock +from unittest.mock import patch + +from google.adk.auth.auth_credential import AuthCredential +from google.adk.auth.auth_credential import AuthCredentialTypes +from google.adk.auth.auth_credential import ServiceAccount +from google.adk.auth.auth_credential import ServiceAccountCredential +from google.adk.auth.service_account_credential_exchanger import ServiceAccountCredentialExchanger +import pytest + + +class TestServiceAccountCredentialExchanger: + """Test cases for ServiceAccountCredentialExchanger.""" + + def test_init_valid_credential(self): + """Test successful initialization with valid service account credential.""" + credential = AuthCredential( + auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, + service_account=ServiceAccount( + service_account_credential=ServiceAccountCredential( + type_="service_account", + project_id="test-project", + private_key_id="key-id", + private_key=( + "-----BEGIN PRIVATE KEY-----\nMOCK_KEY\n-----END PRIVATE" + " KEY-----" + ), + client_email="test@test-project.iam.gserviceaccount.com", + client_id="12345", + auth_uri="https://accounts.google.com/o/oauth2/auth", + token_uri="https://oauth2.googleapis.com/token", + auth_provider_x509_cert_url=( + "https://www.googleapis.com/oauth2/v1/certs" + ), + client_x509_cert_url="https://www.googleapis.com/robot/v1/metadata/x509/test%40test-project.iam.gserviceaccount.com", + universe_domain="googleapis.com", + ), + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ), + ) + + exchanger = ServiceAccountCredentialExchanger(credential) + assert exchanger._credential == credential + + def test_init_invalid_credential_type(self): + """Test initialization with invalid credential type raises ValueError.""" + credential = AuthCredential( + auth_type=AuthCredentialTypes.API_KEY, + api_key="test-key", + ) + + with pytest.raises( + ValueError, match="Credential is not a service account credential" + ): + ServiceAccountCredentialExchanger(credential) + + @patch( + "google.adk.auth.service_account_credential_exchanger.service_account.Credentials.from_service_account_info" + ) + @patch("google.adk.auth.service_account_credential_exchanger.Request") + def test_exchange_with_explicit_credentials_success( + self, mock_request_class, mock_from_service_account_info + ): + """Test successful exchange with explicit service account credentials.""" + # Setup mocks + mock_request = MagicMock() + mock_request_class.return_value = mock_request + + mock_credentials = MagicMock() + mock_credentials.token = "mock_access_token" + mock_from_service_account_info.return_value = mock_credentials + + # Create test credential + service_account_cred = ServiceAccountCredential( + type_="service_account", + project_id="test-project", + private_key_id="key-id", + private_key=( + "-----BEGIN PRIVATE KEY-----\nMOCK_KEY\n-----END PRIVATE KEY-----" + ), + client_email="test@test-project.iam.gserviceaccount.com", + client_id="12345", + auth_uri="https://accounts.google.com/o/oauth2/auth", + token_uri="https://oauth2.googleapis.com/token", + auth_provider_x509_cert_url=( + "https://www.googleapis.com/oauth2/v1/certs" + ), + client_x509_cert_url="https://www.googleapis.com/robot/v1/metadata/x509/test%40test-project.iam.gserviceaccount.com", + universe_domain="googleapis.com", + ) + + credential = AuthCredential( + auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, + service_account=ServiceAccount( + service_account_credential=service_account_cred, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ), + ) + + exchanger = ServiceAccountCredentialExchanger(credential) + result = exchanger.exchange() + + # Verify the result + assert result.auth_type == AuthCredentialTypes.HTTP + assert result.http.scheme == "bearer" + assert result.http.credentials.token == "mock_access_token" + + # Verify mocks were called correctly + mock_from_service_account_info.assert_called_once_with( + service_account_cred.model_dump(), + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ) + mock_credentials.refresh.assert_called_once_with(mock_request) + + @patch( + "google.adk.auth.service_account_credential_exchanger.google.auth.default" + ) + @patch("google.adk.auth.service_account_credential_exchanger.Request") + def test_exchange_with_default_credentials_success( + self, mock_request_class, mock_google_auth_default + ): + """Test successful exchange with default application credentials.""" + # Setup mocks + mock_request = MagicMock() + mock_request_class.return_value = mock_request + + mock_credentials = MagicMock() + mock_credentials.token = "default_access_token" + mock_google_auth_default.return_value = (mock_credentials, "test-project") + + # Create test credential with use_default_credential=True + credential = AuthCredential( + auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, + service_account=ServiceAccount( + use_default_credential=True, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ), + ) + + exchanger = ServiceAccountCredentialExchanger(credential) + result = exchanger.exchange() + + # Verify the result + assert result.auth_type == AuthCredentialTypes.HTTP + assert result.http.scheme == "bearer" + assert result.http.credentials.token == "default_access_token" + + # Verify mocks were called correctly + mock_google_auth_default.assert_called_once() + mock_credentials.refresh.assert_called_once_with(mock_request) + + def test_exchange_missing_service_account(self): + """Test exchange fails when service_account is None.""" + credential = AuthCredential( + auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, + service_account=None, + ) + + exchanger = ServiceAccountCredentialExchanger(credential) + + with pytest.raises( + ValueError, match="Service account credentials are missing" + ): + exchanger.exchange() + + def test_exchange_missing_credentials_and_not_default(self): + """Test exchange fails when credentials are missing and use_default_credential is False.""" + credential = AuthCredential( + auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, + service_account=ServiceAccount( + service_account_credential=None, + use_default_credential=False, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ), + ) + + exchanger = ServiceAccountCredentialExchanger(credential) + + with pytest.raises( + ValueError, match="Service account credentials are missing" + ): + exchanger.exchange() + + @patch( + "google.adk.auth.service_account_credential_exchanger.service_account.Credentials.from_service_account_info" + ) + def test_exchange_credential_creation_failure( + self, mock_from_service_account_info + ): + """Test exchange handles credential creation failure gracefully.""" + # Setup mock to raise exception + mock_from_service_account_info.side_effect = Exception( + "Invalid private key" + ) + + # Create test credential + service_account_cred = ServiceAccountCredential( + type_="service_account", + project_id="test-project", + private_key_id="key-id", + private_key="invalid-key", + client_email="test@test-project.iam.gserviceaccount.com", + client_id="12345", + auth_uri="https://accounts.google.com/o/oauth2/auth", + token_uri="https://oauth2.googleapis.com/token", + auth_provider_x509_cert_url=( + "https://www.googleapis.com/oauth2/v1/certs" + ), + client_x509_cert_url="https://www.googleapis.com/robot/v1/metadata/x509/test%40test-project.iam.gserviceaccount.com", + universe_domain="googleapis.com", + ) + + credential = AuthCredential( + auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, + service_account=ServiceAccount( + service_account_credential=service_account_cred, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ), + ) + + exchanger = ServiceAccountCredentialExchanger(credential) + + with pytest.raises( + ValueError, match="Failed to exchange service account token" + ): + exchanger.exchange() + + @patch( + "google.adk.auth.service_account_credential_exchanger.google.auth.default" + ) + def test_exchange_default_credential_failure(self, mock_google_auth_default): + """Test exchange handles default credential failure gracefully.""" + # Setup mock to raise exception + mock_google_auth_default.side_effect = Exception( + "No default credentials found" + ) + + # Create test credential with use_default_credential=True + credential = AuthCredential( + auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, + service_account=ServiceAccount( + use_default_credential=True, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ), + ) + + exchanger = ServiceAccountCredentialExchanger(credential) + + with pytest.raises( + ValueError, match="Failed to exchange service account token" + ): + exchanger.exchange() + + @patch( + "google.adk.auth.service_account_credential_exchanger.service_account.Credentials.from_service_account_info" + ) + @patch("google.adk.auth.service_account_credential_exchanger.Request") + def test_exchange_refresh_failure( + self, mock_request_class, mock_from_service_account_info + ): + """Test exchange handles credential refresh failure gracefully.""" + # Setup mocks + mock_request = MagicMock() + mock_request_class.return_value = mock_request + + mock_credentials = MagicMock() + mock_credentials.refresh.side_effect = Exception( + "Network error during refresh" + ) + mock_from_service_account_info.return_value = mock_credentials + + # Create test credential + service_account_cred = ServiceAccountCredential( + type_="service_account", + project_id="test-project", + private_key_id="key-id", + private_key=( + "-----BEGIN PRIVATE KEY-----\nMOCK_KEY\n-----END PRIVATE KEY-----" + ), + client_email="test@test-project.iam.gserviceaccount.com", + client_id="12345", + auth_uri="https://accounts.google.com/o/oauth2/auth", + token_uri="https://oauth2.googleapis.com/token", + auth_provider_x509_cert_url=( + "https://www.googleapis.com/oauth2/v1/certs" + ), + client_x509_cert_url="https://www.googleapis.com/robot/v1/metadata/x509/test%40test-project.iam.gserviceaccount.com", + universe_domain="googleapis.com", + ) + + credential = AuthCredential( + auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, + service_account=ServiceAccount( + service_account_credential=service_account_cred, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ), + ) + + exchanger = ServiceAccountCredentialExchanger(credential) + + with pytest.raises( + ValueError, match="Failed to exchange service account token" + ): + exchanger.exchange() + + def test_exchange_none_credential_in_constructor(self): + """Test that passing None credential raises appropriate error during construction.""" + # This test verifies behavior when _credential is None, though this shouldn't + # happen in normal usage due to constructor validation + credential = AuthCredential( + auth_type=AuthCredentialTypes.SERVICE_ACCOUNT, + service_account=ServiceAccount( + use_default_credential=True, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ), + ) + + exchanger = ServiceAccountCredentialExchanger(credential) + # Manually set to None to test the validation logic + exchanger._credential = None + + with pytest.raises( + ValueError, match="Service account credentials are missing" + ): + exchanger.exchange()