bastion agent 0.1.14
This commit is contained in:
@@ -0,0 +1,404 @@
|
||||
"""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 = 2
|
||||
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
|
||||
Reference in New Issue
Block a user