Files
truf-server/tests/test_high_only_round7.py
2026-09-30 20:30:56 +03:00

656 lines
35 KiB
Python

import hashlib
import json
import os
from pathlib import Path
import shutil
import sqlite3
import subprocess
import sys
import tempfile
import threading
import time
from types import SimpleNamespace
import unittest
import uuid
import venv
from unittest import mock
ROOT = Path(__file__).resolve().parents[1]
APP_DIR = ROOT / 'app'
sys.path.insert(0, str(APP_DIR))
import child_bootstrap
import console_runner
from keycheckers import keycheck_common
import keycheck_runner
import lifecycle_authority
from process_identity import current_process_identity
from result_spool import prepare_scan_event
import runtime_security
import scanner
from scanner_db import ScannerDB, migrate_runtime_safety_schema
import supervisor
from supervisor_instance import build_instance_metadata, write_instance_metadata
class ShutdownAuthorityTests(unittest.TestCase):
def test_live_child_enters_failed_hold_and_never_releases_authority(self):
source = mock.Mock(source='stuck')
source.stop.return_value = False
source.is_running.return_value = True
context = {
'lifecycle_phase': lifecycle_authority.PHASE_ACTIVE,
'activation_state': lifecycle_authority.PHASE_ACTIVE,
'start_gate_open': True,
'shutdown_event': threading.Event(),
'supervisor_config': {},
}
self.assertFalse(supervisor.coordinated_shutdown([source], context))
self.assertEqual(context['lifecycle_phase'], lifecycle_authority.PHASE_FAILED_HOLD)
self.assertFalse(context['authority_release_safe'])
self.assertFalse(supervisor.lifecycle_start_allowed(context))
def test_failed_hold_keeps_authenticated_status_and_retry_but_denies_start(self):
source = mock.Mock(source='stuck', runtime_blocked=False, last_action_error='')
source.row.return_value = ['stuck', 'failed', '-', 'loop', '-', '-', '-', '0/0', '-', 'log', 'live', 'stopped']
context = {
'lifecycle_phase': lifecycle_authority.PHASE_FAILED_HOLD,
'activation_state': lifecycle_authority.PHASE_FAILED_HOLD,
'start_gate_open': False,
'shutdown_event': threading.Event(),
'shutdown_retry_event': threading.Event(),
'with_postgres': False,
}
server = supervisor.SupervisorControlServer(
('127.0.0.1', 0), [source], context, threading.RLock(), 'hold', 'h' * 48,
)
threading.Thread(target=server.serve_forever, daemon=True).start()
metadata = {
'instance_id': 'hold', 'token': 'h' * 48,
'control': {'host': '127.0.0.1', 'port': server.server_address[1]},
}
try:
snapshot = supervisor.send_control_request(metadata, 'snapshot')
self.assertEqual(snapshot['activation_state'], lifecycle_authority.PHASE_FAILED_HOLD)
with self.assertRaisesRegex(RuntimeError, 'FAILED_HOLD'):
supervisor.send_control_command(metadata, 'start stuck')
self.assertIn('shutdown requested', supervisor.send_control_request(metadata, 'shutdown'))
self.assertTrue(context['shutdown_retry_event'].is_set())
source.start.assert_not_called()
finally:
server.shutdown()
server.server_close()
def test_bounded_retry_reaches_safe_only_after_exact_child_exit(self):
source = mock.Mock(source='stuck')
source.stop.side_effect = [False, True]
source.is_running.side_effect = [True, False]
retry = threading.Event()
retry.set()
context = {
'lifecycle_phase': lifecycle_authority.PHASE_FAILED_HOLD,
'activation_state': lifecycle_authority.PHASE_FAILED_HOLD,
'start_gate_open': False,
'authority_release_safe': False,
'shutdown_retry_event': retry,
'supervisor_config': {'postgres_stop_failed_retry_sec': 1},
}
self.assertTrue(supervisor.retain_unsafe_authority([source], context, max_attempts=2))
self.assertTrue(context['authority_release_safe'])
def test_replacement_endpoint_authority_is_denied_while_root_holds_lock(self):
with tempfile.TemporaryDirectory() as temp_dir:
runtime_security.ensure_private_directory(temp_dir, reject_reparse=True)
postgres = os.path.join(temp_dir, 'postgres')
runtime_security.ensure_private_directory(postgres, reject_reparse=True)
config = {'global': {'runtime_dir': temp_dir}}
dsn = 'postgresql://role:secret@127.0.0.1:5544/shared'
root = runtime_security.ClusterAuthorityLock(config, endpoint_dsn=dsn).acquire()
try:
with self.assertRaises(BlockingIOError):
runtime_security.ClusterAuthorityLock(
{'global': {'runtime_dir': temp_dir}},
endpoint_dsn='postgresql://replacement:other@localhost:5544/shared',
).acquire()
finally:
root.release()
class OutboxSafetyTests(unittest.TestCase):
def setUp(self):
self.environment = mock.patch.dict(os.environ, {
'SCANNER_DB_URL': '', 'DATABASE_URL': '', 'TRUF_MANAGED_POSTGRES_DSN': '',
'SCAN_OUTBOX_RETRY_BASE_SEC': '1', 'SCAN_OUTBOX_RETRY_MAX_SEC': '4',
})
self.environment.start()
self.temp = tempfile.TemporaryDirectory()
self.db = ScannerDB(db_path=os.path.join(self.temp.name, 'scanner.db'))
self.run_id = self.db.start_run('test', ['test'])
self.cycle_id = self.db.start_source_cycle(
self.run_id, 'source', 'github', 'search', 'q', 1, 1, None, {}, {},
)
def tearDown(self):
self.db.close()
self.temp.cleanup()
self.environment.stop()
def ingest(self, target='https://github.com/example/outbox', blob_size=0):
self.db.enqueue_targets('source', 'github', 'q', [target])
claim = dict(self.db.claim_targets('source', 'github', 1, 'owner', 60, return_rows=True)[0])
event_id = str(uuid.uuid4())
result = {
'scan_event_id': event_id, 'target': target, 'scan_type': 'github',
'timestamp': '2026-07-19T00:00:00+00:00', 'findings': [], 'errors': [],
'large': 'x' * blob_size,
}
event = prepare_scan_event({
'version': 1, 'scan_event_id': event_id, 'run_id': self.run_id,
'cycle_id': self.cycle_id, 'source': 'source', 'query': 'q',
'target': target, 'result': result, 'scan_options': {},
'queue_id': claim['id'], 'claim_lease_token': claim['lease_token'],
'claim_lease_owner': claim['lease_owner'], 'queue_status': 'done',
})
return self.db.ingest_scan_event(event), result
def test_large_payload_is_reference_only_and_reconstructed_on_claim(self):
outcome, result = self.ingest(blob_size=2 * 1024 * 1024)
stored = self.db.conn.execute(
'SELECT payload_json FROM scan_publication_outbox WHERE id = ?', (outcome['outbox_id'],),
).fetchone()
self.assertEqual(stored['payload_json'], '')
rows = self.db.claim_scan_publications('publisher', 1)
self.assertEqual(json.loads(rows[0]['payload_json'])['large'], result['large'])
def test_repeated_failure_never_dead_letters_and_delivery_deletes(self):
outcome, _ = self.ingest()
for attempt in range(25):
self.db.conn.execute(
"UPDATE scan_publication_outbox SET available_after = '2000-01-01T00:00:00+00:00' WHERE id = ?",
(outcome['outbox_id'],),
)
self.db.conn.commit()
row = self.db.claim_scan_publications(f'publisher-{attempt}', 1)[0]
self.assertTrue(self.db.finish_scan_publication(row['id'], f'publisher-{attempt}', False, 'outage'))
state = self.db.conn.execute(
'SELECT status, attempts, payload_json FROM scan_publication_outbox WHERE id = ?',
(outcome['outbox_id'],),
).fetchone()
self.assertEqual(state['status'], 'pending')
self.assertEqual(state['attempts'], 25)
self.assertEqual(state['payload_json'], '')
self.db.conn.execute(
"UPDATE scan_publication_outbox SET available_after = '2000-01-01T00:00:00+00:00' WHERE id = ?",
(outcome['outbox_id'],),
)
self.db.conn.commit()
row = self.db.claim_scan_publications('final-publisher', 1)[0]
self.assertTrue(self.db.finish_scan_publication(row['id'], 'final-publisher', True))
self.assertIsNone(self.db.conn.execute(
'SELECT id FROM scan_publication_outbox WHERE id = ?', (outcome['outbox_id'],),
).fetchone())
def test_backlog_gate_closes_and_recovers_after_delivery(self):
outcome, _ = self.ingest()
blocked = self.db.scan_publication_backlog_health(1, 10 * 1024 * 1024, 86400, additional_items=1)
self.assertTrue(blocked['healthy'])
self.assertFalse(blocked['accepting'])
args = SimpleNamespace(
scan_outbox_max_pending_items=1,
scan_outbox_max_pending_bytes=10 * 1024 * 1024,
scan_outbox_max_pending_age_sec=86400,
)
with self.assertRaisesRegex(RuntimeError, 'backlog gate'):
console_runner.require_scan_publication_capacity(self.db, args)
row = self.db.claim_scan_publications('publisher', 1)[0]
self.assertTrue(self.db.finish_scan_publication(row['id'], 'publisher', True))
self.assertTrue(console_runner.require_scan_publication_capacity(self.db, args)['accepting'])
def test_migration_retries_and_compacts_legacy_dead_rows(self):
outcome, _ = self.ingest()
self.db.conn.execute(
"UPDATE scan_publication_outbox SET status = 'dead', payload_json = ? WHERE id = ?",
('x' * 1024 * 1024, outcome['outbox_id']),
)
self.db.conn.commit()
migrate_runtime_safety_schema(self.db)
row = self.db.conn.execute(
'SELECT status, payload_json FROM scan_publication_outbox WHERE id = ?',
(outcome['outbox_id'],),
).fetchone()
self.assertEqual((row['status'], row['payload_json']), ('pending', ''))
class BootstrapAndManifestTests(unittest.TestCase):
def test_all_supervised_child_commands_require_isolated_no_site_startup(self):
for kind in ('scanner', 'keycheck', 'dashboard', 'keycheck-provider'):
with self.subTest(kind=kind):
command = supervisor.child_bootstrap_command(
kind, [], provider_entrypoint='keycheckers/github/githubKeycheck.py' if kind == 'keycheck-provider' else None,
)
self.assertEqual(command[1:4], ['-I', '-S', '-B'])
def test_shadow_module_cannot_run_or_read_token_before_manifest_rejection(self):
with tempfile.TemporaryDirectory() as temp_dir:
app_dir = os.path.join(temp_dir, 'app')
runtime_dir = os.path.join(temp_dir, 'runtime')
shutil.copytree(APP_DIR, app_dir, ignore=shutil.ignore_patterns('__pycache__', '*.pyc'))
os.makedirs(runtime_dir)
shutil.copy2(ROOT / 'runtime' / 'check-openrouter-keys.ps1', os.path.join(runtime_dir, 'check-openrouter-keys.ps1'))
launcher_paths = []
for name in ('start_runtime.ps1', 'stop_runtime.ps1'):
launcher_path = os.path.join(temp_dir, name)
shutil.copy2(ROOT / name, launcher_path)
launcher_paths.append(launcher_path)
executable = os.path.join(temp_dir, 'trufflehog.exe')
Path(executable).write_bytes(b'fixture')
detector = os.path.join(app_dir, 'trufflehog-custom-detectors.yaml')
for path in (temp_dir, app_dir, runtime_dir):
runtime_security.harden_private_directory(path)
for path in (executable, detector, os.path.join(runtime_dir, 'check-openrouter-keys.ps1'), *launcher_paths):
runtime_security.harden_private_file(path)
manifest = lifecycle_authority.build_code_manifest(
app_dir=app_dir, trufflehog_path=executable, policy_paths=[detector],
)
config_path = os.path.join(app_dir, 'config.yaml')
instance_path = os.path.join(temp_dir, 'instance.json')
identity = current_process_identity()
dsn = 'postgresql://truf:fixture@127.0.0.1:5432/truf'
metadata = build_instance_metadata(
'nonce', os.path.join(app_dir, 'supervisor.py'), config_path,
'127.0.0.1', 9, False, identity=identity, activation_state='ACTIVE',
code_manifest=manifest,
expected_code_manifest_sha256=lifecycle_authority.code_manifest_sha256(manifest),
canonical_dsn_sha256=lifecycle_authority.dsn_sha256(dsn),
lifecycle_mode='foreground', instance_file=instance_path,
)
write_instance_metadata(instance_path, metadata)
marker = os.path.join(temp_dir, 'stolen.txt')
Path(os.path.join(app_dir, 'requests.py')).write_text(
'import os\nopen(os.environ["THEFT_MARKER"], "w").write(os.environ.get("TRUF_SUPERVISOR_TOKEN", ""))\n',
encoding='ascii',
)
environment = {
**os.environ,
**lifecycle_authority.supervised_child_environment(
{**metadata, 'instance_file': instance_path}, dsn, 'scanner',
),
'SCANNER_DB_URL': dsn,
'DATABASE_URL': dsn,
'THEFT_MARKER': marker,
}
completed = subprocess.run(
[sys.executable, '-I', '-S', '-B', os.path.join(app_dir, 'child_bootstrap.py'), 'scanner', '--', '--config', config_path],
cwd=app_dir, env=environment, stdin=subprocess.DEVNULL,
stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, timeout=20, check=False,
)
self.assertNotEqual(completed.returncode, 0)
self.assertIn('file set drifted', completed.stdout)
self.assertFalse(os.path.exists(marker))
def test_site_hooks_never_run_and_authenticated_child_imports_dependencies(self):
with tempfile.TemporaryDirectory() as temp_dir:
venv_dir = os.path.join(temp_dir, 'venv')
venv.EnvBuilder(with_pip=False, system_site_packages=True).create(venv_dir)
if os.name == 'nt':
venv_python = os.path.join(venv_dir, 'Scripts', 'python.exe')
site_packages = os.path.join(venv_dir, 'Lib', 'site-packages')
else:
venv_python = os.path.join(venv_dir, 'bin', 'python')
version = f'python{sys.version_info.major}.{sys.version_info.minor}'
site_packages = os.path.join(venv_dir, 'lib', version, 'site-packages')
site_marker = os.path.join(temp_dir, 'site-token.txt')
pth_marker = os.path.join(temp_dir, 'pth-token.txt')
Path(os.path.join(site_packages, 'sitecustomize.py')).write_text(
'import os\nopen(os.environ["SITE_MARKER"], "w", encoding="utf-8").write(os.environ.get("TRUF_SUPERVISOR_TOKEN", ""))\n',
encoding='ascii',
)
Path(os.path.join(site_packages, 'malicious.pth')).write_text(
'import os; open(os.environ["PTH_MARKER"], "w", encoding="utf-8").write(os.environ.get("TRUF_MANAGED_POSTGRES_DSN", ""))\n',
encoding='ascii',
)
rejected_environment = {
key: value for key, value in os.environ.items()
if key not in lifecycle_authority.PRIVATE_CHILD_ENV_KEYS
}
rejected_environment.update({
'TRUF_SUPERVISOR_TOKEN': 'pre-auth-token',
'TRUF_MANAGED_POSTGRES_DSN': 'postgresql://pre-auth-dsn',
'SITE_MARKER': site_marker,
'PTH_MARKER': pth_marker,
})
rejected_command = supervisor.child_bootstrap_command('scanner', ['--config', 'missing.yaml'])
rejected_command[0] = venv_python
rejected = subprocess.run(
rejected_command, cwd=temp_dir, env=rejected_environment, stdin=subprocess.DEVNULL,
stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, timeout=20, check=False,
)
self.assertNotEqual(rejected.returncode, 0)
self.assertIn('direct mutation is retired', rejected.stdout)
self.assertFalse(os.path.exists(site_marker))
self.assertFalse(os.path.exists(pth_marker))
app_dir = os.path.join(temp_dir, 'app')
runtime_dir = os.path.join(temp_dir, 'runtime')
shutil.copytree(APP_DIR, app_dir, ignore=shutil.ignore_patterns('__pycache__', '*.pyc'))
os.makedirs(runtime_dir)
shutil.copy2(ROOT / 'runtime' / 'check-openrouter-keys.ps1', os.path.join(runtime_dir, 'check-openrouter-keys.ps1'))
launcher_paths = []
for name in ('start_runtime.ps1', 'stop_runtime.ps1'):
launcher_path = os.path.join(temp_dir, name)
shutil.copy2(ROOT / name, launcher_path)
launcher_paths.append(launcher_path)
success_marker = os.path.join(temp_dir, 'child-success.txt')
Path(os.path.join(app_dir, 'console_runner.py')).write_text(
'import os\nimport boto3\nimport botocore\nimport pandas\nimport plotly\nimport psycopg\n'
'import requests\nimport streamlit\nimport urllib3\nimport yaml\n'
'with open(os.environ["CHILD_SUCCESS"], "w", encoding="utf-8") as handle:\n'
' handle.write(":".join(module.__name__ for module in '
'(boto3, botocore, pandas, plotly, psycopg, requests, streamlit, urllib3, yaml)))\n',
encoding='ascii',
)
config_path = os.path.join(app_dir, 'config.yaml')
Path(config_path).write_text('{}\n', encoding='ascii')
executable = os.path.join(temp_dir, 'trufflehog.exe' if os.name == 'nt' else 'trufflehog')
Path(executable).write_bytes(b'fixture')
for path in (temp_dir, app_dir, runtime_dir):
runtime_security.harden_private_directory(path)
for path in (config_path, executable, os.path.join(runtime_dir, 'check-openrouter-keys.ps1'), *launcher_paths):
runtime_security.harden_private_file(path)
manifest = lifecycle_authority.build_code_manifest(app_dir=app_dir, trufflehog_path=executable)
instance_id = 'bootstrap-test-instance'
token = 'b' * 48
dsn = 'postgresql://truf:fixture@127.0.0.1:5432/truf'
context = {
'lifecycle_phase': lifecycle_authority.PHASE_ACTIVE,
'activation_state': lifecycle_authority.PHASE_ACTIVE,
'with_postgres': True,
'authority': {},
'canonical_dsn_sha256': lifecycle_authority.dsn_sha256(dsn),
}
server = supervisor.SupervisorControlServer(
('127.0.0.1', 0), [], context, threading.RLock(), instance_id, token,
)
threading.Thread(target=server.serve_forever, daemon=True).start()
try:
instance_path = os.path.join(temp_dir, 'instance.json')
metadata = build_instance_metadata(
'nonce', os.path.join(app_dir, 'supervisor.py'), config_path,
'127.0.0.1', server.server_address[1], True,
identity=current_process_identity(), instance_id=instance_id, token=token,
activation_state=lifecycle_authority.PHASE_ACTIVE,
code_manifest=manifest,
expected_code_manifest_sha256=lifecycle_authority.code_manifest_sha256(manifest),
canonical_dsn_sha256=lifecycle_authority.dsn_sha256(dsn),
lifecycle_mode='foreground', instance_file=instance_path,
)
context['authority'] = {
key: metadata[key]
for key in ('config_sha256', 'supervisor_sha256', 'code_manifest_sha256')
}
write_instance_metadata(instance_path, metadata)
environment = {
**os.environ,
**lifecycle_authority.supervised_child_environment(metadata, dsn, 'scanner'),
'SCANNER_DB_URL': dsn,
'DATABASE_URL': dsn,
'SITE_MARKER': site_marker,
'PTH_MARKER': pth_marker,
'CHILD_SUCCESS': success_marker,
}
command = supervisor.child_bootstrap_command('scanner', [])
command[0] = venv_python
command[4] = os.path.join(app_dir, 'child_bootstrap.py')
completed = subprocess.run(
command, cwd=app_dir, env=environment, stdin=subprocess.DEVNULL,
stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, timeout=30, check=False,
)
finally:
server.shutdown()
server.server_close()
self.assertEqual(completed.returncode, 0, completed.stdout)
self.assertEqual(
Path(success_marker).read_text(encoding='utf-8'),
'boto3:botocore:pandas:plotly:psycopg:requests:streamlit:urllib3:yaml',
)
self.assertFalse(os.path.exists(site_marker))
self.assertFalse(os.path.exists(pth_marker))
def test_policy_asset_mutation_is_manifest_drift(self):
with tempfile.TemporaryDirectory() as temp_dir:
executable = os.path.join(temp_dir, 'trufflehog.exe')
policy = os.path.join(temp_dir, 'detectors.yaml')
Path(executable).write_bytes(b'fixture')
Path(policy).write_text('detectors: []\n', encoding='ascii')
manifest = lifecycle_authority.build_code_manifest(
trufflehog_path=executable, policy_paths=[policy],
)
Path(policy).write_text('detectors: [changed]\n', encoding='ascii')
with self.assertRaisesRegex(lifecycle_authority.LifecycleAuthorityError, 'asset'):
lifecycle_authority.verify_code_manifest(manifest)
class EndpointIdentityTests(unittest.TestCase):
def test_explicit_dsn_ignores_conflicting_environment_role_and_temp(self):
first = 'postgresql://role_one:a@localhost:6543/shared'
second = 'postgresql://role_two:b@127.0.0.1:6543/shared'
config = {'global': {'runtime_dir': r'D:\one'}}
with mock.patch.dict(os.environ, {
'TRUF_MANAGED_POSTGRES_DSN': 'postgresql://wrong:x@127.0.0.1:9999/wrong',
'TEMP': r'Z:\temp-one', 'TMP': r'Y:\temp-two',
}):
one = runtime_security.canonical_cluster_endpoint_identity(config, database_url=first)
path_one = runtime_security.cluster_endpoint_authority_lock_path(config, endpoint_dsn=first)
with mock.patch.dict(os.environ, {'TEMP': r'Q:\other', 'TMP': r'R:\other'}):
two = runtime_security.canonical_cluster_endpoint_identity(config, database_url=second)
path_two = runtime_security.cluster_endpoint_authority_lock_path(config, endpoint_dsn=second)
self.assertEqual(one, two)
self.assertEqual(path_one, path_two)
self.assertNotIn('role_', one)
with self.assertRaisesRegex(runtime_security.PrivateFileError, 'explicitly'):
runtime_security.canonical_cluster_endpoint_identity(config)
class RotationSafetyTests(unittest.TestCase):
@classmethod
def setUpClass(cls):
scanner.initialize_scanner_runtime(preflight_complete=True, register_cleanup=False)
def test_ledger_recovery_queries_use_supporting_indexes_at_scale(self):
with tempfile.TemporaryDirectory() as temp_dir:
scanner.ensure_private_directory(temp_dir, reject_reparse=True)
path = os.path.join(temp_dir, 'scan_results.jsonl')
connection = scanner._open_projection_ledger(path)
try:
now = time.time()
connection.executemany(
'''INSERT INTO publication_identity(
identity_key, identity_value, payload_sha256, state, file_name,
byte_offset, byte_length, created_at, updated_at
) VALUES ('scan_event_id', ?, ?, 'appended', 'scan_results.jsonl', 0, 2, ?, ?)''',
((f'event-{index}', '0' * 64, now + index, now + index) for index in range(5000)),
)
connection.commit()
prepared_plan = ' '.join(str(value) for row in connection.execute(
"EXPLAIN QUERY PLAN SELECT identity_key, identity_value FROM publication_identity WHERE state = 'prepared' ORDER BY created_at LIMIT 2"
) for value in row)
rotation_plan = ' '.join(str(value) for row in connection.execute(
"EXPLAIN QUERY PLAN SELECT identity_value FROM publication_identity WHERE file_name = 'scan_results.jsonl' AND state = 'appended'"
) for value in row)
finally:
connection.close()
self.assertIn('idx_publication_identity_state_created', prepared_plan)
self.assertIn('idx_publication_identity_file_state', rotation_plan)
def _read_scanner_rows(self, path, state_dir):
with mock.patch.dict(os.environ, {
'KEYCHECK_OUTPUT_DIR': state_dir,
'KEYCHECK_STATE_DIR': state_dir,
'KEYCHECK_SERVICE': 'fixture',
'KEYCHECK_INPUT_TAIL_BYTES': '0',
}):
return [item['data'] for item in keycheck_common.iter_jsonl_input(path)]
def test_every_scanner_rotation_crash_state_keeps_all_rows_visible(self):
for stage in ('orphan', 'manifest_before_truncate', 'truncated_before_final', 'new_append_after_truncate'):
with self.subTest(stage=stage), tempfile.TemporaryDirectory() as temp_dir:
scanner.ensure_private_directory(temp_dir, reject_reparse=True)
state_dir = os.path.join(temp_dir, 'state')
scanner.ensure_private_directory(state_dir, reject_reparse=True)
path = os.path.join(temp_dir, 'found_secrets.jsonl')
segment = os.path.join(temp_dir, 'found_secrets.000001.jsonl')
first = b'{"finding_uid":"first"}\n'
second = b'{"finding_uid":"second"}\n'
Path(segment).write_bytes(first)
Path(path).write_bytes(first)
scanner.harden_private_file(segment)
scanner.harden_private_file(path)
if stage != 'orphan':
manifest = {
'current_path': os.path.abspath(path),
'segments': [{'name': os.path.basename(segment), 'path': segment, 'sequence': 1, 'bytes': len(first)}],
'next_sequence': 2,
'current_skip_bytes': len(first),
'current_skip_signature': scanner._jsonl_file_generation(path),
}
scanner.write_jsonl_manifest(path, manifest)
if stage in ('truncated_before_final', 'new_append_after_truncate'):
Path(path).write_bytes(b'' if stage == 'truncated_before_final' else second)
scanner.harden_private_file(path)
rows = self._read_scanner_rows(path, state_dir)
identities = [row['finding_uid'] for row in rows]
self.assertIn('first', identities)
if stage == 'new_append_after_truncate':
self.assertIn('second', identities)
self.assertLessEqual(identities.count('first'), 2)
self.assertIn(segment, scanner.projection_segment_paths(path))
def test_keycheck_stale_skip_never_hides_new_current_rows(self):
with tempfile.TemporaryDirectory() as temp_dir:
runtime_security.ensure_private_directory(temp_dir, reject_reparse=True)
path = os.path.join(temp_dir, 'geminiResults.jsonl')
segment = os.path.join(temp_dir, 'geminiResults.000001.jsonl')
Path(segment).write_bytes(b'{"event_id":"old"}\n')
Path(path).write_bytes(b'{"event_id":"old"}\n')
runtime_security.harden_private_file(segment)
runtime_security.harden_private_file(path)
signature = keycheck_common._keycheck_file_generation(path)
keycheck_common.write_jsonl_manifest(path, {
'current_path': os.path.abspath(path),
'segments': [{'name': os.path.basename(segment), 'path': segment, 'sequence': 1}],
'current_skip_bytes': os.path.getsize(segment),
'current_skip_signature': signature,
})
Path(path).write_bytes(b'{"event_id":"new"}\n')
self.assertEqual(keycheck_runner.keycheck_result_current_skip(path), 0)
self.assertEqual(keycheck_runner.keycheck_result_paths(path), [segment, os.path.abspath(path)])
class PostmanAndCandidateTests(unittest.TestCase):
@classmethod
def setUpClass(cls):
scanner.initialize_scanner_runtime(preflight_complete=True, register_cleanup=False)
def test_duplicate_key_secret_survives_cache_byte_for_byte(self):
with tempfile.TemporaryDirectory() as temp_dir:
cache = os.path.join(temp_dir, 'cache')
scanner.ensure_private_directory(temp_dir, reject_reparse=True)
scanner.ensure_private_directory(cache, reject_reparse=True)
raw = b'{"token":"decoy","token":"duplicate-key-secret"}'
with mock.patch.object(scanner.scan_config, 'runtime_dir', temp_dir):
path, digest, size = scanner.write_postman_cache(
raw, cache_dir=cache, max_artifact_size_mb=1,
cache_max_items=10, cache_max_bytes=1024 * 1024, cache_min_free_bytes=0,
)
self.assertEqual(Path(path).read_bytes(), raw)
self.assertEqual((digest, size), (hashlib.sha256(raw).hexdigest(), len(raw)))
self.assertIn(b'duplicate-key-secret', Path(path).read_bytes())
def test_hard_json_limit_rejects_before_parse(self):
oversized = b'[' + (b'0,' * (scanner.POSTMAN_JSON_HARD_MAX_INPUT_BYTES // 2)) + b'0]'
with mock.patch.object(scanner.json, 'loads', side_effect=AssertionError('parser called')):
with self.assertRaisesRegex(ValueError, 'pre-parse'):
scanner.write_postman_cache(oversized, max_artifact_size_mb=64)
def test_fifty_thousand_values_are_capped_and_written_once(self):
contexts = [
{'path': f'$.values[{index}]', 'key': 'api_key', 'value': 'AIza' + f'{index:035d}', 'endpoint': '', 'host': ''}
for index in range(50000)
]
calls = []
def append(path, lines):
calls.append((path, list(lines)))
return len(lines)
with tempfile.TemporaryDirectory() as temp_dir, \
mock.patch.object(scanner, 'load_postman_context', return_value=contexts), \
mock.patch.object(scanner, 'keycheck_output_path', side_effect=lambda service, name: os.path.join(temp_dir, service, name)), \
mock.patch.object(scanner, 'append_unique_lines_locked', side_effect=append), \
mock.patch.object(scanner.scan_config, 'keycheck_candidate_artifact_max_items', 250), \
mock.patch.object(scanner.scan_config, 'keycheck_candidate_artifact_max_bytes', 1024 * 1024):
counts = scanner.write_structured_keycheck_candidates('artifact.json', {})
self.assertEqual(len(calls), 1)
self.assertEqual(len(calls[0][1]), 250)
self.assertEqual(counts, {'gemini': 250})
def test_candidate_file_refuses_partial_batch_at_aggregate_cap(self):
with tempfile.TemporaryDirectory() as temp_dir:
scanner.ensure_private_directory(temp_dir, reject_reparse=True)
path = os.path.join(temp_dir, 'gem.txt')
with mock.patch.object(scanner.scan_config, 'keycheck_candidate_file_max_items', 10), \
mock.patch.object(scanner.scan_config, 'keycheck_candidate_file_max_bytes', 1024 * 1024), \
mock.patch.object(scanner.os, 'fsync', wraps=os.fsync) as fsync:
with self.assertRaises(scanner.CandidateQueueCapacityError):
scanner.append_unique_lines_locked(path, [f'key-{index}' for index in range(50000)])
self.assertFalse(os.path.exists(path))
self.assertEqual(fsync.call_count, 0)
def test_keychecker_list_loading_is_bounded(self):
with tempfile.TemporaryDirectory() as temp_dir:
path = os.path.join(temp_dir, 'keys.txt')
Path(path).write_text('x' * 1024, encoding='ascii')
with mock.patch.dict(os.environ, {'KEYCHECK_INPUT_LIST_MAX_BYTES': '128'}):
with self.assertRaisesRegex(RuntimeError, 'byte bound'):
keycheck_common.load_keys_from_file(path)
class ScanSlotSafetyTests(unittest.TestCase):
def test_live_stale_owner_and_child_are_never_evicted_by_age(self):
identity = current_process_identity()
conn = sqlite3.connect(':memory:')
conn.executescript(scanner.SCAN_SLOT_SCHEMA)
old = time.time() - 100000
conn.execute(
'''INSERT INTO scan_slots(
slot_id, owner_pid, owner_thread, owner_source, owner_creation_time,
owner_executable, acquired_at, updated_at
) VALUES ('live-owner', ?, 1, 'test', ?, ?, ?, ?)''',
(identity.pid, identity.creation_time, identity.executable, old, old),
)
conn.execute(
'''INSERT INTO scan_slots(
slot_id, owner_pid, owner_thread, owner_source, owner_creation_time,
owner_executable, child_pid, child_creation_time, child_executable,
acquired_at, updated_at
) VALUES ('live-child', 99999991, 1, 'test', 'gone', ?, ?, ?, ?, ?, ?)''',
(identity.executable, identity.pid, identity.creation_time, identity.executable, old, old),
)
conn.execute(
'''INSERT INTO scan_slots(
slot_id, owner_pid, owner_thread, owner_source, owner_creation_time,
owner_executable, acquired_at, updated_at
) VALUES ('gone', 99999992, 1, 'test', 'gone', ?, ?, ?)''',
(identity.executable, old, old),
)
scanner.cleanup_stale_scan_slots(conn, time.time(), 60)
remaining = {row[0] for row in conn.execute('SELECT slot_id FROM scan_slots')}
self.assertEqual(remaining, {'live-owner', 'live-child'})
conn.close()
if __name__ == '__main__':
unittest.main()