import json from pathlib import Path import sys import tempfile import unittest from unittest import mock ROOT = Path(__file__).resolve().parents[1] APP_DIR = ROOT / 'app' sys.path.insert(0, str(APP_DIR)) from keycheckers.gcp import gcpKeycheck import keycheck_runner from runtime_security import ensure_private_directory, harden_private_file class GCPTokenEndpointTests(unittest.TestCase): def test_adc_rejects_local_and_untrusted_token_uri_without_request(self): base = { 'client_id': 'client-id', 'client_secret': 'client-secret', 'refresh_token': 'refresh-token', } for token_uri in ( 'http://127.0.0.1:8080/token', 'https://evil.example/token', 'https://oauth2.googleapis.com/token?next=https://evil.example', ): with self.subTest(token_uri=token_uri), mock.patch.object(gcpKeycheck.requests, 'post') as post: result = gcpKeycheck.check_adc({**base, 'token_uri': token_uri}, None, 5) self.assertEqual(result['status'], 'UNKNOWN') self.assertIn('untrusted', result['message']) post.assert_not_called() def test_service_account_rejects_untrusted_uri_before_signing_or_request(self): credentials = { 'client_email': 'fixture@example.invalid', 'private_key_id': 'fixture', 'private_key': 'not-read-for-an-untrusted-uri', 'token_uri': 'https://evil.example/token', } with mock.patch.object(gcpKeycheck, 'rsa_pkcs1v15_sha256_sign') as sign, \ mock.patch.object(gcpKeycheck.requests, 'post') as post: result = gcpKeycheck.check_service_account(credentials, None, 5) self.assertEqual(result['status'], 'UNKNOWN') sign.assert_not_called() post.assert_not_called() def test_trusted_endpoint_is_posted_without_redirect_following(self): response = mock.Mock(status_code=400) response.text = 'invalid_grant' response.json.return_value = {'error': 'invalid_grant'} credentials = { 'client_id': 'client-id', 'client_secret': 'client-secret', 'refresh_token': 'refresh-token', 'token_uri': gcpKeycheck.GOOGLE_TOKEN_URL, } with mock.patch.object(gcpKeycheck.requests, 'post', return_value=response) as post: result = gcpKeycheck.check_adc(credentials, None, 5) self.assertEqual(result['status'], 'DEAD') self.assertEqual(post.call_args.args[0], gcpKeycheck.GOOGLE_TOKEN_URL) self.assertFalse(post.call_args.kwargs['allow_redirects']) def test_redirect_is_not_followed_or_forwarded(self): response = mock.Mock(status_code=302, headers={'Location': 'https://evil.example/token'}) credentials = { 'client_id': 'client-id', 'client_secret': 'client-secret', 'refresh_token': 'refresh-token', } with mock.patch.object(gcpKeycheck.requests, 'post', return_value=response) as post: result = gcpKeycheck.check_adc(credentials, None, 5) self.assertEqual(result['status'], 'UNKNOWN') self.assertIn('redirect refused', result['message']) self.assertEqual(post.call_count, 1) self.assertFalse(post.call_args.kwargs['allow_redirects']) class GCPVertexEndpointTests(unittest.TestCase): def test_vertex_api_host_supports_global_multi_region_and_regional_locations(self): self.assertEqual(gcpKeycheck.vertex_api_host('global'), 'aiplatform.googleapis.com') self.assertEqual(gcpKeycheck.vertex_api_host('us'), 'aiplatform.us.rep.googleapis.com') self.assertEqual(gcpKeycheck.vertex_api_host('eu'), 'aiplatform.eu.rep.googleapis.com') self.assertEqual(gcpKeycheck.vertex_api_host('us-east5'), 'us-east5-aiplatform.googleapis.com') def test_gemini_probe_uses_global_and_multi_region_hosts(self): response = mock.Mock(status_code=200) response.json.return_value = {'totalTokens': 1} with mock.patch.object(gcpKeycheck.requests, 'post', return_value=response) as post: result = gcpKeycheck.probe_vertex_llm( 'token', 'project', None, 6, locations=['global', 'us', 'eu'], models=['gemini-3.6-flash'], max_attempts=3, ) self.assertTrue(result['enabled']) self.assertEqual([call.args[0] for call in post.call_args_list], [ 'https://aiplatform.googleapis.com/v1/projects/project/locations/global/publishers/google/models/gemini-3.6-flash:countTokens', 'https://aiplatform.us.rep.googleapis.com/v1/projects/project/locations/us/publishers/google/models/gemini-3.6-flash:countTokens', 'https://aiplatform.eu.rep.googleapis.com/v1/projects/project/locations/eu/publishers/google/models/gemini-3.6-flash:countTokens', ]) def test_anthropic_probe_uses_global_host_and_one_token_generation(self): response = mock.Mock(status_code=200) with mock.patch.object(gcpKeycheck.requests, 'post', return_value=response) as post: result = gcpKeycheck.probe_vertex_anthropic( 'token', 'project', None, 6, locations=['global'], models=['claude-opus-5'], max_attempts=1, ) self.assertTrue(result['enabled']) self.assertEqual( post.call_args.args[0], 'https://aiplatform.googleapis.com/v1/projects/project/locations/global/publishers/anthropic/models/claude-opus-5:rawPredict', ) self.assertEqual(post.call_args.kwargs['json']['max_tokens'], 1) class GCPLegacyVertexImportTests(unittest.TestCase): def test_collect_deduplicates_valid_credentials_across_vertex_files(self): with tempfile.TemporaryDirectory() as temp_dir: keycheck_dir = Path(temp_dir) / 'keychecks' gcp_dir = keycheck_dir / 'gcp' ensure_private_directory(str(keycheck_dir)) ensure_private_directory(str(gcp_dir)) credential = json.dumps({ 'client_email': 'fixture@example.invalid', 'private_key_id': 'fixture-key-id', 'private_key': '-----BEGIN PRIVATE KEY-----\nfixture\n-----END PRIVATE KEY-----', }, sort_keys=True, separators=(',', ':')) for filename in keycheck_runner.LEGACY_GCP_VERTEX_STATUS_FILES: path = gcp_dir / filename path.write_text(f'{credential}\tVERTEX\tlegacy\n', encoding='utf-8') harden_private_file(str(path)) entries, skipped = keycheck_runner.collect_legacy_gcp_vertex_credentials({ 'keycheck_dir': str(keycheck_dir), }) self.assertEqual(skipped, 0) self.assertEqual(len(entries), 1) self.assertEqual(next(iter(entries.values())), set(keycheck_runner.LEGACY_GCP_VERTEX_STATUS_FILES)) if __name__ == '__main__': unittest.main()