Initial server source import
This commit is contained in:
sashatrask
@@ -0,0 +1,143 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user