sashatrask
144 lines
6.8 KiB
Python
144 lines
6.8 KiB
Python
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()
|