405 lines
16 KiB
Python
405 lines
16 KiB
Python
"""bastion agent `sanitize` stage — the salt.
|
|
|
|
Runs the PII sterilizer (pii) over the extract store and writes a
|
|
SECOND store:
|
|
|
|
<out>/sanitize/<aa>/<sha256>.txt sanitized text — PII replaced by
|
|
stable [PII_TYPE_n] tokens
|
|
<out>/sanitize/vault.db token <-> value map, 0600, local
|
|
<out>/sanitize.jsonl one record per extracted file:
|
|
status, mode, stats (counts only —
|
|
never a value)
|
|
<out>/sanitize.summary.json
|
|
|
|
From here on, later stages read ONLY the sanitized store. The extract
|
|
store keeps the plaintext and stays on this machine; the vault is the
|
|
only way back (`/rehydrate` semantics, local).
|
|
|
|
Two ways to reach the sterilizer, chosen from the environment, never the
|
|
manifest (no addresses in manifests — docs/pii-sterilization.md):
|
|
|
|
in-process the default. the pii package runs inside this process: the rules
|
|
pass, the recall backend selected by PII_RECALL
|
|
(llm | gliner — the laptop path), the vault under <out>.
|
|
service PII_ADDR (+ PII_API_KEY) set: POST /sterilize to the
|
|
running PII service, which owns its own vault. The
|
|
server path; the bastion advertises the capability.
|
|
|
|
Modes (manifest `sanitize.mode`):
|
|
sterilize rules + recall pass (+ optional adversarial verify).
|
|
flag rules pass ONLY — deterministic, model-free, always
|
|
in-process. Redacts the precision anchors (emails, cards,
|
|
keys, IBANs, …); contextual PII (names, schools) is NOT
|
|
caught. The manifest doc says so; this docstring says so.
|
|
|
|
Fail closed: the first recall failure (model unreachable, package
|
|
missing) stops the stage — remaining files are recorded `not_run` and
|
|
nothing half-sanitized is written. Long texts are split on paragraph
|
|
boundaries per call (SANITIZE_CHUNK_CHARS); tokens are stable per scope,
|
|
so a value seen in two chunks gets one token.
|
|
"""
|
|
|
|
import datetime as _dt
|
|
import json
|
|
import os
|
|
import re
|
|
import tempfile
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
from .extract import store_path as extract_store_path
|
|
from .manifest import signing_digest
|
|
|
|
# Part of the reuse key: bump it whenever the pipeline would sterilize the
|
|
# same text differently, or a stored result from older code is kept as is.
|
|
SANITIZE_VERSION = 3
|
|
SANITIZE_CHUNK_CHARS = 20000 # per sterilize call; under the service's 2 MiB body cap
|
|
SERVICE_TIMEOUT = 600.0
|
|
|
|
STATUS_OK, STATUS_ERROR, STATUS_SKIPPED, STATUS_NOT_RUN = "ok", "error", "skipped", "not_run"
|
|
|
|
|
|
class SanitizeError(Exception):
|
|
"""Fail-closed stop: the recall pass is unavailable."""
|
|
|
|
|
|
def pii_modules():
|
|
"""The sterilizer, a SIBLING top-level package here and inside the built
|
|
.pyz alike (build.py writes it as `pii/`). One import path: an ImportError
|
|
here means the sterilizer is genuinely absent, which is a fail-closed stop,
|
|
never a silent skip.
|
|
|
|
pii.pipeline.sterilize(text, scope, vault, llm_extract, llm_verify=None)
|
|
returning (text, result), pii.recall.build_backend (which raises
|
|
RecallError rather than degrading) and pii.vault.Vault."""
|
|
from pii import pipeline, recall
|
|
from pii.vault import Vault
|
|
return pipeline, recall, Vault
|
|
|
|
|
|
# Test seam: when set, used instead of recall.build_backend(env).
|
|
BUILD_BACKEND = None
|
|
|
|
|
|
# ── sterilizer front-ends ───────────────────────────────────────────────────
|
|
|
|
class InProcess:
|
|
"""pii inside this process. Vault lives under the output dir."""
|
|
|
|
name = "in-process"
|
|
|
|
def __init__(self, out_dir, mode, env, verify=False):
|
|
pipeline, recall, Vault = pii_modules()
|
|
self._pipeline = pipeline
|
|
self._recall_error = recall.RecallError
|
|
self.mode = mode
|
|
self.verify = verify
|
|
os.umask(0o077)
|
|
# FREEAGENT_VAULT_DIR keeps the vault (plaintext PII values) OUT of the
|
|
# output directory when that directory sits inside a folder that may
|
|
# be cloud-synced (<root>/.agent). The bastion sets it to its own
|
|
# state dir; the default stays beside the output for a hand-run agent.
|
|
vault_dir = (env or {}).get("FREEAGENT_VAULT_DIR", "").strip()
|
|
vault_path = (os.path.join(vault_dir, "vault.db") if vault_dir
|
|
else os.path.join(out_dir, "sanitize", "vault.db"))
|
|
self.vault = Vault(vault_path)
|
|
try:
|
|
os.chmod(self.vault.path, 0o600)
|
|
except OSError:
|
|
pass
|
|
if mode == "flag":
|
|
self.backend = None
|
|
self.describe = "flag:rules-only"
|
|
else:
|
|
build = BUILD_BACKEND or recall.build_backend
|
|
try:
|
|
self.backend = build(env)
|
|
except recall.RecallError as e:
|
|
raise SanitizeError(str(e))
|
|
self.describe = "sterilize:" + self.backend.describe()
|
|
|
|
def sterilize(self, text, scope):
|
|
if self.backend is None:
|
|
extract_fn, verify_fn = (lambda t: []), None
|
|
else:
|
|
extract_fn = self.backend.extract
|
|
verify_fn = self.backend.verify if self.verify else None
|
|
try:
|
|
# documents are raw files, never this pipeline's output: every
|
|
# placeholder-shaped literal in one is escaped
|
|
return self._pipeline.sterilize(text, scope, self.vault, extract_fn, verify_fn,
|
|
escape_held_literals=True)
|
|
except self._recall_error as e:
|
|
raise SanitizeError("recall pass failed: %s" % e)
|
|
|
|
|
|
class Service:
|
|
"""The running PII service over HTTP. It owns the vault."""
|
|
|
|
name = "service"
|
|
|
|
def __init__(self, addr, api_key, verify=False, timeout=SERVICE_TIMEOUT):
|
|
if not api_key:
|
|
raise SanitizeError("PII_ADDR set but PII_API_KEY empty — the service is fail-closed")
|
|
self.url = ("http://" + addr if "://" not in addr else addr).rstrip("/")
|
|
self.api_key = api_key
|
|
self.verify = verify
|
|
self.timeout = timeout
|
|
self.describe = "service:" + self.url
|
|
|
|
def sterilize(self, text, scope):
|
|
body = json.dumps({"text": text, "scope": scope, "verify": self.verify,
|
|
"escape_held_literals": True}).encode("utf-8")
|
|
req = urllib.request.Request(
|
|
self.url + "/sterilize", data=body,
|
|
headers={"Content-Type": "application/json",
|
|
"Authorization": "Bearer " + self.api_key})
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=self.timeout) as resp:
|
|
payload = json.load(resp)
|
|
except urllib.error.HTTPError as e:
|
|
if e.code == 502:
|
|
raise SanitizeError("pii service: recall pass unavailable (502)")
|
|
raise RuntimeError("pii service HTTP %d" % e.code)
|
|
except (OSError, ValueError) as e:
|
|
raise SanitizeError("pii service unreachable: %s" % e)
|
|
stats = dict(payload.get("stats", {}))
|
|
# a service that predates the field ignores it and applies the narrower
|
|
# rule silently; without its echo the document cannot be trusted
|
|
if stats.pop("escape_held_literals", None) is not True:
|
|
raise SanitizeError("the pii service did not honour escape_held_literals "
|
|
"(no echo in its stats); upgrade the service")
|
|
return payload["text"], {"findings": payload.get("findings", []), "stats": stats}
|
|
|
|
|
|
def build_front_end(out_dir, sanitize_cfg, env):
|
|
if sanitize_cfg.mode == "flag":
|
|
return InProcess(out_dir, "flag", env)
|
|
addr = env.get("PII_ADDR", "").strip()
|
|
if addr:
|
|
return Service(addr, env.get("PII_API_KEY", ""), verify=sanitize_cfg.verify)
|
|
return InProcess(out_dir, "sterilize", env, verify=sanitize_cfg.verify)
|
|
|
|
|
|
# ── helpers ─────────────────────────────────────────────────────────────────
|
|
|
|
# Same shape as pii.pipeline.TOKEN_RE (a test keeps them equal). Kept here so
|
|
# splitting needs no pii import; the sterilizer is loaded fail-closed elsewhere.
|
|
_PLACEHOLDER_RE = re.compile(r"\[PII_([A-Z_]+)_(\d+)\]")
|
|
|
|
|
|
def _clear_of_placeholder(piece, limit):
|
|
"""A hard cut at limit, moved back to before the '[' of a placeholder it
|
|
would split. Each half of a split placeholder would reach the sterilizer
|
|
as unrecognised text and be stored unescaped. A placeholder at the very
|
|
start is kept whole instead, so the chunk may run over the limit by the
|
|
rest of the placeholder.
|
|
|
|
A placeholder has no '[' after its first character, so the only one that
|
|
can straddle the cut starts at the last '[' before it: exact at any
|
|
placeholder length, with no search window."""
|
|
start = piece.rfind("[", 0, limit)
|
|
if start >= 0:
|
|
m = _PLACEHOLDER_RE.match(piece, start)
|
|
if m and m.end() > limit:
|
|
return start if start > 0 else m.end()
|
|
return limit
|
|
|
|
|
|
def split_chunks(text, limit=None):
|
|
"""Split on paragraph, then line, boundaries; a single oversized line
|
|
is cut hard. Concatenation of the chunks is byte-identical to text.
|
|
The default limit is read when called, so a changed
|
|
SANITIZE_CHUNK_CHARS takes effect."""
|
|
if limit is None:
|
|
limit = SANITIZE_CHUNK_CHARS
|
|
if not text:
|
|
return []
|
|
if len(text) <= limit:
|
|
return [text]
|
|
parts = text.split("\n\n")
|
|
paras = [p + ("\n\n" if i < len(parts) - 1 else "") for i, p in enumerate(parts)]
|
|
out, buf = [], ""
|
|
for piece in paras:
|
|
if len(buf) + len(piece) <= limit:
|
|
buf += piece
|
|
continue
|
|
if buf:
|
|
out.append(buf)
|
|
buf = ""
|
|
while len(piece) > limit:
|
|
cut = piece.rfind("\n", 0, limit)
|
|
cut = _clear_of_placeholder(piece, limit) if cut <= 0 else cut + 1
|
|
out.append(piece[:cut])
|
|
piece = piece[cut:]
|
|
buf = piece
|
|
if buf:
|
|
out.append(buf)
|
|
return [c for c in out if c]
|
|
|
|
|
|
def _atomic_write_text(path, text):
|
|
d = os.path.dirname(path)
|
|
os.makedirs(d, exist_ok=True)
|
|
fd, tmp = tempfile.mkstemp(prefix=".tmp-", dir=d)
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8", newline="\n") as fh:
|
|
fh.write(text)
|
|
os.replace(tmp, path)
|
|
except BaseException:
|
|
try:
|
|
os.unlink(tmp)
|
|
except OSError:
|
|
pass
|
|
raise
|
|
|
|
|
|
def store_path(out_dir, sha):
|
|
return os.path.join(out_dir, "sanitize", sha[:2], sha + ".txt")
|
|
|
|
|
|
def load_extract_index(out_dir):
|
|
path = os.path.join(out_dir, "extract.jsonl")
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError("no extract index at %s — run the extract stage first" % path)
|
|
with open(path, encoding="utf-8") as fh:
|
|
return [json.loads(l) for l in fh]
|
|
|
|
|
|
def load_index(out_dir):
|
|
idx = {}
|
|
try:
|
|
with open(os.path.join(out_dir, "sanitize.jsonl"), encoding="utf-8") as fh:
|
|
for line in fh:
|
|
try:
|
|
rec = json.loads(line)
|
|
except ValueError:
|
|
continue
|
|
if rec.get("sha256"):
|
|
idx[rec["sha256"]] = rec
|
|
except OSError:
|
|
pass
|
|
return idx
|
|
|
|
|
|
def _cache_key(front, cfg):
|
|
return "%s|%s|%s|%s|v%d" % (front.name, cfg.mode, cfg.scope, cfg.verify, SANITIZE_VERSION)
|
|
|
|
|
|
def _merge_stats(acc, stats):
|
|
for k, v in stats.items():
|
|
if isinstance(v, int):
|
|
acc[k] = acc.get(k, 0) + v
|
|
else:
|
|
acc[k] = v
|
|
|
|
|
|
# ── the stage ───────────────────────────────────────────────────────────────
|
|
|
|
def sanitize(manifest, out_dir, env=None, now=None, front_end=None):
|
|
cfg = manifest.sanitize
|
|
if cfg is None:
|
|
raise SanitizeError("manifest has no sanitize section / stage")
|
|
env = dict(os.environ) if env is None else env
|
|
front = front_end or build_front_end(out_dir, cfg, env)
|
|
key = _cache_key(front, cfg)
|
|
extracted = load_extract_index(out_dir)
|
|
prev = load_index(out_dir)
|
|
counts = {k: 0 for k in ("files", "ok", "reused", "skipped", "error", "not_run")}
|
|
type_counts = {}
|
|
index = []
|
|
done = {}
|
|
stopped = None
|
|
|
|
for rec in extracted:
|
|
counts["files"] += 1
|
|
sha = rec.get("sha256")
|
|
base = {"root": rec["root"], "path": rec["path"], "sha256": sha,
|
|
"mode": cfg.mode, "v": SANITIZE_VERSION}
|
|
if rec.get("status") not in ("ok",):
|
|
counts["skipped"] += 1
|
|
index.append(dict(base, status=STATUS_SKIPPED, error="extract: " + str(rec.get("status"))))
|
|
continue
|
|
if sha in done:
|
|
dup = done[sha]
|
|
index.append(dict(dup, root=rec["root"], path=rec["path"]))
|
|
counts["reused" if dup.get("reused") else dup["status"]] += 1
|
|
continue
|
|
if stopped is not None:
|
|
out = dict(base, status=STATUS_NOT_RUN, error=stopped)
|
|
counts["not_run"] += 1
|
|
index.append(out)
|
|
done[sha] = out
|
|
continue
|
|
spath = store_path(out_dir, sha)
|
|
old = prev.get(sha)
|
|
if old and old.get("status") == STATUS_OK and old.get("key") == key and os.path.exists(spath):
|
|
out = dict(base, status=STATUS_OK, key=key, stats=old.get("stats", {}),
|
|
types=old.get("types", {}), chars=old.get("chars", 0), reused=True)
|
|
counts["reused"] += 1
|
|
index.append(out)
|
|
done[sha] = out
|
|
continue
|
|
try:
|
|
with open(extract_store_path(out_dir, sha), encoding="utf-8") as fh:
|
|
text = fh.read()
|
|
except OSError as e:
|
|
out = dict(base, status=STATUS_ERROR, error="extract store: %s" % e)
|
|
counts["error"] += 1
|
|
index.append(out)
|
|
done[sha] = out
|
|
continue
|
|
pieces, stats, types = [], {}, {}
|
|
try:
|
|
for chunk in split_chunks(text):
|
|
clean, res = front.sterilize(chunk, cfg.scope)
|
|
pieces.append(clean)
|
|
_merge_stats(stats, res.get("stats", {}))
|
|
for f in res.get("findings", []):
|
|
t = f.get("type", "?")
|
|
types[t] = types.get(t, 0) + 1
|
|
except SanitizeError as e:
|
|
stopped = str(e)
|
|
out = dict(base, status=STATUS_NOT_RUN, error=stopped)
|
|
counts["not_run"] += 1
|
|
index.append(out)
|
|
done[sha] = out
|
|
continue
|
|
except (RuntimeError, ValueError, KeyError) as e:
|
|
out = dict(base, status=STATUS_ERROR, error="%s: %s" % (type(e).__name__, e))
|
|
counts["error"] += 1
|
|
index.append(out)
|
|
done[sha] = out
|
|
continue
|
|
clean = "".join(pieces)
|
|
_atomic_write_text(spath, clean)
|
|
out = dict(base, status=STATUS_OK, key=key, stats=stats, types=types, chars=len(clean))
|
|
counts["ok"] += 1
|
|
for t, n in types.items():
|
|
type_counts[t] = type_counts.get(t, 0) + n
|
|
index.append(out)
|
|
done[sha] = out
|
|
|
|
index.sort(key=lambda r: (r["root"], r["path"]))
|
|
_atomic_write_text(os.path.join(out_dir, "sanitize.jsonl"), "".join(
|
|
json.dumps(r, ensure_ascii=False, sort_keys=True) + "\n" for r in index))
|
|
summary = {
|
|
"freeagent": 1,
|
|
"stage": "sanitize",
|
|
"sanitize_version": SANITIZE_VERSION,
|
|
"manifest_id": manifest.id,
|
|
"manifest_digest": "sha256:" + signing_digest(manifest.raw),
|
|
"sanitized_at": now or _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
|
|
"mode": cfg.mode,
|
|
"scope": cfg.scope,
|
|
"verify": cfg.verify,
|
|
"sterilizer": front.describe,
|
|
"counts": counts,
|
|
"types": type_counts,
|
|
"stopped": stopped,
|
|
"entries": len(index),
|
|
}
|
|
_atomic_write_text(os.path.join(out_dir, "sanitize.summary.json"),
|
|
json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True) + "\n")
|
|
return summary
|