import builtins import hashlib import json import os import sqlite3 import sys import tempfile import unittest from pathlib import Path from unittest import mock ROOT = Path(__file__).resolve().parents[1] APP_DIR = ROOT / "app" sys.path.insert(0, str(APP_DIR)) import keycheck_runner import runtime_security import scanner from keycheckers import keycheck_common from keycheckers.gemini import geminiKeycheck as gemini from scanner_db import ScannerDB class JsonlGenerationTests(unittest.TestCase): @classmethod def setUpClass(cls): scanner.initialize_scanner_runtime(preflight_complete=True, register_cleanup=False) @staticmethod def reader_environment(state_dir): return mock.patch.dict(os.environ, { "KEYCHECK_OUTPUT_DIR": state_dir, "KEYCHECK_STATE_DIR": state_dir, "KEYCHECK_SERVICE": "fixture", "KEYCHECK_INPUT_TAIL_BYTES": "0", }) @staticmethod def write_private(path, payload): Path(path).write_bytes(payload) runtime_security.harden_private_file(path) def test_rotations_and_equal_prefix_recovery_publish_new_generations(self): cases = ( ("found_secrets.jsonl", scanner.rotate_jsonl_if_needed, scanner.reconcile_jsonl_segments), ("geminiResults.jsonl", keycheck_common.rotate_jsonl_if_needed, keycheck_common.reconcile_keycheck_jsonl_segments), ) for name, rotate, reconcile in cases: with self.subTest(kind=name), tempfile.TemporaryDirectory() as temp_dir: runtime_security.ensure_private_directory(temp_dir, reject_reparse=True) path = os.path.join(temp_dir, name) self.write_private(path, b'{"id":"rotation"}\n') before = os.stat(path, follow_symlinks=False) rotate(path, 1) after = os.stat(path, follow_symlinks=False) self.assertNotEqual((before.st_dev, before.st_ino), (after.st_dev, after.st_ino)) self.assertEqual(Path(path).read_bytes(), b"") next_segment = ( scanner.next_jsonl_segment_path if rotate is scanner.rotate_jsonl_if_needed else keycheck_common.next_jsonl_segment_path ) segment = next_segment(path, {})[0] self.write_private(segment, b'{"id":"prefix"}\n') self.write_private(path, b'{"id":"prefix"}\n') prefix_before = os.stat(path, follow_symlinks=False) reconcile(path) prefix_after = os.stat(path, follow_symlinks=False) self.assertNotEqual( (prefix_before.st_dev, prefix_before.st_ino), (prefix_after.st_dev, prefix_after.st_ino), ) self.assertEqual(Path(path).read_bytes(), b"") def test_same_size_truncate_and_regrow_replays_from_zero(self): with tempfile.TemporaryDirectory() as temp_dir: runtime_security.ensure_private_directory(temp_dir, reject_reparse=True) state_dir = os.path.join(temp_dir, "state") runtime_security.ensure_private_directory(state_dir, reject_reparse=True) path = os.path.join(temp_dir, "found_secrets.jsonl") old_row = b'{"id":"old"}\n' new_row = b'{"id":"new"}\n' self.assertEqual(len(old_row), len(new_row)) self.write_private(path, old_row) with self.reader_environment(state_dir): self.assertEqual( [item["data"]["id"] for item in keycheck_common.iter_jsonl_input(path)], ["old"], ) state_path = os.path.join(state_dir, "input_state.json") checkpoint = json.loads(Path(state_path).read_text(encoding="utf-8")) with open(path, "r+b") as handle: handle.truncate(0) handle.write(new_row) handle.flush() os.fsync(handle.fileno()) changed_mtime = int(checkpoint["mtime_ns"]) + 2_000_000_000 os.utime(path, ns=(changed_mtime, changed_mtime)) with self.reader_environment(state_dir): rows = [item["data"]["id"] for item in keycheck_common.iter_jsonl_input(path)] self.assertEqual(rows, ["new"]) def test_plain_reader_fstats_the_descriptor_opened_after_a_swap(self): with tempfile.TemporaryDirectory() as temp_dir: runtime_security.ensure_private_directory(temp_dir, reject_reparse=True) state_dir = os.path.join(temp_dir, "state") runtime_security.ensure_private_directory(state_dir, reject_reparse=True) path = os.path.join(temp_dir, "found_secrets.jsonl") replacement = os.path.join(temp_dir, "replacement.jsonl") self.write_private(path, b'{"id":"old"}\n') with self.reader_environment(state_dir): list(keycheck_common.iter_jsonl_input(path)) self.write_private(replacement, b'{"id":"new"}\n') real_open = builtins.open swapped = {"done": False} def swap_before_open(candidate, mode="r", *args, **kwargs): if ( not swapped["done"] and mode == "rb" and os.path.normcase(os.path.abspath(os.fspath(candidate))) == os.path.normcase(os.path.abspath(path)) ): swapped["done"] = True os.replace(replacement, path) return real_open(candidate, mode, *args, **kwargs) with self.reader_environment(state_dir), mock.patch("builtins.open", side_effect=swap_before_open): rows = [item["data"]["id"] for item in keycheck_common.iter_jsonl_input(path)] self.assertEqual(rows, ["new"]) def test_segmented_reader_fstats_the_descriptor_opened_after_a_swap(self): with tempfile.TemporaryDirectory() as temp_dir: runtime_security.ensure_private_directory(temp_dir, reject_reparse=True) state_dir = os.path.join(temp_dir, "state") runtime_security.ensure_private_directory(state_dir, reject_reparse=True) current = os.path.join(temp_dir, "found_secrets.jsonl") segment = os.path.join(temp_dir, "found_secrets.000001.jsonl") replacement = os.path.join(temp_dir, "replacement.jsonl") self.write_private(segment, b'{"id":"segment"}\n') self.write_private(current, b'{"id":"old"}\n') with self.reader_environment(state_dir): list(keycheck_common.iter_jsonl_input(current)) self.write_private(replacement, b'{"id":"new"}\n') real_open = builtins.open swapped = {"done": False} def swap_before_open(candidate, mode="r", *args, **kwargs): if ( not swapped["done"] and mode == "rb" and os.path.normcase(os.path.abspath(os.fspath(candidate))) == os.path.normcase(os.path.abspath(current)) ): swapped["done"] = True os.replace(replacement, current) return real_open(candidate, mode, *args, **kwargs) with self.reader_environment(state_dir), mock.patch("builtins.open", side_effect=swap_before_open): rows = [item["data"]["id"] for item in keycheck_common.iter_jsonl_input(current)] self.assertEqual(rows, ["new"]) def test_result_ingestion_observes_a_rotated_active_generation(self): clean_database_env = { "SCANNER_DB_URL": "", "DATABASE_URL": "", "TRUF_MANAGED_POSTGRES_DSN": "", "KEYCHECK_DB_URL": "", } with tempfile.TemporaryDirectory() as temp_dir, mock.patch.dict( os.environ, clean_database_env, clear=False, ): runtime_security.ensure_private_directory(temp_dir, reject_reparse=True) keycheck_dir = os.path.join(temp_dir, "keychecks") service_dir = os.path.join(keycheck_dir, "openai") runtime_security.ensure_private_directory(keycheck_dir, reject_reparse=True) runtime_security.ensure_private_directory(service_dir, reject_reparse=True) db_path = os.path.join(temp_dir, "scanner.db") db = ScannerDB(db_path=db_path) db.close() results_path = os.path.join(service_dir, "openaiResults.jsonl") def event(event_id, key): digest = hashlib.sha256(key.encode("utf-8")).hexdigest() return { "event_id": event_id, "status": "ALIVE", "key_hash": digest, "secret_hash": digest, "key_masked": key, "checked_at": "2026-07-19T00:00:00+00:00", } old_event = event("old-generation", "old-key") self.write_private(results_path, (json.dumps(old_event) + "\n").encode("utf-8")) layout = {"database_path": db_path, "keycheck_dir": keycheck_dir} self.assertEqual(keycheck_runner.ingest_keycheck_results_to_db(layout, ["openai"], 10), 1) before = os.stat(results_path, follow_symlinks=False) keycheck_common.rotate_jsonl_if_needed(results_path, 1) after = os.stat(results_path, follow_symlinks=False) self.assertNotEqual((before.st_dev, before.st_ino), (after.st_dev, after.st_ino)) new_event = event("new-generation", "new-key") with open(results_path, "ab") as handle: handle.write((json.dumps(new_event) + "\n").encode("utf-8")) handle.flush() os.fsync(handle.fileno()) runtime_security.harden_private_file(results_path) self.assertEqual(keycheck_runner.ingest_keycheck_results_to_db(layout, ["openai"], 10), 1) connection = sqlite3.connect(db_path) try: event_ids = {row[0] for row in connection.execute("SELECT event_id FROM keycheck_results")} finally: connection.close() self.assertEqual(event_ids, {"old-generation", "new-generation"}) def test_scanner_rotation_fault_stages_keep_the_closed_rows_visible(self): for stage in ("active_replace", "final_manifest"): with self.subTest(stage=stage), tempfile.TemporaryDirectory() as temp_dir: runtime_security.ensure_private_directory(temp_dir, reject_reparse=True) path = os.path.join(temp_dir, "scan_results.jsonl") self.write_private(path, b'{"scan_event_id":"retained"}\n') if stage == "active_replace": patch = mock.patch.object(scanner, "durable_replace", side_effect=OSError("crash")) else: real_write = scanner.write_jsonl_manifest calls = {"count": 0} def fail_final_manifest(candidate, manifest): calls["count"] += 1 if calls["count"] == 2: raise OSError("crash") return real_write(candidate, manifest) patch = mock.patch.object(scanner, "write_jsonl_manifest", side_effect=fail_final_manifest) with patch, self.assertRaisesRegex(OSError, "crash"): scanner.rotate_jsonl_if_needed(path, 1) rows = [] for candidate in scanner.projection_segment_paths(path): rows.extend(Path(candidate).read_text(encoding="utf-8").splitlines()) self.assertTrue(any(json.loads(row)["scan_event_id"] == "retained" for row in rows)) class GeminiStatusAtomicityTests(unittest.TestCase): KEY = "AIza" + ("a" * 35) @staticmethod def status_paths(root): return { status: os.path.join(root, f"{index:02d}-{status}.txt") for index, status in enumerate(gemini.STATUS_FILES) } def seed_every_status(self, paths): for path in paths.values(): Path(path).write_text(f"{self.KEY}:[old-model]:old:RATE_LIMITED\n", encoding="utf-8") runtime_security.harden_private_file(path) def count_classifications(self, paths): matches = [] for status, path in paths.items(): if not os.path.exists(path): continue for line in Path(path).read_text(encoding="utf-8").splitlines(): if gemini.key_from_line(line) == self.KEY: matches.append((status, line)) return matches @staticmethod def invalid_result(): return {"status": "INVALID", "error": {"message": "invalid fixture"}} def test_normal_reclassification_leaves_exactly_one_target_row(self): with tempfile.TemporaryDirectory() as temp_dir: runtime_security.ensure_private_directory(temp_dir, reject_reparse=True) paths = self.status_paths(temp_dir) self.seed_every_status(paths) with mock.patch.object(gemini, "STATUS_FILES", paths): gemini.append_status_file(self.KEY, self.invalid_result()) matches = self.count_classifications(paths) self.assertEqual(matches, [("INVALID", f"{self.KEY}\tINVALID\tinvalid fixture")]) def test_every_snapshot_write_stage_retains_an_existing_classification(self): stage_count = len(gemini.STATUS_FILES) for failed_stage in range(1, stage_count + 1): with self.subTest(stage=failed_stage), tempfile.TemporaryDirectory() as temp_dir: runtime_security.ensure_private_directory(temp_dir, reject_reparse=True) paths = self.status_paths(temp_dir) self.seed_every_status(paths) real_replace_snapshot = gemini._replace_status_snapshot calls = {"count": 0} def fail_snapshot(path, lines): calls["count"] += 1 if calls["count"] == failed_stage: raise OSError("injected snapshot write failure") return real_replace_snapshot(path, lines) with mock.patch.object(gemini, "STATUS_FILES", paths), \ mock.patch.object(gemini, "_replace_status_snapshot", side_effect=fail_snapshot): with self.assertRaisesRegex(OSError, "snapshot write failure"): gemini.append_status_file(self.KEY, self.invalid_result()) self.assertGreaterEqual(len(self.count_classifications(paths)), 1) def test_every_durable_replace_stage_retains_an_existing_classification(self): stage_count = len(gemini.STATUS_FILES) for failed_stage in range(1, stage_count + 1): with self.subTest(stage=failed_stage), tempfile.TemporaryDirectory() as temp_dir: runtime_security.ensure_private_directory(temp_dir, reject_reparse=True) paths = self.status_paths(temp_dir) self.seed_every_status(paths) writer_module = sys.modules[gemini.private_atomic_writer.__module__] real_durable_replace = writer_module.durable_replace calls = {"count": 0} def fail_replace(source, destination): calls["count"] += 1 if calls["count"] == failed_stage: raise OSError("injected durable replace failure") return real_durable_replace(source, destination) with mock.patch.object(gemini, "STATUS_FILES", paths), \ mock.patch.object(writer_module, "durable_replace", side_effect=fail_replace): with self.assertRaisesRegex(OSError, "durable replace failure"): gemini.append_status_file(self.KEY, self.invalid_result()) self.assertGreaterEqual(len(self.count_classifications(paths)), 1) if __name__ == "__main__": unittest.main()