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