import os from pathlib import Path import sqlite3 import subprocess import sys import tempfile 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 owned_process import scanner class ScanSlotCleanupTests(unittest.TestCase): def test_fresh_dead_slot_is_removed_but_live_and_unknown_slots_remain(self): now = time.time() conn = sqlite3.connect(':memory:') self.addCleanup(conn.close) conn.executescript(scanner.SCAN_SLOT_SCHEMA) for slot_id, owner_pid, child_pid in ( ('dead', 101, 102), ('live', 201, 202), ('unknown', 301, 302), ): 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 (?, ?, 1, 'test', 'owner-created', 'owner.exe', ?, 'child-created', 'child.exe', ?, ?)''', (slot_id, owner_pid, child_pid, now - 1, now - 1), ) identity_state = { 101: False, 102: False, 201: True, 202: False, 301: False, 302: None, } with mock.patch.object( scanner, 'exact_process_identity_live', side_effect=lambda pid, creation, executable: identity_state[pid], ): scanner.cleanup_stale_scan_slots(conn, now, stale_sec=60) remaining = {row[0] for row in conn.execute('SELECT slot_id FROM scan_slots')} self.assertEqual(remaining, {'live', 'unknown'}) class WindowsJobMemoryConfigTests(unittest.TestCase): def test_job_limit_structure_combines_memory_and_kill_on_close_flags(self): limit = 4096 * 1024 * 1024 limits = owned_process._windows_job_limit_information(limit) flags = limits.BasicLimitInformation.LimitFlags self.assertEqual(limits.JobMemoryLimit, limit) self.assertEqual( flags, owned_process._JOB_OBJECT_LIMIT_JOB_MEMORY | owned_process._JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, ) unlimited = owned_process._windows_job_limit_information(0) self.assertEqual( unlimited.BasicLimitInformation.LimitFlags, owned_process._JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, ) self.assertEqual(unlimited.JobMemoryLimit, 0) def test_owned_process_rejects_invalid_job_memory_limits_before_launch(self): with mock.patch.object(owned_process.subprocess, 'Popen') as popen: for value in (-1, None, True, 1.5, '4096'): with self.subTest(value=value): with self.assertRaises(ValueError): owned_process.OwnedProcess(['fixture'], job_memory_limit_bytes=value) popen.assert_not_called() def test_owned_process_rejects_invalid_resource_policy_before_launch(self): with mock.patch.object(owned_process.subprocess, 'Popen') as popen: for name, values in ( ('job_cpu_weight', (-1, 10, True, '2')), ('process_memory_priority', (-1, 6, True, '4')), ): for value in values: with self.subTest(name=name, value=value): with self.assertRaises(ValueError): owned_process.OwnedProcess(['fixture'], **{name: value}) popen.assert_not_called() def test_canonical_job_memory_limit_is_applied(self): limit = 4096 * 1024 * 1024 config = console_runner.load_config(str(APP_DIR / 'config.yaml')) self.assertEqual(config['global']['trufflehog_job_memory_limit_bytes'], limit) self.assertEqual(config['global']['trufflehog_windows_job_cpu_weight'], 2) self.assertEqual(config['global']['trufflehog_windows_memory_priority'], 4) docker_args = console_runner.build_args_from_source_config( 'dockerhub', config['sources']['dockerhub'], config['global'], 'fixture', ) self.assertEqual( docker_args.trufflehog_job_memory_limit_bytes, 6 * 1024 * 1024 * 1024, ) previous = scanner.scan_config.__dict__.copy() try: console_runner.apply_global_config({ 'trufflehog_job_memory_limit_bytes': limit, 'trufflehog_windows_job_cpu_weight': 2, 'trufflehog_windows_memory_priority': 4, }) self.assertEqual(scanner.scan_config.trufflehog_job_memory_limit_bytes, limit) self.assertEqual(scanner.scan_config.trufflehog_windows_job_cpu_weight, 2) self.assertEqual(scanner.scan_config.trufflehog_windows_memory_priority, 4) finally: scanner.scan_config.__dict__.clear() scanner.scan_config.__dict__.update(previous) class RunCommandDiskPressureTests(unittest.TestCase): def test_disk_pressure_kills_owned_process_and_reports_source_resource(self): class MockProcess: pid = 4321 job_membership_verified = True def __init__(self): self.returncode = None self.killed = False def poll(self): return self.returncode def kill(self): self.killed = True self.returncode = 1 def wait(self, timeout=None): return self.returncode process = MockProcess() memory_limit = 4096 * 1024 * 1024 with tempfile.TemporaryDirectory() as temp_dir: command_work_dir = os.path.join(temp_dir, 'trufflehog-run-test') os.mkdir(command_work_dir) with mock.patch.object(scanner, 'acquire_scan_slot', return_value=None), \ mock.patch.object(scanner, 'create_command_work_dir', return_value=command_work_dir), \ mock.patch.object(scanner, '_shared_staging_owners', return_value=[]), \ mock.patch.object(scanner, 'cleanup_command_work_dir'), \ mock.patch.object(scanner, 'require_trufflehog_launch_authority'), \ mock.patch.object(scanner, 'write_temp_owner'), \ mock.patch.object(scanner, 'OwnedProcess', return_value=process) as process_factory, \ mock.patch.object(scanner.shutil, 'disk_usage', return_value=SimpleNamespace(free=1)), \ mock.patch.object(scanner.scan_config, 'min_free_gb', 1), \ mock.patch.object(scanner.scan_config, 'trufflehog_job_memory_limit_bytes', memory_limit), \ mock.patch.object(scanner.scan_config, 'trufflehog_windows_job_cpu_weight', 2), \ mock.patch.object(scanner.scan_config, 'trufflehog_windows_memory_priority', 4): stdout, stderr, returncode = scanner.run_command(['fixture'], timeout_sec=30) process_factory.assert_called_once() self.assertTrue(process.killed) self.assertEqual(stdout, '') self.assertEqual(returncode, -1) self.assertIn('Not enough free space on configured TruffleHog work volume', stderr) if os.name == 'nt': self.assertEqual(process_factory.call_args.kwargs['job_memory_limit_bytes'], memory_limit) self.assertEqual(process_factory.call_args.kwargs['job_cpu_weight'], 2) self.assertEqual(process_factory.call_args.kwargs['process_memory_priority'], 4) creationflags = process_factory.call_args.kwargs['creationflags'] self.assertTrue(creationflags & subprocess.CREATE_NEW_PROCESS_GROUP) self.assertTrue(creationflags & subprocess.CREATE_NO_WINDOW) self.assertFalse( creationflags & getattr(subprocess, 'BELOW_NORMAL_PRIORITY_CLASS', 0x00004000) ) result = scanner.apply_trufflehog_diagnostics({}, stderr, returncode, 'git') self.assertEqual(result['error_class'], 'source_resource') self.assertEqual(result['source_failure_category'], 'source_resource') def test_opportunistic_capacity_gate_reserves_job_and_overhead(self): gib = 1024 ** 3 previous = scanner.scan_config.__dict__.copy() try: scanner.scan_config.opportunistic_scan_slots = 1 scanner.scan_config.opportunistic_scan_sources = ['github'] scanner.scan_config.trufflehog_job_memory_limit_bytes = 4 * gib scanner.scan_config.opportunistic_scan_reserve_overhead_bytes = 1 * gib scanner.scan_config.opportunistic_scan_min_available_after_reserve_bytes = 4 * gib scanner.scan_config.opportunistic_scan_min_commit_after_reserve_bytes = 6 * gib with mock.patch.object(scanner, 'windows_scan_capacity_snapshot', return_value={ 'available_physical_bytes': 9 * gib, 'commit_headroom_bytes': 11 * gib, }): self.assertTrue(scanner.opportunistic_scan_slot_allowed('github')) with mock.patch.object(scanner, 'windows_scan_capacity_snapshot', return_value={ 'available_physical_bytes': 9 * gib - 1, 'commit_headroom_bytes': 11 * gib, }): self.assertFalse(scanner.opportunistic_scan_slot_allowed('github')) with mock.patch.object(scanner, 'windows_scan_capacity_snapshot', side_effect=OSError('probe')): self.assertFalse(scanner.opportunistic_scan_slot_allowed('github')) self.assertFalse(scanner.opportunistic_scan_slot_allowed('dockerhub')) finally: scanner.scan_config.__dict__.clear() scanner.scan_config.__dict__.update(previous) def test_disk_monitor_failure_kills_owned_process_fail_closed(self): class MockProcess: pid = 4322 job_membership_verified = True def __init__(self): self.returncode = None self.killed = False def poll(self): return self.returncode def kill(self): self.killed = True self.returncode = 1 def wait(self, timeout=None): return self.returncode process = MockProcess() with tempfile.TemporaryDirectory() as temp_dir: command_work_dir = os.path.join(temp_dir, 'trufflehog-run-test') os.mkdir(command_work_dir) with mock.patch.object(scanner, 'acquire_scan_slot', return_value=None), \ mock.patch.object(scanner, 'create_command_work_dir', return_value=command_work_dir), \ mock.patch.object(scanner, '_shared_staging_owners', return_value=[]), \ mock.patch.object(scanner, 'cleanup_command_work_dir'), \ mock.patch.object(scanner, 'require_trufflehog_launch_authority'), \ mock.patch.object(scanner, 'write_temp_owner'), \ mock.patch.object(scanner, 'OwnedProcess', return_value=process) as process_factory, \ mock.patch.object(scanner.shutil, 'disk_usage', side_effect=OSError('probe failed')), \ mock.patch.object(scanner.scan_config, 'min_free_gb', 1), \ mock.patch.object(scanner.scan_config, 'trufflehog_job_memory_limit_bytes', 4096 * 1024 * 1024): _, stderr, returncode = scanner.run_command(['fixture'], timeout_sec=30) process_factory.assert_called_once() self.assertTrue(process.killed) self.assertEqual(returncode, -1) self.assertIn('Unable to monitor configured TruffleHog work volume', stderr) if __name__ == '__main__': unittest.main()