sashatrask
338 lines
16 KiB
Python
338 lines
16 KiB
Python
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()
|