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