Initial server source import
This commit is contained in:
sashatrask
@@ -0,0 +1,696 @@
|
||||
import ast
|
||||
import io
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
APP_DIR = ROOT / 'app'
|
||||
sys.path.insert(0, str(APP_DIR))
|
||||
|
||||
import console_runner
|
||||
import keycheck_runner
|
||||
import lifecycle_authority
|
||||
import scanner
|
||||
import supervisor
|
||||
from keycheckers import keycheck_common
|
||||
from lifecycle_authority import (
|
||||
LifecycleAuthorityError,
|
||||
PHASE_ACTIVE,
|
||||
build_code_manifest,
|
||||
dsn_sha256,
|
||||
require_active_supervisor_child,
|
||||
supervised_child_environment,
|
||||
)
|
||||
from runtime_security import ensure_private_directory, harden_private_file
|
||||
from supervisor_instance import build_instance_metadata, write_instance_metadata
|
||||
|
||||
|
||||
class DirectEntrypointTests(unittest.TestCase):
|
||||
def test_production_provider_scheduler_serializes_handshakes(self):
|
||||
config = keycheck_runner.load_config(str(APP_DIR / 'config.yaml'))
|
||||
self.assertEqual(config['keychecks']['scheduler_workers'], 1)
|
||||
active = 0
|
||||
peak = 0
|
||||
lock = threading.Lock()
|
||||
|
||||
def run_service(*_args, **_kwargs):
|
||||
nonlocal active, peak
|
||||
with lock:
|
||||
active += 1
|
||||
peak = max(peak, active)
|
||||
time.sleep(0.01)
|
||||
with lock:
|
||||
active -= 1
|
||||
return 0
|
||||
|
||||
with mock.patch.object(keycheck_runner, 'run_service', side_effect=run_service):
|
||||
code, failures = keycheck_runner.run_provider_services(
|
||||
list(keycheck_runner.SERVICES)[:4], SimpleNamespace(max_keys=1), {},
|
||||
config['keychecks'], work_probe=lambda _service: True,
|
||||
)
|
||||
self.assertEqual((code, failures), (0, []))
|
||||
self.assertEqual(peak, 1)
|
||||
|
||||
def test_provider_scheduler_is_fair_bounded_and_max_four(self):
|
||||
requested = list(keycheck_runner.SERVICES)[:8]
|
||||
active = 0
|
||||
peak = 0
|
||||
seen = []
|
||||
lock = threading.Lock()
|
||||
|
||||
def run_service(service, _script, args, _layout, _extra):
|
||||
nonlocal active, peak
|
||||
with lock:
|
||||
active += 1
|
||||
peak = max(peak, active)
|
||||
seen.append((service, args.max_keys))
|
||||
time.sleep(0.02)
|
||||
with lock:
|
||||
active -= 1
|
||||
return 0
|
||||
|
||||
with mock.patch.object(keycheck_runner, 'run_service', side_effect=run_service):
|
||||
code, failures = keycheck_runner.run_provider_services(
|
||||
requested, SimpleNamespace(max_keys=0), {}, {'scheduler_batch_keys': 7},
|
||||
)
|
||||
self.assertEqual(code, 0)
|
||||
self.assertEqual(failures, [])
|
||||
self.assertLessEqual(peak, 4)
|
||||
self.assertGreaterEqual(peak, 2)
|
||||
self.assertEqual({service for service, _ in seen}, set(requested))
|
||||
self.assertEqual({limit for _, limit in seen}, {7})
|
||||
|
||||
def test_provider_scheduler_drains_more_than_one_slice_fairly(self):
|
||||
requested = list(keycheck_runner.SERVICES)[:5]
|
||||
remaining = {service: (2505 if index == 0 else 3) for index, service in enumerate(requested)}
|
||||
submissions = []
|
||||
active = 0
|
||||
peak = 0
|
||||
lock = threading.Lock()
|
||||
|
||||
def run_service(service, _script, args, _layout, _extra):
|
||||
nonlocal active, peak
|
||||
with lock:
|
||||
active += 1
|
||||
peak = max(peak, active)
|
||||
submissions.append(service)
|
||||
time.sleep(0.01)
|
||||
with lock:
|
||||
remaining[service] = max(0, remaining[service] - args.max_keys)
|
||||
active -= 1
|
||||
return 0
|
||||
|
||||
with mock.patch.object(keycheck_runner, 'run_service', side_effect=run_service):
|
||||
code, failures = keycheck_runner.run_provider_services(
|
||||
requested, SimpleNamespace(max_keys=0), {},
|
||||
{'scheduler_batch_keys': 1000, 'scheduler_deadline_sec': 30},
|
||||
work_probe=lambda service: remaining[service] > 0,
|
||||
)
|
||||
self.assertEqual((code, failures), (0, []))
|
||||
self.assertEqual(remaining, {service: 0 for service in requested})
|
||||
self.assertEqual(submissions.count(requested[0]), 3)
|
||||
self.assertEqual(set(submissions[:5]), set(requested))
|
||||
self.assertLessEqual(peak, 4)
|
||||
|
||||
def test_provider_scheduler_does_not_hot_loop_capacity_blocked_service(self):
|
||||
calls = []
|
||||
|
||||
def blocked(service, *_args, **_kwargs):
|
||||
calls.append(service)
|
||||
return keycheck_runner.KEYCHECK_CAPACITY_BLOCKED_EXIT
|
||||
|
||||
with mock.patch.object(keycheck_runner, 'run_service', side_effect=blocked):
|
||||
code, failures = keycheck_runner.run_provider_services(
|
||||
['openai'], SimpleNamespace(max_keys=0), {},
|
||||
{'scheduler_batch_keys': 1000, 'scheduler_deadline_sec': 30},
|
||||
work_probe=lambda _service: True,
|
||||
)
|
||||
self.assertEqual((code, failures), (0, []))
|
||||
self.assertEqual(calls, ['openai'])
|
||||
|
||||
def test_console_runner_refuses_before_config_database_dependencies_or_scan(self):
|
||||
args = SimpleNamespace(show_state=False, config='missing-config.yaml')
|
||||
with mock.patch.object(console_runner, 'parse_args', return_value=args), \
|
||||
mock.patch.object(console_runner, 'run_config_mode') as run_config, \
|
||||
mock.patch.object(console_runner, 'ScannerDB') as scanner_db, \
|
||||
mock.patch.object(console_runner, 'check_dependencies') as dependencies, \
|
||||
mock.patch.object(console_runner, 'scan_targets_batch') as scan, \
|
||||
mock.patch.dict(os.environ, {'SCANNER_SUPERVISED': '1'}, clear=True):
|
||||
with self.assertRaisesRegex(SystemExit, 'direct mutation is retired'):
|
||||
console_runner.main()
|
||||
run_config.assert_not_called()
|
||||
scanner_db.assert_not_called()
|
||||
dependencies.assert_not_called()
|
||||
scan.assert_not_called()
|
||||
|
||||
def test_keycheck_runner_refuses_before_config_database_or_files(self):
|
||||
args = SimpleNamespace(print_plan=False, config='missing-config.yaml')
|
||||
with mock.patch.object(keycheck_runner, 'parse_args', return_value=args), \
|
||||
mock.patch.object(keycheck_runner, 'load_config') as load_config, \
|
||||
mock.patch.object(keycheck_runner, 'ScannerDB') as scanner_db, \
|
||||
mock.patch.object(keycheck_runner, 'run_service') as run_service, \
|
||||
mock.patch.dict(os.environ, {'SCANNER_SUPERVISED': '1'}, clear=True):
|
||||
with self.assertRaisesRegex(SystemExit, 'direct mutation is retired'):
|
||||
keycheck_runner.main()
|
||||
load_config.assert_not_called()
|
||||
scanner_db.assert_not_called()
|
||||
run_service.assert_not_called()
|
||||
|
||||
def test_print_plan_sync_reset_requires_authority_before_any_mutation(self):
|
||||
with mock.patch.object(sys, 'argv', [
|
||||
'keycheck_runner.py', '--print-plan', '--sync-keychecks-to-db', '--reset-keycheck-results',
|
||||
]):
|
||||
args = keycheck_runner.parse_args()
|
||||
with mock.patch.object(keycheck_runner, 'parse_args', return_value=args), \
|
||||
mock.patch.object(
|
||||
keycheck_runner, 'require_active_supervisor_child',
|
||||
side_effect=LifecycleAuthorityError('direct mutation is retired'),
|
||||
) as authority, \
|
||||
mock.patch.object(keycheck_runner, 'load_config') as load_config, \
|
||||
mock.patch.object(keycheck_runner, 'ScannerDB') as scanner_db, \
|
||||
mock.patch.object(keycheck_runner, 'sync_keycheck_results_to_db') as sync:
|
||||
with self.assertRaisesRegex(SystemExit, 'direct mutation is retired'):
|
||||
keycheck_runner.main()
|
||||
authority.assert_called_once()
|
||||
load_config.assert_not_called()
|
||||
scanner_db.assert_not_called()
|
||||
sync.assert_not_called()
|
||||
|
||||
def test_authenticated_mutating_plan_is_strictly_read_only(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir, mock.patch.object(sys, 'argv', [
|
||||
'keycheck_runner.py', '--print-plan', '--sync-keychecks-to-db', '--reset-keycheck-results',
|
||||
]):
|
||||
args = keycheck_runner.parse_args()
|
||||
layout = {
|
||||
'project_dir': str(APP_DIR),
|
||||
'results_dir': os.path.join(temp_dir, 'results'),
|
||||
'keycheck_dir': os.path.join(temp_dir, 'keychecks'),
|
||||
'proxy_file': os.path.join(temp_dir, 'proxy.txt'),
|
||||
}
|
||||
with mock.patch.object(keycheck_runner, 'parse_args', return_value=args), \
|
||||
mock.patch.object(keycheck_runner, 'require_active_supervisor_child') as authority, \
|
||||
mock.patch.object(keycheck_runner, 'load_config', return_value={'global': layout, 'keychecks': {}}), \
|
||||
mock.patch.object(keycheck_runner, 'preflight_lifecycle_paths') as preflight, \
|
||||
mock.patch.object(keycheck_runner, 'require_sensitive_runtime_paths') as sensitive_paths, \
|
||||
mock.patch.object(keycheck_runner, 'load_postgres_env') as postgres_env, \
|
||||
mock.patch.object(keycheck_runner, 'ScannerDB') as scanner_db, \
|
||||
mock.patch.object(keycheck_runner, 'sync_keycheck_results_to_db') as sync, \
|
||||
mock.patch.object(keycheck_runner, 'repair_keycheck_links') as repair, \
|
||||
mock.patch.object(keycheck_runner, 'ingest_keycheck_results_to_db') as ingest, \
|
||||
mock.patch.object(keycheck_runner, 'write_summary') as summary, \
|
||||
mock.patch.object(keycheck_runner, 'run_service') as run_service:
|
||||
keycheck_runner.main()
|
||||
authority.assert_called_once()
|
||||
preflight.assert_not_called()
|
||||
sensitive_paths.assert_not_called()
|
||||
postgres_env.assert_not_called()
|
||||
scanner_db.assert_not_called()
|
||||
sync.assert_not_called()
|
||||
repair.assert_not_called()
|
||||
ingest.assert_not_called()
|
||||
summary.assert_not_called()
|
||||
run_service.assert_not_called()
|
||||
self.assertEqual(list(Path(temp_dir).iterdir()), [])
|
||||
|
||||
def test_plain_print_plan_needs_no_authority_and_touches_no_runtime_paths(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir, mock.patch.object(sys, 'argv', [
|
||||
'keycheck_runner.py', '--print-plan', '--service', 'github', '--no-summary',
|
||||
]):
|
||||
args = keycheck_runner.parse_args()
|
||||
layout = {
|
||||
'project_dir': str(APP_DIR),
|
||||
'results_dir': os.path.join(temp_dir, 'results'),
|
||||
'keycheck_dir': os.path.join(temp_dir, 'keychecks'),
|
||||
'proxy_file': os.path.join(temp_dir, 'proxy.txt'),
|
||||
}
|
||||
with mock.patch.object(keycheck_runner, 'parse_args', return_value=args), \
|
||||
mock.patch.object(keycheck_runner, 'require_active_supervisor_child') as authority, \
|
||||
mock.patch.object(keycheck_runner, 'load_config', return_value={'global': layout, 'keychecks': {}}), \
|
||||
mock.patch.object(keycheck_runner, 'preflight_lifecycle_paths') as preflight, \
|
||||
mock.patch.object(keycheck_runner, 'require_sensitive_runtime_paths') as sensitive_paths, \
|
||||
mock.patch.object(keycheck_runner, 'load_postgres_env') as postgres_env, \
|
||||
mock.patch.object(keycheck_runner, 'ScannerDB') as scanner_db, \
|
||||
mock.patch.object(keycheck_runner, 'write_summary') as summary, \
|
||||
mock.patch.object(keycheck_runner, 'run_service') as run_service:
|
||||
keycheck_runner.main()
|
||||
authority.assert_not_called()
|
||||
preflight.assert_not_called()
|
||||
sensitive_paths.assert_not_called()
|
||||
postgres_env.assert_not_called()
|
||||
scanner_db.assert_not_called()
|
||||
summary.assert_not_called()
|
||||
run_service.assert_not_called()
|
||||
self.assertEqual(list(Path(temp_dir).iterdir()), [])
|
||||
|
||||
def test_empty_and_foreign_dsn_fail_before_handshake_or_application_work(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
ensure_private_directory(temp_dir, reject_reparse=True)
|
||||
config_path = os.path.join(temp_dir, 'config.yaml')
|
||||
instance_path = os.path.join(temp_dir, 'supervisor.instance.json')
|
||||
Path(config_path).write_text('{}\n', encoding='ascii')
|
||||
harden_private_file(config_path)
|
||||
|
||||
for dsn in ('', 'postgresql://truf:secret@attacker.invalid:5432/truf'):
|
||||
with self.subTest(dsn=dsn):
|
||||
if os.path.exists(instance_path):
|
||||
os.remove(instance_path)
|
||||
metadata = build_instance_metadata(
|
||||
'nonce',
|
||||
supervisor.__file__,
|
||||
config_path,
|
||||
'127.0.0.1',
|
||||
12345,
|
||||
False,
|
||||
activation_state=PHASE_ACTIVE,
|
||||
lifecycle_mode='foreground',
|
||||
canonical_dsn_sha256=dsn_sha256(dsn),
|
||||
instance_file=instance_path,
|
||||
)
|
||||
write_instance_metadata(instance_path, metadata)
|
||||
child = dict(metadata, instance_file=instance_path)
|
||||
environment = supervised_child_environment(child, dsn, 'scanner')
|
||||
environment['SCANNER_DB_URL'] = dsn
|
||||
environment['DATABASE_URL'] = dsn
|
||||
with mock.patch.dict(os.environ, environment, clear=True), \
|
||||
mock.patch.object(lifecycle_authority, '_verify_supervisor_process'), \
|
||||
mock.patch.object(lifecycle_authority, '_send_handshake') as handshake:
|
||||
with self.assertRaises(LifecycleAuthorityError):
|
||||
require_active_supervisor_child(config_path, 'scanner', require_dsn=True)
|
||||
handshake.assert_not_called()
|
||||
|
||||
def test_invalid_preflight_with_docker_token_has_no_import_cleanup_or_credential_files(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
ensure_private_directory(temp_dir, reject_reparse=True)
|
||||
work_dir = os.path.join(temp_dir, 'work')
|
||||
ensure_private_directory(work_dir, reject_reparse=True)
|
||||
stale = os.path.join(work_dir, 'trufflehog-stale')
|
||||
ensure_private_directory(stale, reject_reparse=True)
|
||||
marker = os.path.join(stale, '.scanner-owner.json')
|
||||
Path(marker).write_text('{"owner_pid": 99999999}\n', encoding='ascii')
|
||||
harden_private_file(marker)
|
||||
config_path = os.path.join(temp_dir, 'config.yaml')
|
||||
Path(config_path).write_text(
|
||||
'global:\n'
|
||||
f' root_dir: "{temp_dir.replace(os.sep, "/")}"\n'
|
||||
f' project_dir: "{temp_dir.replace(os.sep, "/")}"\n'
|
||||
f' runtime_dir: "{os.path.join(temp_dir, "missing-runtime").replace(os.sep, "/")}"\n'
|
||||
f' work_dir: "{work_dir.replace(os.sep, "/")}"\n'
|
||||
'sources: {}\n',
|
||||
encoding='ascii',
|
||||
)
|
||||
harden_private_file(config_path)
|
||||
before = sorted(str(path.relative_to(temp_dir)) for path in Path(temp_dir).rglob('*'))
|
||||
environment = os.environ.copy()
|
||||
environment.update({
|
||||
'DOCKER_TOKENS': 'example:token-value',
|
||||
'TRUFFLEHOG_WORK_DIR': work_dir,
|
||||
})
|
||||
completed = subprocess.run(
|
||||
[sys.executable, str(APP_DIR / 'supervisor.py'), '--config', config_path, '--background-status'],
|
||||
cwd=temp_dir,
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
env=environment,
|
||||
timeout=20,
|
||||
check=False,
|
||||
)
|
||||
after = sorted(str(path.relative_to(temp_dir)) for path in Path(temp_dir).rglob('*'))
|
||||
self.assertNotEqual(completed.returncode, 0)
|
||||
self.assertIn('preflight failed', completed.stdout)
|
||||
self.assertEqual(after, before)
|
||||
self.assertTrue(os.path.exists(stale))
|
||||
self.assertFalse(os.path.exists(os.path.join(work_dir, 'docker-config')))
|
||||
|
||||
def test_direct_leaf_provider_refuses_before_args_files_or_database(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
output_dir = os.path.join(temp_dir, 'outputs')
|
||||
environment = {
|
||||
key: value for key, value in os.environ.items()
|
||||
if not key.startswith('TRUF_SUPERVISOR_') and key not in ('SCANNER_SUPERVISED', 'TRUF_MANAGED_POSTGRES_DSN')
|
||||
}
|
||||
environment['KEYCHECK_OUTPUT_DIR'] = output_dir
|
||||
completed = subprocess.run(
|
||||
[sys.executable, str(APP_DIR / 'keycheckers' / 'github' / 'githubKeycheck.py'), '--input', 'missing.jsonl'],
|
||||
cwd=temp_dir,
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
env=environment,
|
||||
timeout=20,
|
||||
check=False,
|
||||
)
|
||||
self.assertNotEqual(completed.returncode, 0)
|
||||
self.assertIn('direct mutation is retired', completed.stdout)
|
||||
self.assertFalse(os.path.exists(output_dir))
|
||||
|
||||
def test_runner_delegates_distinct_provider_capability_and_managed_dsn(self):
|
||||
canonical = 'postgresql://truf:secret@127.0.0.1:5432/truf'
|
||||
metadata = {
|
||||
'instance_file': r'D:\fixture\instance.json',
|
||||
'instance_id': 'instance',
|
||||
'token': 't' * 48,
|
||||
'config_sha256': '1' * 64,
|
||||
'supervisor_sha256': '2' * 64,
|
||||
'code_manifest_sha256': '3' * 64,
|
||||
'canonical_dsn_sha256': lifecycle_authority.dsn_sha256(canonical),
|
||||
}
|
||||
|
||||
class CapturedProcess:
|
||||
def wait(self):
|
||||
return 0
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
results = os.path.join(temp_dir, 'results')
|
||||
keychecks = os.path.join(temp_dir, 'keychecks')
|
||||
ensure_private_directory(temp_dir, reject_reparse=True)
|
||||
ensure_private_directory(results, reject_reparse=True)
|
||||
ensure_private_directory(keychecks, reject_reparse=True)
|
||||
layout = {
|
||||
'project_dir': str(APP_DIR),
|
||||
'results_dir': results,
|
||||
'keycheck_dir': keychecks,
|
||||
'proxy_file': os.path.join(temp_dir, 'proxy.txt'),
|
||||
'database_path': '',
|
||||
'database_url': canonical,
|
||||
}
|
||||
args = SimpleNamespace(
|
||||
config='config.yaml', input=None, proxy_file=None, max_keys=0,
|
||||
retry_network=False, retry_limited=False, retry_unknown=False,
|
||||
retry_restricted=False, retry_no_balance=False, retry_valid=False,
|
||||
no_resource_probe=False, recheck_all=False,
|
||||
)
|
||||
with mock.patch.object(keycheck_runner, 'require_active_supervisor_child', return_value=metadata), \
|
||||
mock.patch.object(keycheck_runner, 'OwnedProcess', return_value=CapturedProcess()) as owned, \
|
||||
mock.patch.dict(os.environ, {
|
||||
'TRUF_MANAGED_POSTGRES_DSN': canonical,
|
||||
'SCANNER_DB_URL': canonical,
|
||||
'DATABASE_URL': canonical,
|
||||
}, clear=False):
|
||||
self.assertEqual(keycheck_runner.run_service('github', keycheck_runner.SERVICES['github'], args, layout), 0)
|
||||
child_env = owned.call_args.kwargs['env']
|
||||
child_command = owned.call_args.args[0]
|
||||
self.assertEqual(child_command[1:4], ['-I', '-S', '-B'])
|
||||
provider_index = child_command.index('keycheck-provider')
|
||||
self.assertEqual(child_command[provider_index + 2], '--')
|
||||
self.assertNotIn('--input', child_command)
|
||||
self.assertEqual(child_env[lifecycle_authority.CHILD_KIND_ENV], 'keycheck-provider')
|
||||
self.assertEqual(child_env[lifecycle_authority.CHILD_TOKEN_ENV], metadata['token'])
|
||||
self.assertEqual(child_env['SCANNER_DB_URL'], canonical)
|
||||
self.assertEqual(child_env['DATABASE_URL'], canonical)
|
||||
self.assertEqual(child_env['KEYCHECK_DB_URL'], canonical)
|
||||
|
||||
def test_runner_reports_the_exact_provider_nonzero_exit(self):
|
||||
canonical = 'postgresql://truf:secret@127.0.0.1:5432/truf'
|
||||
metadata = {
|
||||
'instance_file': r'D:\fixture\instance.json', 'instance_id': 'instance',
|
||||
'token': 't' * 48, 'config_sha256': '1' * 64,
|
||||
'supervisor_sha256': '2' * 64, 'code_manifest_sha256': '3' * 64,
|
||||
'canonical_dsn_sha256': lifecycle_authority.dsn_sha256(canonical),
|
||||
}
|
||||
|
||||
class FailedProcess:
|
||||
stdout = io.BytesIO(b'supervised child runtime failed: fixture database timeout\n')
|
||||
|
||||
def wait(self):
|
||||
return 7
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
results = os.path.join(temp_dir, 'results')
|
||||
keychecks = os.path.join(temp_dir, 'keychecks')
|
||||
for path in (temp_dir, results, keychecks):
|
||||
ensure_private_directory(path, reject_reparse=True)
|
||||
layout = {
|
||||
'project_dir': str(APP_DIR), 'results_dir': results, 'keycheck_dir': keychecks,
|
||||
'proxy_file': os.path.join(temp_dir, 'proxy.txt'), 'database_path': '',
|
||||
'database_url': canonical,
|
||||
}
|
||||
args = SimpleNamespace(
|
||||
config='config.yaml', input=None, proxy_file=None, max_keys=0,
|
||||
retry_network=False, retry_limited=False, retry_unknown=False,
|
||||
retry_restricted=False, retry_no_balance=False, retry_valid=False,
|
||||
no_resource_probe=False, recheck_all=False,
|
||||
)
|
||||
with mock.patch.object(keycheck_runner, 'require_active_supervisor_child', return_value=metadata), \
|
||||
mock.patch.object(keycheck_runner, 'OwnedProcess', return_value=FailedProcess()) as owned, \
|
||||
mock.patch('builtins.print') as output, \
|
||||
mock.patch.dict(os.environ, {'TRUF_MANAGED_POSTGRES_DSN': canonical}, clear=False):
|
||||
self.assertEqual(
|
||||
keycheck_runner.run_service('gcp', keycheck_runner.SERVICES['gcp'], args, layout),
|
||||
7,
|
||||
)
|
||||
diagnostic = Path(keychecks, 'gcp', '.provider-last-run.log')
|
||||
self.assertEqual(diagnostic.read_bytes(), FailedProcess.stdout.getvalue())
|
||||
self.assertIs(owned.call_args.kwargs['stdout'], subprocess.PIPE)
|
||||
self.assertIs(owned.call_args.kwargs['stderr'], subprocess.STDOUT)
|
||||
self.assertTrue(any(
|
||||
'provider_exit: service=gcp code=7 outcome=infrastructure_failure' in str(call.args[0])
|
||||
and '.provider-last-run.log' in str(call.args[0])
|
||||
for call in output.call_args_list
|
||||
))
|
||||
|
||||
def test_secondary_handshake_retries_one_timeout_then_succeeds(self):
|
||||
expected = {'instance_id': 'fixture'}
|
||||
with mock.patch.object(
|
||||
lifecycle_authority, '_send_handshake',
|
||||
side_effect=(TimeoutError('timed out'), expected),
|
||||
) as handshake:
|
||||
result = lifecycle_authority._send_handshake_with_timeout_retry({}, 1)
|
||||
self.assertEqual(result, expected)
|
||||
self.assertEqual(handshake.call_count, 2)
|
||||
|
||||
def test_secondary_handshake_two_timeouts_fail_after_two_attempts(self):
|
||||
with mock.patch.object(
|
||||
lifecycle_authority, '_send_handshake',
|
||||
side_effect=(TimeoutError('first'), TimeoutError('second')),
|
||||
) as handshake:
|
||||
with self.assertRaisesRegex(TimeoutError, 'second'):
|
||||
lifecycle_authority._send_handshake_with_timeout_retry({}, 1)
|
||||
self.assertEqual(handshake.call_count, 2)
|
||||
|
||||
def test_secondary_handshake_non_timeout_failure_is_not_retried(self):
|
||||
failure = LifecycleAuthorityError('authority rejected')
|
||||
with mock.patch.object(
|
||||
lifecycle_authority, '_send_handshake', side_effect=failure,
|
||||
) as handshake:
|
||||
with self.assertRaisesRegex(LifecycleAuthorityError, 'authority rejected'):
|
||||
lifecycle_authority._send_handshake_with_timeout_retry({}, 1)
|
||||
self.assertEqual(handshake.call_count, 1)
|
||||
|
||||
def test_provider_authority_opts_into_one_secondary_handshake_retry(self):
|
||||
with mock.patch.object(
|
||||
keycheck_common, 'require_active_supervisor_child', return_value={},
|
||||
) as authority, mock.patch.dict(os.environ, {
|
||||
'KEYCHECK_SERVICE': 'different-service',
|
||||
}, clear=False):
|
||||
with self.assertRaisesRegex(LifecycleAuthorityError, 'capability'):
|
||||
keycheck_common.require_provider_authority('aws')
|
||||
authority.assert_called_once_with(
|
||||
child_kind='keycheck-provider', require_dsn=True,
|
||||
handshake_timeout_retries=1,
|
||||
)
|
||||
|
||||
def test_provider_runtime_startup_failure_is_sanitized_and_stopped(self):
|
||||
canonical = 'postgresql://truf:secret@127.0.0.1:5432/truf'
|
||||
metadata = {
|
||||
'instance_file': r'D:\fixture\instance.json', 'instance_id': 'instance',
|
||||
'token': 't' * 48, 'config_sha256': '1' * 64,
|
||||
'supervisor_sha256': '2' * 64, 'code_manifest_sha256': '3' * 64,
|
||||
'canonical_dsn_sha256': lifecycle_authority.dsn_sha256(canonical),
|
||||
}
|
||||
|
||||
class PartialProcess:
|
||||
stdout = io.BytesIO()
|
||||
|
||||
def __init__(self):
|
||||
self.terminated = False
|
||||
self.killed = False
|
||||
|
||||
def wait(self, timeout=None):
|
||||
if timeout is None:
|
||||
raise OSError('sensitive command payload')
|
||||
return 0 if self.terminated or self.killed else None
|
||||
|
||||
def terminate(self):
|
||||
self.terminated = True
|
||||
|
||||
def kill(self):
|
||||
self.killed = True
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
results = os.path.join(temp_dir, 'results')
|
||||
keychecks = os.path.join(temp_dir, 'keychecks')
|
||||
for path in (temp_dir, results, keychecks):
|
||||
ensure_private_directory(path, reject_reparse=True)
|
||||
layout = {
|
||||
'project_dir': str(APP_DIR), 'results_dir': results,
|
||||
'keycheck_dir': keychecks, 'proxy_file': os.path.join(temp_dir, 'proxy.txt'),
|
||||
'database_path': '', 'database_url': canonical,
|
||||
}
|
||||
args = SimpleNamespace(
|
||||
config='config.yaml', input=None, proxy_file=None, max_keys=0,
|
||||
retry_network=False, retry_limited=False, retry_unknown=False,
|
||||
retry_restricted=False, retry_no_balance=False, retry_valid=False,
|
||||
no_resource_probe=False, recheck_all=False,
|
||||
)
|
||||
with mock.patch.object(keycheck_runner, 'require_active_supervisor_child', return_value=metadata), \
|
||||
mock.patch.object(
|
||||
keycheck_runner, 'OwnedProcess',
|
||||
side_effect=subprocess.TimeoutExpired(['sensitive-command'], 15),
|
||||
), \
|
||||
mock.patch('builtins.print') as output, \
|
||||
mock.patch.dict(os.environ, {'TRUF_MANAGED_POSTGRES_DSN': canonical}, clear=False):
|
||||
code = keycheck_runner.run_service(
|
||||
'aws', keycheck_runner.SERVICES['aws'], args, layout,
|
||||
)
|
||||
diagnostic = Path(keychecks, 'aws', '.provider-last-run.log')
|
||||
payload = diagnostic.read_text(encoding='ascii')
|
||||
process = PartialProcess()
|
||||
self.assertTrue(keycheck_runner._stop_failed_provider_process(process))
|
||||
self.assertEqual(code, 1)
|
||||
self.assertTrue(process.terminated)
|
||||
self.assertFalse(process.killed)
|
||||
self.assertEqual(payload, 'provider startup failed: TimeoutExpired\n')
|
||||
self.assertNotIn('sensitive', payload)
|
||||
self.assertFalse(any('sensitive' in str(call) for call in output.call_args_list))
|
||||
|
||||
def test_provider_failure_continues_services_and_reaches_finalization(self):
|
||||
args = SimpleNamespace(
|
||||
config='config.yaml', service='all', input=None, proxy_file=None,
|
||||
max_keys=0, retry_network=True, retry_limited=False,
|
||||
retry_unknown=False, retry_restricted=False, retry_no_balance=False,
|
||||
retry_valid=False, no_resource_probe=False, recheck_all=False,
|
||||
print_plan=False, summary_only=False, no_summary=False,
|
||||
no_db_ingest=False, ingest_keychecks_to_db=False, ingest_max_rows=None,
|
||||
no_auto_repair_links=False, sync_keychecks_to_db=False,
|
||||
reset_keycheck_results=False, repair_keycheck_links=False,
|
||||
import_legacy_gcp_vertex=False,
|
||||
repair_batch_size=25, repair_max_rows=0, repair_max_attempts=3,
|
||||
repair_since_days=14, summary_tsv=None, summary_json=None,
|
||||
alive_summary_tsv=None,
|
||||
)
|
||||
layout = {
|
||||
'results_dir': 'results', 'keycheck_dir': 'keychecks',
|
||||
'proxy_file': 'proxy', 'project_dir': str(APP_DIR),
|
||||
}
|
||||
requested = ['aws', 'azure', 'deepseek']
|
||||
with mock.patch.object(keycheck_runner, 'parse_args', return_value=args), \
|
||||
mock.patch.object(keycheck_runner, 'require_active_supervisor_child'), \
|
||||
mock.patch.object(keycheck_runner, 'load_config', return_value={
|
||||
'global': layout, 'keychecks': {},
|
||||
}), mock.patch.object(keycheck_runner, 'preflight_lifecycle_paths') as preflight, \
|
||||
mock.patch.object(keycheck_runner, 'require_sensitive_runtime_paths'), \
|
||||
mock.patch.object(keycheck_runner, 'load_postgres_env'), \
|
||||
mock.patch.object(keycheck_runner, 'apply_keycheck_config_defaults'), \
|
||||
mock.patch.object(keycheck_runner, 'service_list', return_value=requested), \
|
||||
mock.patch.object(keycheck_runner, 'run_service', side_effect=(1, 0, 0)) as run, \
|
||||
mock.patch.object(keycheck_runner, 'write_summary') as summary, \
|
||||
mock.patch.object(keycheck_runner, 'maybe_ingest_keycheck_results') as ingest, \
|
||||
mock.patch.object(keycheck_runner, 'maybe_auto_repair_links') as repair:
|
||||
with self.assertRaises(SystemExit) as raised:
|
||||
keycheck_runner.main()
|
||||
self.assertEqual(raised.exception.code, 1)
|
||||
preflight.assert_called_once_with(
|
||||
args.config,
|
||||
{'global': layout, 'keychecks': {}},
|
||||
authority_profile='server',
|
||||
)
|
||||
self.assertEqual([call.args[0] for call in run.call_args_list], requested)
|
||||
summary.assert_called_once()
|
||||
self.assertEqual(summary.call_args.args[1], list(keycheck_runner.SERVICES))
|
||||
ingest.assert_not_called()
|
||||
repair.assert_not_called()
|
||||
|
||||
|
||||
class ManifestCoverageTests(unittest.TestCase):
|
||||
def test_manifest_covers_every_lifecycle_authority_module(self):
|
||||
manifest = build_code_manifest()
|
||||
required = {
|
||||
'owned_process.py', 'supervisor.py', 'console_runner.py', 'scanner.py',
|
||||
'docker_shadow.py',
|
||||
'keycheck_runner.py', 'dashboard.py', 'postgres_runtime.py',
|
||||
'process_identity.py', 'runtime_security.py', 'scanner_db.py',
|
||||
'db_backend.py', 'result_spool.py', 'target_identity.py',
|
||||
'audit_github_tokens.py', 'sync_alive_github_tokens.py',
|
||||
'runtime_document.py', 'capacity_model.py', 'runtime_document_io.py',
|
||||
'managed_files.py',
|
||||
'host_agent_client.py', 'host_agent_protocol.py',
|
||||
'host_agent_reconcile.py', 'host_agent_server.py',
|
||||
'host_agent_apply.py', 'host_agent_lifecycle.py',
|
||||
'host_agent_runtime.py', 'host_agent_state.py',
|
||||
'../runtime/check-openrouter-keys.ps1',
|
||||
'../start_core_runtime.ps1', '../start_runtime.ps1', '../stop_runtime.ps1',
|
||||
}
|
||||
self.assertTrue(required.issubset(manifest['files']))
|
||||
|
||||
def test_manifest_covers_every_leaf_provider_common_module_and_configured_binary(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
executable = os.path.join(temp_dir, 'trufflehog.exe')
|
||||
Path(executable).write_bytes(b'fixture')
|
||||
git_executable = os.path.join(temp_dir, 'git.exe')
|
||||
Path(git_executable).write_bytes(b'git fixture')
|
||||
manifest = build_code_manifest(trufflehog_path=executable, git_path=git_executable)
|
||||
expected = {
|
||||
path.relative_to(APP_DIR).as_posix()
|
||||
for path in (APP_DIR / 'keycheckers').rglob('*.py')
|
||||
}
|
||||
self.assertTrue(expected.issubset(manifest['files']))
|
||||
self.assertEqual(
|
||||
manifest['executables']['trufflehog']['path'],
|
||||
lifecycle_authority.canonical_path(executable),
|
||||
)
|
||||
self.assertEqual(
|
||||
manifest['executables']['git']['path'],
|
||||
lifecycle_authority.canonical_path(git_executable),
|
||||
)
|
||||
|
||||
def test_every_leaf_provider_authenticates_as_its_first_main_action(self):
|
||||
violations = []
|
||||
for relative in keycheck_runner.SERVICES.values():
|
||||
path = APP_DIR / relative
|
||||
tree = ast.parse(path.read_text(encoding='utf-8'))
|
||||
main = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == 'main')
|
||||
first = main.body[0] if main.body else None
|
||||
call = first.value if isinstance(first, ast.Expr) else None
|
||||
if not isinstance(call, ast.Call) or not isinstance(call.func, ast.Name) or call.func.id != 'require_provider_authority':
|
||||
violations.append(path.relative_to(ROOT).as_posix())
|
||||
self.assertEqual(violations, [])
|
||||
|
||||
def test_unmanifested_trufflehog_binary_is_rejected_before_launch(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
trusted = os.path.join(temp_dir, 'trusted.exe')
|
||||
replacement = os.path.join(temp_dir, 'replacement.exe')
|
||||
Path(trusted).write_bytes(b'trusted')
|
||||
Path(replacement).write_bytes(b'replacement')
|
||||
metadata = {'code_manifest': build_code_manifest(trufflehog_path=trusted)}
|
||||
with mock.patch.object(scanner, 'require_active_supervisor_child', return_value=metadata):
|
||||
with self.assertRaisesRegex(RuntimeError, 'immutable supervisor authority'):
|
||||
scanner.require_trufflehog_launch_authority([replacement, '--version'])
|
||||
|
||||
def test_exact_acl_verification_covers_manifest_binary(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
executable = os.path.join(temp_dir, 'trufflehog.exe')
|
||||
Path(executable).write_bytes(b'fixture')
|
||||
manifest = build_code_manifest(trufflehog_path=executable)
|
||||
|
||||
def acl_ready(path):
|
||||
return lifecycle_authority.canonical_path(path) != lifecycle_authority.canonical_path(executable)
|
||||
|
||||
with mock.patch.object(lifecycle_authority, 'private_file_ready', side_effect=acl_ready):
|
||||
with self.assertRaisesRegex(LifecycleAuthorityError, 'executable:trufflehog'):
|
||||
lifecycle_authority.verify_code_manifest(manifest, require_private_acl=True)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user