Privatae
400 lines
18 KiB
Python
400 lines
18 KiB
Python
"""Unit tests for the sanitize stage. Run: python3 -m unittest discover .
|
|
|
|
The real pii pipeline and vault run in-process; the recall backend
|
|
is a fake (no model), the service path is a stub HTTP server speaking the
|
|
real /sterilize wire format. Nothing here needs a model or the network.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import shutil
|
|
import stat
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from unittest import mock
|
|
|
|
from . import sanitize as sz
|
|
from .extract import extract
|
|
from .manifest import validate_document
|
|
from .walk import walk
|
|
|
|
IS_WINDOWS = os.name == "nt"
|
|
|
|
|
|
def manifest(root, mode="flag", verify=False, scope=None):
|
|
san = {"mode": mode, "verify": verify}
|
|
if scope:
|
|
san["scope"] = scope
|
|
return validate_document({
|
|
"freeagent": 1, "id": "t", "owner": {"host": "h"},
|
|
"scope": {"roots": [root]},
|
|
"mask": ["read"], "stages": ["walk", "extract", "sanitize"], "sweep": "manual",
|
|
"sanitize": san,
|
|
"sink": {"catalog": "local", "vectors": "local", "plaintext": "never"},
|
|
})
|
|
|
|
|
|
def put(root, rel, text):
|
|
path = os.path.join(root, *rel.split("/"))
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
with open(path, "w", encoding="utf-8") as fh:
|
|
fh.write(text)
|
|
|
|
|
|
class FakeBackend:
|
|
name = "fake"
|
|
calls = 0
|
|
|
|
def __init__(self, findings=(), fail=False):
|
|
self.findings = list(findings)
|
|
self.fail = fail
|
|
|
|
def describe(self):
|
|
return "fake"
|
|
|
|
def extract(self, text):
|
|
FakeBackend.calls += 1
|
|
if self.fail:
|
|
_, recall, _ = sz.pii_modules()
|
|
raise recall.RecallError("model unreachable")
|
|
return list(self.findings) # proposes everything; grounding must drop the ghosts
|
|
|
|
def verify(self, text):
|
|
return []
|
|
|
|
|
|
class Case(unittest.TestCase):
|
|
def setUp(self):
|
|
self.root = tempfile.mkdtemp(prefix="fa-sz-")
|
|
self.out = tempfile.mkdtemp(prefix="fa-szout-")
|
|
put(self.root, "note.txt", "Marisol's email is m@example.com, card 4111 1111 1111 1111.\n")
|
|
put(self.root, "twin.txt", "Marisol's email is m@example.com, card 4111 1111 1111 1111.\n")
|
|
put(self.root, "clean.txt", "nothing personal here\n")
|
|
put(self.root, "blob.bin", "\x00\x01")
|
|
FakeBackend.calls = 0
|
|
|
|
def tearDown(self):
|
|
shutil.rmtree(self.root, ignore_errors=True)
|
|
shutil.rmtree(self.out, ignore_errors=True)
|
|
|
|
def prep(self, m):
|
|
walk(m, self.out)
|
|
extract(m, self.out, use_external=False)
|
|
|
|
def run_stage(self, m, env=None, backend=None):
|
|
self.prep(m)
|
|
if backend is not None:
|
|
with mock.patch.object(sz, "BUILD_BACKEND", lambda e: backend):
|
|
s = sz.sanitize(m, self.out, env=env or {})
|
|
else:
|
|
s = sz.sanitize(m, self.out, env=env or {})
|
|
with open(os.path.join(self.out, "sanitize.jsonl"), encoding="utf-8") as fh:
|
|
idx = {json.loads(l)["path"]: json.loads(l) for l in fh}
|
|
return s, idx
|
|
|
|
def text_of(self, rec):
|
|
with open(sz.store_path(self.out, rec["sha256"]), encoding="utf-8") as fh:
|
|
return fh.read()
|
|
|
|
|
|
class FlagModeTests(Case):
|
|
def test_rules_only_redaction(self):
|
|
s, idx = self.run_stage(manifest(self.root, "flag"))
|
|
out = self.text_of(idx["note.txt"])
|
|
self.assertNotIn("m@example.com", out)
|
|
for card in ("4111 1111 1111 1111", "4111111111111111"):
|
|
self.assertNotIn(card, out)
|
|
self.assertIn("Marisol", out) # contextual PII is NOT caught by flag mode
|
|
self.assertIn("[PII_EMAIL_1]", out)
|
|
self.assertIn("[PII_CREDIT_CARD_1]", out)
|
|
self.assertEqual(idx["note.txt"]["types"], {"EMAIL": 1, "CREDIT_CARD": 1})
|
|
self.assertEqual(s["sterilizer"], "flag:rules-only")
|
|
self.assertEqual(s["types"], {"EMAIL": 1, "CREDIT_CARD": 1})
|
|
# twin shares content -> deduped, counted once as ok
|
|
self.assertEqual(idx["twin.txt"]["sha256"], idx["note.txt"]["sha256"])
|
|
self.assertEqual(s["counts"]["ok"], 3)
|
|
self.assertEqual(idx["blob.bin"]["status"], "skipped")
|
|
self.assertEqual(s["counts"]["skipped"], 1)
|
|
|
|
def test_index_and_summary_carry_no_values(self):
|
|
s, idx = self.run_stage(manifest(self.root, "flag"))
|
|
blob = json.dumps(idx) + json.dumps(s)
|
|
self.assertNotIn("m@example.com", blob)
|
|
# the whole card: four digits alone also turn up inside hashes and times
|
|
for card in ("4111 1111 1111 1111", "4111111111111111"):
|
|
self.assertNotIn(card, blob)
|
|
|
|
@unittest.skipIf(IS_WINDOWS, "POSIX permission bits")
|
|
def test_vault_is_private(self):
|
|
self.run_stage(manifest(self.root, "flag"))
|
|
mode = stat.S_IMODE(os.stat(os.path.join(self.out, "sanitize", "vault.db")).st_mode)
|
|
self.assertEqual(mode, 0o600)
|
|
|
|
def test_reuse_and_invalidate(self):
|
|
s1, _ = self.run_stage(manifest(self.root, "flag"))
|
|
s2, idx = self.run_stage(manifest(self.root, "flag"))
|
|
self.assertEqual(s2["counts"]["reused"], s1["counts"]["ok"])
|
|
self.assertTrue(idx["note.txt"]["reused"])
|
|
# a different scope is a different cache key -> redone
|
|
s3, _ = self.run_stage(manifest(self.root, "flag", scope="other"))
|
|
self.assertEqual(s3["counts"]["reused"], 0)
|
|
# a changed file is redone
|
|
put(self.root, "clean.txt", "now with a@b.co\n")
|
|
s4, idx4 = self.run_stage(manifest(self.root, "flag"))
|
|
self.assertIn("[PII_EMAIL_", self.text_of(idx4["clean.txt"]))
|
|
|
|
def test_result_from_an_older_pipeline_version_is_not_reused(self):
|
|
# the version is the only reuse-key part that changes with the
|
|
# pipeline: output written by older code must be redone on upgrade
|
|
with mock.patch.object(sz, "SANITIZE_VERSION", 1):
|
|
s1, _ = self.run_stage(manifest(self.root, "flag"))
|
|
self.assertGreaterEqual(sz.SANITIZE_VERSION, 2)
|
|
s2, idx = self.run_stage(manifest(self.root, "flag"))
|
|
self.assertEqual(s2["counts"]["reused"], 0)
|
|
self.assertEqual(s2["counts"]["ok"], s1["counts"]["ok"])
|
|
self.assertFalse(idx["note.txt"].get("reused", False))
|
|
|
|
def test_text_kept_by_the_older_rules_is_redone(self):
|
|
# version 2 kept an international number and a lowercase IBAN in the
|
|
# clear; a store it wrote must not be reused by the wider rules
|
|
put(self.root, "intl.txt", "Mobile +44 7700 900123, iban gb82west12345698765432\n")
|
|
with mock.patch.object(sz, "SANITIZE_VERSION", 2):
|
|
self.run_stage(manifest(self.root, "flag"))
|
|
s, idx = self.run_stage(manifest(self.root, "flag"))
|
|
self.assertEqual(s["counts"]["reused"], 0)
|
|
out = self.text_of(idx["intl.txt"])
|
|
self.assertNotIn("7700 900123", out)
|
|
self.assertNotIn("gb82west", out)
|
|
self.assertEqual(idx["intl.txt"]["types"], {"PHONE": 1, "BANK_ACCOUNT": 1})
|
|
|
|
def test_tokens_stable_across_files_in_scope(self):
|
|
put(self.root, "again.txt", "ping m@example.com again\n")
|
|
_, idx = self.run_stage(manifest(self.root, "flag"))
|
|
self.assertIn("[PII_EMAIL_1]", self.text_of(idx["again.txt"]))
|
|
self.assertIn("[PII_EMAIL_1]", self.text_of(idx["note.txt"]))
|
|
|
|
def test_vault_dir_env_keeps_vault_out_of_out(self):
|
|
vdir = tempfile.mkdtemp(prefix="fa-vault-")
|
|
try:
|
|
s, idx = self.run_stage(manifest(self.root, "flag"), env={"FREEAGENT_VAULT_DIR": vdir})
|
|
self.assertTrue(os.path.exists(os.path.join(vdir, "vault.db")))
|
|
self.assertFalse(os.path.exists(os.path.join(self.out, "sanitize", "vault.db")))
|
|
self.assertIn("[PII_EMAIL_1]", self.text_of(idx["note.txt"]))
|
|
finally:
|
|
shutil.rmtree(vdir, ignore_errors=True)
|
|
|
|
def test_requires_extract_first(self):
|
|
with self.assertRaises(FileNotFoundError):
|
|
sz.sanitize(manifest(self.root, "flag"), self.out, env={})
|
|
|
|
|
|
class SterilizeModeTests(Case):
|
|
def test_recall_backend_findings_are_tokenized(self):
|
|
fake = FakeBackend([("Marisol", "PERSON"), ("Ghost", "PERSON")])
|
|
s, idx = self.run_stage(manifest(self.root, "sterilize"), backend=fake)
|
|
out = self.text_of(idx["note.txt"])
|
|
self.assertNotIn("Marisol", out)
|
|
self.assertIn("[PII_PERSON_1]", out)
|
|
self.assertEqual(idx["note.txt"]["types"]["PERSON"], 1)
|
|
self.assertEqual(idx["note.txt"]["stats"]["llm_ungrounded"], 1)
|
|
self.assertEqual(s["sterilizer"], "sterilize:fake")
|
|
self.assertEqual(FakeBackend.calls, 2) # note (twin deduped) + clean
|
|
|
|
def test_fail_closed_stops_the_stage(self):
|
|
s, idx = self.run_stage(manifest(self.root, "sterilize"), backend=FakeBackend(fail=True))
|
|
self.assertEqual(s["counts"]["ok"], 0)
|
|
self.assertEqual(s["counts"]["not_run"], 3)
|
|
self.assertIn("model unreachable", s["stopped"])
|
|
self.assertEqual(idx["note.txt"]["status"], "not_run")
|
|
self.assertFalse(os.path.exists(sz.store_path(self.out, idx["note.txt"]["sha256"])))
|
|
self.assertEqual(FakeBackend.calls, 1) # stopped after the first failure
|
|
|
|
def test_backend_that_cannot_start_raises(self):
|
|
self.prep(manifest(self.root, "sterilize"))
|
|
with self.assertRaises(sz.SanitizeError):
|
|
sz.sanitize(manifest(self.root, "sterilize"), self.out, env={"PII_RECALL": "nope"})
|
|
|
|
def test_chunking_round_trip(self):
|
|
big = ("para %d with m@example.com\n\n" % i for i in range(2000))
|
|
text = "".join(big)
|
|
put(self.root, "big.txt", text)
|
|
chunks = sz.split_chunks(text, limit=5000)
|
|
self.assertGreater(len(chunks), 5)
|
|
self.assertEqual("".join(chunks), text)
|
|
self.assertTrue(all(len(c) <= 5000 for c in chunks))
|
|
sizes = []
|
|
real = sz.InProcess.sterilize
|
|
|
|
def spy(front, chunk, scope):
|
|
if chunk.startswith("para "): # big.txt's chunks, not the other files
|
|
sizes.append(len(chunk))
|
|
return real(front, chunk, scope)
|
|
|
|
with mock.patch.object(sz, "SANITIZE_CHUNK_CHARS", 5000), \
|
|
mock.patch.object(sz.InProcess, "sterilize", spy):
|
|
_, idx = self.run_stage(manifest(self.root, "flag"))
|
|
# the stage itself must have split at the patched limit, not the default
|
|
self.assertLessEqual(max(sizes), 5000)
|
|
self.assertEqual(len(sizes), len(chunks))
|
|
out = self.text_of(idx["big.txt"])
|
|
self.assertNotIn("m@example.com", out)
|
|
self.assertEqual(out.count("[PII_EMAIL_1]"), 2000) # one token, every chunk
|
|
|
|
def test_hard_cut_never_splits_a_placeholder(self):
|
|
_, _, Vault = sz.pii_modules()
|
|
from pii import pipeline
|
|
self.assertEqual(sz._PLACEHOLDER_RE.pattern, pipeline.TOKEN_RE.pattern)
|
|
vault = Vault(os.path.join(self.out, "v.db"))
|
|
vault.token_for("s", "EMAIL", "ana@x.org") # [PII_EMAIL_1] is live
|
|
# the corpus stage escapes every placeholder in a document, but only
|
|
# one that reaches the sterilizer whole: split in half, neither half
|
|
# is recognised and the literal is stored raw, to rehydrate later as
|
|
# the vault's value
|
|
text = "x" * 15 + "[PII_EMAIL_1]" + "y" * 10 # one line, no newline
|
|
chunks = sz.split_chunks(text, limit=20)
|
|
self.assertEqual("".join(chunks), text)
|
|
self.assertTrue(any("[PII_EMAIL_1]" in c for c in chunks), chunks)
|
|
out = "".join(pipeline.sterilize(c, "s", vault, lambda t: [], escape_held_literals=True)[0]
|
|
for c in chunks)
|
|
self.assertNotIn("[PII_EMAIL_1]", out)
|
|
self.assertEqual(pipeline.rehydrate(out, "s", vault)[0], text)
|
|
# a placeholder at the very start of an over-long piece stays whole
|
|
self.assertEqual(sz.split_chunks("[PII_PERSON_1]" + "z" * 30, limit=10)[0], "[PII_PERSON_1]")
|
|
|
|
def test_hard_cut_clears_a_placeholder_of_any_length(self):
|
|
long_ph = "[PII_" + "OTHER_" * 20 + "123456789]" # well over 64 characters
|
|
for before in (30, 90, 150):
|
|
with self.subTest(before=before):
|
|
text = "x" * before + long_ph + "y" * 40
|
|
chunks = sz.split_chunks(text, limit=100)
|
|
self.assertEqual("".join(chunks), text)
|
|
self.assertTrue(any(long_ph in c for c in chunks), [len(c) for c in chunks])
|
|
|
|
def test_corpus_document_quoting_a_live_placeholder_round_trips(self):
|
|
# one manifest-wide scope: a.txt mints [PII_EMAIL_1]; b.txt quotes that
|
|
# placeholder literally and must come back as written, not as a's email
|
|
for name in ("note.txt", "twin.txt", "clean.txt", "blob.bin"):
|
|
os.remove(os.path.join(self.root, name))
|
|
put(self.root, "a.txt", "mail ana@x.org\n")
|
|
quote = "the export shows [PII_EMAIL_1] where the address was\n"
|
|
put(self.root, "b.txt", quote)
|
|
_, idx = self.run_stage(manifest(self.root, "flag"))
|
|
_, _, Vault = sz.pii_modules()
|
|
from pii import pipeline
|
|
vault = Vault(os.path.join(self.out, "sanitize", "vault.db"))
|
|
self.assertEqual(vault.value_for("t", "[PII_EMAIL_1]"), "ana@x.org")
|
|
stored = self.text_of(idx["b.txt"])
|
|
self.assertNotIn("[PII_EMAIL_1]", stored)
|
|
self.assertEqual(pipeline.rehydrate(stored, "t", vault)[0], quote)
|
|
|
|
def test_split_chunks_hard_cuts_a_giant_line(self):
|
|
text = "x" * 12000 + "\n" + "y" * 100
|
|
chunks = sz.split_chunks(text, limit=5000)
|
|
self.assertEqual("".join(chunks), text)
|
|
self.assertTrue(all(len(c) <= 5000 for c in chunks))
|
|
self.assertEqual(sz.split_chunks("", 10), [])
|
|
self.assertEqual(sz.split_chunks("short", 10), ["short"])
|
|
|
|
|
|
class _StubPii(BaseHTTPRequestHandler):
|
|
seen = []
|
|
status = 200
|
|
echo = "honour" # "honour", "strip" (a service that predates the field) or "false"
|
|
|
|
def do_POST(self):
|
|
n = int(self.headers.get("Content-Length", "0"))
|
|
req = json.loads(self.rfile.read(n))
|
|
_StubPii.seen.append((self.path, self.headers.get("Authorization"), req))
|
|
if _StubPii.status != 200:
|
|
self.send_response(_StubPii.status)
|
|
self.send_header("Content-Length", "0")
|
|
self.end_headers()
|
|
return
|
|
body = json.dumps({"text": req["text"].replace("Marisol", "[PII_PERSON_1]"),
|
|
"findings": [{"token": "[PII_PERSON_1]", "type": "PERSON"}],
|
|
"stats": dict({"rules": 0, "llm": 1, "recall": "llm"}, **(
|
|
{} if _StubPii.echo == "strip" else
|
|
{"escape_held_literals": _StubPii.echo == "honour"
|
|
and req.get("escape_held_literals") is True}))}).encode()
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def log_message(self, *a):
|
|
pass
|
|
|
|
|
|
class ServiceModeTests(Case):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.srv = ThreadingHTTPServer(("127.0.0.1", 0), _StubPii)
|
|
threading.Thread(target=cls.srv.serve_forever, daemon=True).start()
|
|
cls.addr = "127.0.0.1:%d" % cls.srv.server_address[1]
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
cls.srv.shutdown()
|
|
cls.srv.server_close()
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
_StubPii.seen, _StubPii.status, _StubPii.echo = [], 200, "honour"
|
|
|
|
def test_service_path(self):
|
|
env = {"PII_ADDR": self.addr, "PII_API_KEY": "k1"}
|
|
s, idx = self.run_stage(manifest(self.root, "sterilize", scope="conv9"), env=env)
|
|
self.assertEqual(s["sterilizer"], "service:http://" + self.addr)
|
|
self.assertNotIn("Marisol", self.text_of(idx["note.txt"]))
|
|
path, auth, req = _StubPii.seen[0]
|
|
self.assertEqual((path, auth, req["scope"], req["verify"]), ("/sterilize", "Bearer k1", "conv9", False))
|
|
self.assertIs(req["escape_held_literals"], True) # documents: escape every literal
|
|
# the echo is a handshake, not a count: it never reaches the index
|
|
self.assertNotIn("escape_held_literals", idx["note.txt"]["stats"])
|
|
with open(os.path.join(self.out, "sanitize.jsonl"), encoding="utf-8") as fh:
|
|
self.assertNotIn("escape_held_literals", fh.read())
|
|
self.assertNotIn("escape_held_literals", json.dumps(s))
|
|
self.assertEqual(idx["note.txt"]["stats"]["recall"], "llm")
|
|
# no local vault in service mode — the service owns it
|
|
self.assertFalse(os.path.exists(os.path.join(self.out, "sanitize", "vault.db")))
|
|
|
|
def test_service_502_fails_closed(self):
|
|
_StubPii.status = 502
|
|
env = {"PII_ADDR": self.addr, "PII_API_KEY": "k1"}
|
|
s, idx = self.run_stage(manifest(self.root, "sterilize"), env=env)
|
|
self.assertEqual(s["counts"]["not_run"], 3)
|
|
self.assertIn("502", s["stopped"])
|
|
|
|
def test_service_that_does_not_echo_the_literal_rule_stops_the_stage(self):
|
|
# an older service ignores escape_held_literals and answers 200 under
|
|
# the narrower rule; the stage must not store what it cannot trust
|
|
env = {"PII_ADDR": self.addr, "PII_API_KEY": "k1"}
|
|
for echo in ("strip", "false"):
|
|
with self.subTest(echo):
|
|
_StubPii.echo = echo
|
|
shutil.rmtree(self.out, ignore_errors=True)
|
|
os.makedirs(self.out)
|
|
s, idx = self.run_stage(manifest(self.root, "sterilize"), env=env)
|
|
self.assertIn("did not honour escape_held_literals", s["stopped"])
|
|
self.assertEqual(s["counts"]["ok"], 0)
|
|
self.assertEqual(s["counts"]["not_run"], 3)
|
|
self.assertFalse(os.path.exists(sz.store_path(self.out, idx["note.txt"]["sha256"])))
|
|
|
|
def test_service_without_key_refused(self):
|
|
self.prep(manifest(self.root, "sterilize"))
|
|
with self.assertRaises(sz.SanitizeError):
|
|
sz.sanitize(manifest(self.root, "sterilize"), self.out, env={"PII_ADDR": self.addr})
|
|
|
|
def test_flag_mode_never_uses_the_service(self):
|
|
env = {"PII_ADDR": self.addr, "PII_API_KEY": "k1"}
|
|
s, _ = self.run_stage(manifest(self.root, "flag"), env=env)
|
|
self.assertEqual(_StubPii.seen, [])
|
|
self.assertEqual(s["sterilizer"], "flag:rules-only")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|