sashatrask
821 lines
38 KiB
Python
821 lines
38 KiB
Python
import builtins
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
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 import keycheck_common
|
|
from keycheckers.azure import azureKeycheck
|
|
from keycheckers.gemini import geminiKeycheck
|
|
from keycheckers.huggingface import huggingfaceKeycheck
|
|
from keycheckers.openai import Keycheck as openaiKeycheck
|
|
from keycheckers.replicate import replicateKeycheck
|
|
from keycheckers.xai import xaiKeycheck
|
|
from keycheck_candidates import (
|
|
extract_candidates,
|
|
extract_structured_candidates,
|
|
stored_provider_key_hash,
|
|
)
|
|
from keycheck_runner import SERVICES
|
|
import migrate_runtime_safety
|
|
import supervisor
|
|
import console_runner
|
|
|
|
|
|
class PipelineCutoverInvariantTests(unittest.TestCase):
|
|
def test_supervisor_import_graph_has_no_scanner_or_legacy_spool_import(self):
|
|
source = (APP_DIR / 'supervisor.py').read_text(encoding='utf-8')
|
|
self.assertNotIn('from scanner import', source)
|
|
self.assertNotIn('from result_spool import', source)
|
|
self.assertNotIn('from scanner_db import', source)
|
|
self.assertNotIn('maybe_run_temp_cleanup', source)
|
|
|
|
def test_janitor_receives_no_database_authority(self):
|
|
supervisor_source = (APP_DIR / 'supervisor.py').read_text(encoding='utf-8')
|
|
bootstrap_source = (APP_DIR / 'child_bootstrap.py').read_text(encoding='utf-8')
|
|
self.assertIn(
|
|
"'janitor': supervised_child_environment(child_metadata, '', 'janitor')",
|
|
supervisor_source,
|
|
)
|
|
self.assertIn('janitor child must not receive database mutation capability', bootstrap_source)
|
|
worker = SimpleNamespace(
|
|
source='janitor', dependency_gate=None,
|
|
child_environment={
|
|
'TRUF_MANAGED_POSTGRES_DSN': '', 'TRUF_SUPERVISOR_CHILD_KIND': 'janitor',
|
|
},
|
|
)
|
|
with mock.patch.dict(os.environ, {
|
|
'TRUF_MANAGED_POSTGRES_DSN': 'postgresql://forbidden',
|
|
'SCANNER_DB_URL': 'postgresql://forbidden',
|
|
'DATABASE_URL': 'postgresql://forbidden',
|
|
'PGPASSWORD': 'forbidden',
|
|
}, clear=False):
|
|
child_env = supervisor.ManagedPipelineWorker.build_env(worker)
|
|
self.assertNotIn('TRUF_MANAGED_POSTGRES_DSN', child_env)
|
|
self.assertNotIn('SCANNER_DB_URL', child_env)
|
|
self.assertNotIn('DATABASE_URL', child_env)
|
|
self.assertNotIn('PGPASSWORD', child_env)
|
|
|
|
def test_normal_source_cutover_does_not_call_legacy_publication(self):
|
|
source = (APP_DIR / 'console_runner.py').read_text(encoding='utf-8')
|
|
start = source.index('def run_cycle_v2(')
|
|
end = source.index('\ndef _run_cycle_legacy_compat(', start)
|
|
cutover = source[start:end]
|
|
for forbidden in (
|
|
'drain_result_spool(', 'wait_for_result_spool_ready(',
|
|
'drain_scan_publication_outbox(', 'save_scan_result(',
|
|
'write_foundry_keycheck_candidates_from_findings(',
|
|
):
|
|
self.assertNotIn(forbidden, cutover)
|
|
self.assertNotIn('staged_results', cutover)
|
|
|
|
def test_admission_resolution_is_bounded_and_precedes_scan_permit(self):
|
|
source = (APP_DIR / 'console_runner.py').read_text(encoding='utf-8')
|
|
helper_start = source.index('def reserve_v2_admission_with_recovery(')
|
|
cycle_start = source.index('\ndef run_cycle_v2(', helper_start)
|
|
helper = source[helper_start:cycle_start]
|
|
self.assertIn('resolution_attempts=8', helper)
|
|
self.assertIn('resolution_seconds=30', helper)
|
|
self.assertNotIn('acquire_scan_slot(', helper)
|
|
self.assertIn('release_permit()', helper)
|
|
self.assertEqual(helper.count('reserve_and_claim_target('), 1)
|
|
self.assertIn('recover_result_reservation_claim(', helper)
|
|
cycle_end = source.index('\ndef _run_cycle_legacy_compat(', cycle_start)
|
|
cycle = source[cycle_start:cycle_end]
|
|
self.assertLess(
|
|
cycle.index('lease = acquire_scan_slot('),
|
|
cycle.index('outcome = reserve_with_exact_recovery(permit_box)'),
|
|
)
|
|
|
|
def test_container_capacity_model_covers_slots_and_headroom(self):
|
|
config = migrate_runtime_safety.load_config(str(APP_DIR / 'config.linux.yaml'))
|
|
global_config = config['global']
|
|
physical_slots = (
|
|
int(global_config['max_active_scans'])
|
|
+ int(global_config['opportunistic_scan_slots'])
|
|
)
|
|
model = console_runner.validate_v2_capacity_model(
|
|
physical_slots,
|
|
global_config['result_bundle_max_event_bytes'],
|
|
global_config['projection_backlog_max_bytes'],
|
|
global_config['projection_backlog_headroom_bytes'],
|
|
)
|
|
self.assertEqual(model['physical_slots'], 1)
|
|
self.assertEqual(model['per_scan_projection_bytes'], 134217728)
|
|
self.assertGreaterEqual(
|
|
global_config['projection_backlog_max_bytes'],
|
|
model['required_projection_bytes'],
|
|
)
|
|
|
|
def test_outage_ambiguity_releases_permit_and_never_resubmits_claim(self):
|
|
calls = {'factory': 0, 'reserve': 0, 'recover': 0, 'release': 0}
|
|
|
|
class DB:
|
|
enabled = True
|
|
|
|
@staticmethod
|
|
def set_application_name(_name):
|
|
return None
|
|
|
|
def reserve_and_claim_target(self, *args, **kwargs):
|
|
calls['reserve'] += 1
|
|
raise OSError('commit response unavailable')
|
|
|
|
def recover_result_reservation_claim(self, *args, **kwargs):
|
|
calls['recover'] += 1
|
|
raise OSError('database outage')
|
|
|
|
@staticmethod
|
|
def close():
|
|
return None
|
|
|
|
def factory(**_kwargs):
|
|
calls['factory'] += 1
|
|
return DB()
|
|
|
|
with self.assertRaisesRegex(RuntimeError, 'bounded exact-token resolution'):
|
|
console_runner.reserve_v2_admission_with_recovery(
|
|
'postgresql://fixture', 'github', 'github',
|
|
{'pid': 1, 'creation_time': 'time', 'executable': 'python'},
|
|
'supervisor', 1024, 2048, 1, 1024,
|
|
lease_seconds=60, max_attempts=3, capacity_limits={},
|
|
run_id=1, cycle_id=1, reservation_token='token',
|
|
bundle_id='a' * 32, scan_event_id='b' * 32,
|
|
resolution_attempts=3, resolution_seconds=30,
|
|
db_factory=factory, stop_requested=lambda: False,
|
|
sleep=lambda _delay: None,
|
|
release_permit=lambda: calls.__setitem__('release', calls['release'] + 1),
|
|
)
|
|
self.assertEqual(calls['reserve'], 1)
|
|
self.assertEqual(calls['recover'], 3)
|
|
self.assertEqual(calls['release'], 1)
|
|
|
|
def test_pre_admission_connection_outage_closes_cycle_without_recovery(self):
|
|
calls = {'factory': 0, 'reserve': 0, 'recover': 0, 'release': 0}
|
|
|
|
class DB:
|
|
enabled = False
|
|
|
|
def reserve_and_claim_target(self, *args, **kwargs):
|
|
calls['reserve'] += 1
|
|
|
|
def recover_result_reservation_claim(self, *args, **kwargs):
|
|
calls['recover'] += 1
|
|
|
|
@staticmethod
|
|
def close():
|
|
return None
|
|
|
|
def factory(**_kwargs):
|
|
calls['factory'] += 1
|
|
return DB()
|
|
|
|
outcome = console_runner.reserve_v2_admission_with_recovery(
|
|
'postgresql://fixture', 'github', 'github',
|
|
{'pid': 1, 'creation_time': 'time', 'executable': 'python'},
|
|
'supervisor', 1024, 2048, 1, 1024,
|
|
lease_seconds=60, max_attempts=3, capacity_limits={},
|
|
run_id=1, cycle_id=1, reservation_token='token',
|
|
bundle_id='a' * 32, scan_event_id='b' * 32,
|
|
resolution_attempts=3, resolution_seconds=30,
|
|
db_factory=factory,
|
|
release_permit=lambda: calls.__setitem__('release', calls['release'] + 1),
|
|
)
|
|
|
|
self.assertIsNone(outcome.claim)
|
|
self.assertTrue(outcome.permit_released)
|
|
self.assertFalse(outcome.retry_without_claim)
|
|
self.assertEqual(calls, {
|
|
'factory': 1, 'reserve': 0, 'recover': 0, 'release': 1,
|
|
})
|
|
|
|
def test_transient_admission_outage_recovers_exact_reservation_without_resubmit(self):
|
|
calls = {'reserve': 0, 'recover': 0, 'release': 0, 'sleep': [], 'expected': None}
|
|
recovered_claim = {'reservation_id': 41}
|
|
|
|
class DB:
|
|
enabled = True
|
|
|
|
@staticmethod
|
|
def set_application_name(_name):
|
|
return None
|
|
|
|
def reserve_and_claim_target(self, *args, **kwargs):
|
|
calls['reserve'] += 1
|
|
raise OSError('commit response unavailable')
|
|
|
|
def recover_result_reservation_claim(self, _token, expected):
|
|
calls['recover'] += 1
|
|
calls['expected'] = expected
|
|
if calls['recover'] < 3:
|
|
raise OSError('transient database outage')
|
|
return recovered_claim
|
|
|
|
@staticmethod
|
|
def close():
|
|
return None
|
|
|
|
outcome = console_runner.reserve_v2_admission_with_recovery(
|
|
'postgresql://fixture', 'github', 'github',
|
|
{'pid': 1, 'creation_time': 'time', 'executable': 'python'},
|
|
'supervisor', 1024, 2048, 1, 1024,
|
|
lease_seconds=60, max_attempts=3, capacity_limits={},
|
|
run_id=1, cycle_id=1, reservation_token='token',
|
|
bundle_id='a' * 32, scan_event_id='b' * 32,
|
|
resolution_attempts=90, resolution_seconds=90, retry_delay=1,
|
|
db_factory=lambda **_kwargs: DB(), stop_requested=lambda: False,
|
|
sleep=calls['sleep'].append,
|
|
release_permit=lambda: calls.__setitem__('release', calls['release'] + 1),
|
|
)
|
|
self.assertIs(outcome.claim, recovered_claim)
|
|
self.assertTrue(outcome.permit_released)
|
|
self.assertEqual(calls['reserve'], 1)
|
|
self.assertEqual(calls['recover'], 3)
|
|
self.assertEqual(calls['release'], 1)
|
|
self.assertEqual(calls['sleep'], [1, 1])
|
|
self.assertEqual(calls['expected']['assignment_kind'], 'local')
|
|
self.assertEqual(calls['expected']['reserved_bundle_bytes'], 1024)
|
|
self.assertIsNone(calls['expected']['remote_user_id'])
|
|
self.assertIsNone(calls['expected']['remote_execution_snapshot_sha256'])
|
|
|
|
def test_projector_rotation_never_copies_active_history(self):
|
|
source = (APP_DIR / 'jsonl_projector.py').read_text(encoding='utf-8')
|
|
self.assertNotIn('copyfileobj', source)
|
|
self.assertIn('durable_publish(active, segment)', source)
|
|
|
|
def test_bundle_quarantine_transfers_credit_before_physical_move(self):
|
|
source = (APP_DIR / 'result_ingester.py').read_text(encoding='utf-8')
|
|
start = source.index(' def quarantine(')
|
|
end = source.index('\n def recover(', start)
|
|
method = source[start:end]
|
|
prepare = method.index('quarantine_id = self.db.quarantine_result_bundle(')
|
|
move = method.index('durable_publish(ready.path, quarantine_path)')
|
|
confirm = method.rindex('physical_confirmed=True')
|
|
self.assertLess(prepare, move)
|
|
self.assertLess(move, confirm)
|
|
|
|
def test_ci_and_trufflehog_production_paths_are_streamed(self):
|
|
source = (APP_DIR / 'scanner.py').read_text(encoding='utf-8')
|
|
self.assertNotIn('def download_bytes(', source)
|
|
self.assertNotIn("b''.join(chunks)", source)
|
|
self.assertNotIn('stdout_file.read(max_stdout', source)
|
|
self.assertIn('def download_to_file(', source)
|
|
self.assertIn('def run_command_streamed(', source)
|
|
|
|
def test_postgres_candidate_mode_never_opens_found_secrets_jsonl(self):
|
|
candidate = {
|
|
'id': 7,
|
|
'service': 'openai',
|
|
'attempts': 1,
|
|
'lease_token': 'lease',
|
|
'detector_name': 'OpenAI',
|
|
'secret_text': 'sk-fixture',
|
|
'secret_json': None,
|
|
'finding': {'DetectorName': 'OpenAI', 'Raw': 'sk-fixture'},
|
|
}
|
|
|
|
class DB:
|
|
def __init__(self):
|
|
self.returned = False
|
|
|
|
def claim_keycheck_candidate(self, service, owner, lease_seconds, **kwargs):
|
|
if self.returned:
|
|
return None
|
|
self.returned = True
|
|
return candidate
|
|
|
|
def keycheck_candidate_cached_status(self, candidate_id, lease_token):
|
|
return {'probe_required': True, 'reason': 'no_current_state'}
|
|
|
|
def defer_keycheck_candidate(self, *args):
|
|
raise AssertionError('completed candidate was deferred')
|
|
|
|
real_open = builtins.open
|
|
|
|
def guarded_open(path, *args, **kwargs):
|
|
if os.path.basename(os.fspath(path)).lower() == 'found_secrets.jsonl':
|
|
raise AssertionError('normal keycheck opened found_secrets.jsonl')
|
|
return real_open(path, *args, **kwargs)
|
|
|
|
database = DB()
|
|
with mock.patch.dict(os.environ, {
|
|
'KEYCHECK_INPUT_MODE': 'postgres', 'KEYCHECK_SERVICE': 'openai',
|
|
}, clear=False), mock.patch.object(
|
|
keycheck_common, '_postgres_candidate_db', return_value=database,
|
|
), mock.patch('builtins.open', side_effect=guarded_open):
|
|
iterator = keycheck_common.iter_findings(
|
|
os.path.join('unavailable', 'found_secrets.jsonl'), ['OpenAI'],
|
|
)
|
|
item = next(iterator)
|
|
self.assertEqual(item['candidate_id'], 7)
|
|
candidate['_completed'] = True
|
|
with self.assertRaises(StopIteration):
|
|
next(iterator)
|
|
|
|
def test_unconsumed_postgres_candidate_is_quarantined_after_bound(self):
|
|
candidate = {
|
|
'id': 8, 'service': 'azure', 'attempts': 3, 'lease_token': 'lease',
|
|
'detector_name': 'Azure', 'candidate_kind': 'azure_foundry',
|
|
'secret_text': 'fixture', 'secret_json': None, 'endpoint': '',
|
|
'finding': {'DetectorName': 'Azure', 'Raw': 'fixture'},
|
|
}
|
|
|
|
class DB:
|
|
def __init__(self):
|
|
self.returned = False
|
|
self.quarantined = []
|
|
|
|
def claim_keycheck_candidate(self, service, owner, lease_seconds, **kwargs):
|
|
if self.returned:
|
|
return None
|
|
self.returned = True
|
|
return candidate
|
|
|
|
def keycheck_candidate_cached_status(self, candidate_id, lease_token):
|
|
return {'probe_required': True, 'reason': 'no_current_state'}
|
|
|
|
def quarantine_keycheck_candidate(self, *args):
|
|
self.quarantined.append(args)
|
|
return True
|
|
|
|
def defer_keycheck_candidate(self, *args):
|
|
raise AssertionError('bounded unconsumed candidate was deferred again')
|
|
|
|
database = DB()
|
|
with mock.patch.dict(os.environ, {
|
|
'KEYCHECK_INPUT_MODE': 'postgres', 'KEYCHECK_SERVICE': 'azure',
|
|
'KEYCHECK_CANDIDATE_MAX_UNCONSUMED_ATTEMPTS': '3',
|
|
}, clear=False), mock.patch.object(
|
|
keycheck_common, '_postgres_candidate_db', return_value=database,
|
|
):
|
|
iterator = keycheck_common.iter_findings('unused.jsonl', ['Azure'])
|
|
item = next(iterator)
|
|
self.assertEqual(item['candidate_kind'], 'azure_foundry')
|
|
with self.assertRaises(StopIteration):
|
|
next(iterator)
|
|
self.assertEqual(database.quarantined[0][2], 'provider_candidate_unconsumed')
|
|
|
|
def test_postgres_provider_paths_never_touch_plain_or_retry_files(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
legacy = os.path.join(temp_dir, 'legacy.txt')
|
|
Path(legacy).write_text('nonempty-legacy-key\n', encoding='ascii')
|
|
with mock.patch.dict(os.environ, {
|
|
'KEYCHECK_INPUT_MODE': 'postgres', 'KEYCHECK_SERVICE': 'gemini',
|
|
}, clear=False), mock.patch(
|
|
'builtins.open', side_effect=AssertionError('compatibility file opened'),
|
|
), mock.patch.object(
|
|
geminiKeycheck, 'iter_findings', return_value=iter(()),
|
|
), mock.patch.object(
|
|
azureKeycheck, 'iter_findings', return_value=iter(()),
|
|
):
|
|
self.assertEqual(list(keycheck_common.read_plain_keys([legacy], mock.Mock())), [])
|
|
self.assertEqual(list(geminiKeycheck.iter_candidate_keys('forbidden.jsonl', [legacy])), [])
|
|
self.assertEqual(list(azureKeycheck.extract_candidates('forbidden.jsonl')), [])
|
|
|
|
def test_postgres_probe_requires_active_fenced_candidate(self):
|
|
with mock.patch.dict(os.environ, {'KEYCHECK_INPUT_MODE': 'postgres'}, clear=False), \
|
|
mock.patch.object(keycheck_common, '_ACTIVE_DB_CANDIDATE', None):
|
|
with self.assertRaisesRegex(RuntimeError, 'active fenced candidate'):
|
|
keycheck_common.should_skip_key(
|
|
'fixture', {}, set(), SimpleNamespace(), service='openai',
|
|
)
|
|
|
|
def test_real_provider_probe_runs_only_under_active_database_fence(self):
|
|
key = 'sk-database-fenced-fixture'
|
|
leaf_common = sys.modules[openaiKeycheck.iter_findings.__module__]
|
|
|
|
class DB:
|
|
def __init__(self):
|
|
self.returned = False
|
|
self.completed = 0
|
|
|
|
def claim_keycheck_candidate(self, service, owner, lease_seconds, **kwargs):
|
|
if self.returned:
|
|
return None
|
|
self.returned = True
|
|
return {
|
|
'id': 17, 'service': service, 'attempts': 1,
|
|
'lease_token': 'lease', 'detector_name': 'OpenAI',
|
|
'secret_text': key, 'secret_json': None,
|
|
'metadata_json': '{}',
|
|
'finding': {'DetectorName': 'OpenAI', 'Raw': key},
|
|
}
|
|
|
|
def keycheck_candidate_cached_status(self, candidate_id, lease_token):
|
|
return {'probe_required': True, 'reason': 'no_current_state'}
|
|
|
|
def complete_keycheck_candidate(self, candidate_id, lease_token, *args, **kwargs):
|
|
self.completed += 1
|
|
self.assert_active(candidate_id, lease_token)
|
|
return {'completed': True, 'keycheck_result_id': 1, 'event_id': args[0]}
|
|
|
|
@staticmethod
|
|
def assert_active(candidate_id, lease_token):
|
|
active = leaf_common._ACTIVE_DB_CANDIDATE
|
|
if not active or active['id'] != candidate_id or lease_token != active['lease_token']:
|
|
raise AssertionError('probe/result had no active fenced candidate')
|
|
|
|
def defer_keycheck_candidate(self, *args):
|
|
raise AssertionError('completed candidate was deferred')
|
|
|
|
database = DB()
|
|
|
|
def probe(probe_key, _proxy):
|
|
database.assert_active(17, 'lease')
|
|
self.assertEqual(probe_key, key)
|
|
return 'dead', 'fixture'
|
|
|
|
args = SimpleNamespace(
|
|
input='forbidden.jsonl', plain=['forbidden.txt'], proxy_file='proxy.txt',
|
|
max_keys=0, retry_network=False, retry_limited=False,
|
|
retry_unknown=False, retry_restricted=False, retry_no_balance=False,
|
|
recheck_all=False,
|
|
)
|
|
with mock.patch.dict(os.environ, {
|
|
'KEYCHECK_INPUT_MODE': 'postgres', 'KEYCHECK_SERVICE': 'openai',
|
|
}, clear=False), mock.patch.object(
|
|
openaiKeycheck, 'require_provider_authority', return_value={},
|
|
), mock.patch.object(
|
|
openaiKeycheck, 'parse_args', return_value=args,
|
|
), mock.patch.object(
|
|
openaiKeycheck, 'load_proxies', return_value=None,
|
|
), mock.patch.object(
|
|
openaiKeycheck, 'check_authentication', side_effect=probe,
|
|
), mock.patch.object(
|
|
leaf_common, '_postgres_candidate_db', return_value=database,
|
|
), mock.patch(
|
|
'builtins.open', side_effect=AssertionError('provider opened compatibility input'),
|
|
):
|
|
leaf_common._ACTIVE_DB_CANDIDATE = None
|
|
openaiKeycheck.main()
|
|
self.assertEqual(database.completed, 1)
|
|
|
|
def test_duplicate_postgres_credentials_complete_each_occurrence_without_cached_state(self):
|
|
cases = (
|
|
(
|
|
'openai', openaiKeycheck, 'OpenAI', 'sk-' + ('a' * 40),
|
|
SimpleNamespace(
|
|
input='forbidden.jsonl', plain=['forbidden.txt'], proxy_file='proxy.txt',
|
|
max_keys=0, retry_network=False, retry_limited=False,
|
|
retry_unknown=False, retry_restricted=False, retry_no_balance=False,
|
|
recheck_all=False,
|
|
),
|
|
'check_authentication', ('dead', 'fixture invalid key'),
|
|
),
|
|
(
|
|
'gemini', geminiKeycheck, 'GoogleAI', 'AIza' + ('A' * 35),
|
|
SimpleNamespace(
|
|
input='forbidden.jsonl', plain=['forbidden.txt'], proxy_file='proxy.txt',
|
|
timeout=1, max_keys=0, probe_generation=False,
|
|
retry_limited=False, retry_unknown=False, retry_network=False,
|
|
retry_valid=False, recheck_all=False, debug=False,
|
|
),
|
|
'check_key', {
|
|
'status': 'INVALID', 'error': {'message': 'fixture invalid key'},
|
|
'models': [], 'model_infos': [], 'notable_models': [],
|
|
'generation_models': [], 'model_class': 'none',
|
|
},
|
|
),
|
|
)
|
|
for service, provider, detector, key, args, probe_name, probe_result in cases:
|
|
with self.subTest(service=service):
|
|
leaf_common = sys.modules[provider.iter_findings.__module__]
|
|
test_case = self
|
|
|
|
class DB:
|
|
def __init__(self):
|
|
self.candidates = [
|
|
{
|
|
'id': candidate_id, 'service': service, 'attempts': 1,
|
|
'lease_token': f'lease-{candidate_id}', 'credential_id': 99,
|
|
'detector_name': detector, 'secret_text': key,
|
|
'secret_json': None, 'metadata_json': '{}',
|
|
'finding': {'DetectorName': detector, 'Raw': key},
|
|
}
|
|
for candidate_id in (301, 302)
|
|
]
|
|
self.completed = []
|
|
|
|
def claim_keycheck_candidate(self, claimed_service, owner, lease_seconds, **kwargs):
|
|
test_case.assertEqual(claimed_service, service)
|
|
return self.candidates.pop(0) if self.candidates else None
|
|
|
|
def keycheck_candidate_cached_status(self, candidate_id, lease_token):
|
|
return {'probe_required': True, 'reason': 'no_current_state'}
|
|
|
|
def complete_keycheck_candidate(self, candidate_id, lease_token, *values, **kwargs):
|
|
active = leaf_common._ACTIVE_DB_CANDIDATE
|
|
test_case.assertIsNotNone(active)
|
|
test_case.assertEqual(
|
|
(candidate_id, lease_token), (active['id'], active['lease_token']),
|
|
)
|
|
test_case.assertNotIn(candidate_id, self.completed)
|
|
self.completed.append(candidate_id)
|
|
return {
|
|
'completed': True, 'keycheck_result_id': candidate_id,
|
|
'event_id': values[0],
|
|
}
|
|
|
|
def defer_keycheck_candidate(self, *values):
|
|
raise AssertionError(f'candidate was deferred: {values}')
|
|
|
|
database = DB()
|
|
probe_side_effect = (
|
|
lambda probe_key, _proxy, _args: probe_result
|
|
if service == 'gemini'
|
|
else probe_result
|
|
)
|
|
patches = [
|
|
mock.patch.dict(os.environ, {
|
|
'KEYCHECK_INPUT_MODE': 'postgres', 'KEYCHECK_SERVICE': service,
|
|
}, clear=False),
|
|
mock.patch.object(provider, 'require_provider_authority', return_value={}),
|
|
mock.patch.object(provider, 'parse_args', return_value=args),
|
|
mock.patch.object(provider, 'load_proxies', return_value=None),
|
|
mock.patch.object(provider, 'record_cached_keycheck_occurrence', side_effect=AssertionError('cached-file path executed')),
|
|
mock.patch.object(leaf_common, '_postgres_candidate_db', return_value=database),
|
|
mock.patch('builtins.open', side_effect=AssertionError('provider opened compatibility state')),
|
|
]
|
|
if service == 'openai':
|
|
patches.append(mock.patch.object(provider, probe_name, return_value=probe_result))
|
|
else:
|
|
patches.append(mock.patch.object(provider, probe_name, side_effect=probe_side_effect))
|
|
patches.append(mock.patch.object(provider.time, 'sleep'))
|
|
with patches[0], patches[1], patches[2], patches[3], patches[4], patches[5], patches[6], patches[7] as probe:
|
|
if len(patches) == 9:
|
|
with patches[8]:
|
|
leaf_common._ACTIVE_DB_CANDIDATE = None
|
|
provider.main()
|
|
else:
|
|
leaf_common._ACTIVE_DB_CANDIDATE = None
|
|
provider.main()
|
|
self.assertEqual(database.completed, [301, 302])
|
|
self.assertEqual(probe.call_count, 2)
|
|
self.assertIsNone(leaf_common._ACTIVE_DB_CANDIDATE)
|
|
|
|
def test_all_enabled_providers_use_the_central_cached_skip_gate(self):
|
|
for service, relative_path in SERVICES.items():
|
|
with self.subTest(service=service):
|
|
source = (APP_DIR / relative_path).read_text(encoding='utf-8')
|
|
self.assertNotIn('\ndef should_skip_key(', source)
|
|
|
|
def test_final_cutover_refuses_nonempty_legacy_outbox(self):
|
|
class Cursor:
|
|
@staticmethod
|
|
def fetchone():
|
|
return {'count': 1}
|
|
|
|
class Connection:
|
|
is_postgres = False
|
|
|
|
@staticmethod
|
|
def table_exists(name):
|
|
return name == 'scan_publication_outbox'
|
|
|
|
@staticmethod
|
|
def execute(sql, params=None):
|
|
return Cursor()
|
|
|
|
db = SimpleNamespace(conn=Connection())
|
|
with self.assertRaisesRegex(RuntimeError, 'scan_publication_outbox'):
|
|
migrate_runtime_safety.require_legacy_cutover_clear(db, {'global': {}})
|
|
|
|
def test_final_cutover_refuses_bounded_legacy_spool_object(self):
|
|
class Cursor:
|
|
@staticmethod
|
|
def fetchone():
|
|
return {'count': 0}
|
|
|
|
class Connection:
|
|
is_postgres = False
|
|
|
|
@staticmethod
|
|
def table_exists(name):
|
|
return name == 'scan_publication_outbox'
|
|
|
|
@staticmethod
|
|
def execute(sql, params=None):
|
|
return Cursor()
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
Path(os.path.join(temp_dir, 'event.json')).write_text('{}', encoding='ascii')
|
|
db = SimpleNamespace(conn=Connection())
|
|
with self.assertRaisesRegex(RuntimeError, 'legacy result spool'):
|
|
migrate_runtime_safety.require_legacy_cutover_clear(
|
|
db, {'global': {'legacy_result_spool_dir': temp_dir}},
|
|
)
|
|
|
|
def test_structured_postman_candidates_route_without_file_outboxes(self):
|
|
contexts = [
|
|
{'value': 'AIza' + ('A' * 32), 'key': 'GEMINI_API_KEY', 'path': '$.gemini'},
|
|
{
|
|
'value': 'a' * 32,
|
|
'key': 'AZURE_OPENAI_KEY',
|
|
'endpoint': 'fixture.openai.azure.com',
|
|
'path': '$.azure',
|
|
},
|
|
]
|
|
candidates = list(extract_structured_candidates({
|
|
'contexts': contexts, 'origin': 'fixture:postman',
|
|
}))
|
|
self.assertEqual({candidate.service for candidate in candidates}, {'gemini', 'azure'})
|
|
azure = next(candidate for candidate in candidates if candidate.service == 'azure')
|
|
self.assertEqual(azure.metadata['raw_v2'], f'{"a" * 32}:fixture.openai.azure.com')
|
|
|
|
def test_aq_gemini_candidates_route_from_findings_and_structured_contexts(self):
|
|
key = 'AQ.Ab' + ('A' * 48)
|
|
direct = list(extract_candidates({
|
|
'DetectorName': 'GoogleAIStudio', 'Raw': key,
|
|
}))
|
|
structured = list(extract_structured_candidates({
|
|
'contexts': [{'value': key, 'key': 'GEMINI_API_KEY', 'path': '$.gemini'}],
|
|
'origin': 'fixture:postman',
|
|
}))
|
|
|
|
self.assertEqual(len(direct), 1)
|
|
self.assertEqual(len(structured), 1)
|
|
self.assertEqual(direct[0].service, 'gemini')
|
|
self.assertEqual(direct[0].secret_text, key)
|
|
self.assertEqual(direct[0].provider_key_hash, structured[0].provider_key_hash)
|
|
for malformed in ('AQ.Ab' + ('A' * 47), 'AQ.Ab' + ('A' * 49)):
|
|
self.assertEqual(list(extract_candidates({
|
|
'DetectorName': 'GoogleAIStudio', 'Raw': malformed,
|
|
})), [])
|
|
|
|
def test_provider_key_hash_matches_every_enabled_checker_input(self):
|
|
fixtures = [
|
|
({'DetectorName': 'OpenAI', 'Raw': 'sk-openai-fixture'}, 'sk-openai-fixture'),
|
|
({'DetectorName': 'Anthropic', 'Raw': 'sk-ant-fixture'}, 'sk-ant-fixture'),
|
|
({'DetectorName': 'DeepSeek', 'Raw': 'sk-deepseek-fixture'}, 'sk-deepseek-fixture'),
|
|
({'DetectorName': 'OpenRouter', 'Raw': 'sk-or-fixture'}, 'sk-or-fixture'),
|
|
({'DetectorName': 'Groq', 'Raw': 'gsk_fixture'}, 'gsk_fixture'),
|
|
({'DetectorName': 'Replicate', 'Raw': 'r8_fixture'}, 'r8_fixture'),
|
|
({'DetectorName': 'XAI', 'Raw': 'xai-fixture'}, 'xai-fixture'),
|
|
({'DetectorName': 'HuggingFace', 'Raw': 'hf_fixture'}, 'hf_fixture'),
|
|
({'DetectorName': 'GitHub', 'Raw': 'ghp_' + ('A' * 36)}, 'ghp_' + ('A' * 36)),
|
|
({'DetectorName': 'GitLab', 'Raw': 'glpat-' + ('A' * 24)}, 'glpat-' + ('A' * 24)),
|
|
({'DetectorName': 'QwenDashScope', 'Raw': 'sk-' + ('Q' * 24)}, 'sk-' + ('Q' * 24)),
|
|
({'DetectorName': 'KimiMoonshot', 'Raw': 'sk-' + ('K' * 40)}, 'sk-' + ('K' * 40)),
|
|
({
|
|
'DetectorName': 'ZaiGLM',
|
|
'Raw': ('a' * 32) + '.' + ('Z' * 16),
|
|
'ScannerContext': {'provider_hint': 'zai'},
|
|
}, ('a' * 32) + '.' + ('Z' * 16)),
|
|
({
|
|
'DetectorName': 'DeepSeek',
|
|
'Raw': 'sk-' + ('R' * 32),
|
|
'ScannerContext': {
|
|
'provider_hint': 'ambiguous_generic_sk',
|
|
'provider_candidates': ['deepseek', 'zai', 'qwen', 'kimi'],
|
|
},
|
|
}, 'sk-' + ('R' * 32)),
|
|
({'DetectorName': 'GoogleAIStudio', 'Raw': 'AIza' + ('G' * 32)}, 'AIza' + ('G' * 32)),
|
|
({'DetectorName': 'AWS', 'RawV2': 'AKIA' + ('A' * 16) + ':' + ('s' * 40)}, 'AKIA' + ('A' * 16) + ':' + ('s' * 40)),
|
|
]
|
|
gcp_value = {
|
|
'client_email': 'fixture@example.invalid',
|
|
'private_key': '-----BEGIN PRIVATE KEY-----\nfixture\n-----END PRIVATE KEY-----',
|
|
'private_key_id': 'ключ',
|
|
}
|
|
fixtures.append((
|
|
{'DetectorName': 'GCP', 'RawV2': json.dumps(gcp_value, ensure_ascii=False)},
|
|
json.dumps(gcp_value, ensure_ascii=False, sort_keys=True, separators=(',', ':')),
|
|
))
|
|
docker_token = 'dckr_pat_' + ('D' * 27)
|
|
fixtures.append((
|
|
{'DetectorName': 'Dockerhub', 'Raw': docker_token, 'RawV2': f'user:{docker_token}'},
|
|
f'user:{docker_token}',
|
|
))
|
|
azure_key = 'a' * 32
|
|
fixtures.append((
|
|
{
|
|
'DetectorName': 'AzureOpenAI', 'Raw': azure_key,
|
|
'RawV2': f'{azure_key}:fixture.openai.azure.com',
|
|
},
|
|
f'fixture.openai.azure.com:{azure_key}',
|
|
))
|
|
|
|
candidates = []
|
|
for finding, checker_key in fixtures:
|
|
candidate = next(iter(extract_candidates(finding)))
|
|
self.assertEqual(
|
|
candidate.provider_key_hash,
|
|
hashlib.sha256(checker_key.encode('utf-8')).hexdigest(),
|
|
candidate.service,
|
|
)
|
|
self.assertEqual(
|
|
stored_provider_key_hash(
|
|
candidate.service, candidate.candidate_kind,
|
|
candidate.secret_text, candidate.secret_json,
|
|
candidate.endpoint, candidate.principal,
|
|
),
|
|
candidate.provider_key_hash,
|
|
)
|
|
candidates.append(candidate)
|
|
self.assertEqual({candidate.service for candidate in candidates}, set(SERVICES))
|
|
|
|
def test_structured_and_detector_credentials_share_canonical_identity(self):
|
|
key = 'a' * 32
|
|
direct = next(iter(extract_candidates({
|
|
'DetectorName': 'AzureOpenAI', 'Raw': key,
|
|
'RawV2': f'{key}:fixture.openai.azure.com',
|
|
})))
|
|
structured = next(candidate for candidate in extract_structured_candidates({
|
|
'origin': 'fixture',
|
|
'contexts': [{
|
|
'value': key, 'key': 'AZURE_OPENAI_KEY',
|
|
'endpoint': 'fixture.openai.azure.com',
|
|
}],
|
|
}) if candidate.service == 'azure')
|
|
self.assertEqual(direct.provider_key_hash, structured.provider_key_hash)
|
|
self.assertEqual(direct.credential_hash, structured.credential_hash)
|
|
self.assertEqual(direct.secret_hash, structured.secret_hash)
|
|
|
|
def test_foundry_detector_candidate_extracts_only_key_from_context(self):
|
|
key = 'FoundryFixtureKey1234567890'
|
|
endpoint = 'fixture.models.ai.azure.com'
|
|
candidate = next(iter(extract_candidates({
|
|
'DetectorName': 'AzureFoundryEndpointBeforeKey',
|
|
'Raw': f'https://{endpoint}\napi_key = {key}',
|
|
})))
|
|
self.assertEqual(candidate.candidate_kind, 'azure_foundry')
|
|
self.assertEqual(candidate.secret_text, key)
|
|
self.assertEqual(candidate.endpoint, endpoint)
|
|
self.assertEqual(
|
|
candidate.provider_key_hash,
|
|
hashlib.sha256(f'{endpoint}:{key}'.encode('utf-8')).hexdigest(),
|
|
)
|
|
|
|
def test_malformed_azure_service_principal_does_not_become_foundry(self):
|
|
self.assertEqual(list(extract_candidates({
|
|
'DetectorName': 'Azure', 'Raw': 'not-a-service-principal-fixture-1234567890',
|
|
})), [])
|
|
|
|
def test_postgres_typed_azure_openai_candidate_bypasses_detector_mismatch(self):
|
|
key = 'a' * 32
|
|
item = {
|
|
'detector': 'Azure', 'raw': key, 'raw_v2': '',
|
|
'source': 'postgres:keycheck_candidates:1',
|
|
'finding': {'DetectorName': 'Azure', 'Raw': key},
|
|
'candidate_kind': 'azure_openai',
|
|
'credential_secret_text': key, 'credential_secret_json': '',
|
|
'credential_endpoint': '', 'credential_principal': '',
|
|
}
|
|
with mock.patch.dict(os.environ, {'KEYCHECK_INPUT_MODE': 'postgres'}, clear=False), \
|
|
mock.patch.object(azureKeycheck, 'iter_findings', return_value=iter([item])):
|
|
candidates = list(azureKeycheck.extract_candidates('unused.jsonl'))
|
|
self.assertEqual(len(candidates), 1)
|
|
self.assertEqual(candidates[0][0], key)
|
|
self.assertEqual(candidates[0][1], 'AzureOpenAI')
|
|
|
|
def test_postgres_typed_azure_foundry_uses_normalized_key_and_endpoint(self):
|
|
key = 'FoundryNormalizedFixtureKey1234567890'
|
|
endpoint = 'fixture.models.ai.azure.com'
|
|
item = {
|
|
'detector': 'AzureFoundryEndpointBeforeKey', 'raw': 'legacy-context', 'raw_v2': '',
|
|
'source': 'postgres:keycheck_candidates:2',
|
|
'finding': {'DetectorName': 'AzureFoundryEndpointBeforeKey', 'Raw': 'legacy-context'},
|
|
'candidate_kind': 'azure_foundry',
|
|
'credential_secret_text': key, 'credential_secret_json': '',
|
|
'credential_endpoint': endpoint, 'credential_principal': '',
|
|
}
|
|
with mock.patch.dict(os.environ, {'KEYCHECK_INPUT_MODE': 'postgres'}, clear=False), \
|
|
mock.patch.object(azureKeycheck, 'iter_findings', return_value=iter([item])):
|
|
candidates = list(azureKeycheck.extract_candidates('unused.jsonl'))
|
|
self.assertEqual(len(candidates), 1)
|
|
self.assertEqual(candidates[0][0], f'{endpoint}:{key}')
|
|
self.assertEqual(candidates[0][1], 'AzureFoundry')
|
|
|
|
def test_postgres_provider_keys_prefer_normalized_credential_material(self):
|
|
cases = (
|
|
(huggingfaceKeycheck, 'hf_' + ('a' * 39)),
|
|
(replicateKeycheck, 'r8_' + ('b' * 37)),
|
|
(xaiKeycheck, 'xai-' + ('c' * 40)),
|
|
)
|
|
for provider, key in cases:
|
|
item = {
|
|
'detector': provider.DETECTOR_NAMES[0] if hasattr(provider, 'DETECTOR_NAMES') else provider.DETECTOR,
|
|
'raw': 'legacy-context', 'raw_v2': '', 'source': 'postgres:keycheck_candidates:3',
|
|
'finding': {}, 'candidate_kind': 'provider_key',
|
|
'credential_secret_text': key, 'credential_secret_json': '',
|
|
'credential_endpoint': '', 'credential_principal': '',
|
|
}
|
|
with self.subTest(provider=provider.SERVICE), \
|
|
mock.patch.object(provider, 'iter_findings', return_value=iter([item])):
|
|
candidates = list(provider.extract_candidates('unused.jsonl', []))
|
|
self.assertEqual(candidates[0][0], key)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|