bastion agent 0.1.14
This commit is contained in:
@@ -0,0 +1 @@
|
||||
0.1.14
|
||||
@@ -0,0 +1,12 @@
|
||||
from .manifest import ( # noqa: F401
|
||||
STAGES,
|
||||
MASK_OPS,
|
||||
SWEEPS,
|
||||
GRANT_OPS,
|
||||
Manifest,
|
||||
ManifestError,
|
||||
load_manifest,
|
||||
signing_digest,
|
||||
signing_payload,
|
||||
validate_document,
|
||||
)
|
||||
@@ -0,0 +1,367 @@
|
||||
"""CLI: python3 -m bastion_agent {validate,show,payload,walk} <manifest.yaml> [...]"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
from . import doctor as _doctor
|
||||
from .catalog import CatalogError, catalog
|
||||
from .embed import EmbedError, chunk, embed, search
|
||||
from .extract import extract
|
||||
from .manifest import ManifestError, load_manifest, signing_digest, signing_payload
|
||||
from .run import GateError, RunError, default_out, sign_manifest_file, sweep, write_signature
|
||||
from .sanitize import SanitizeError, sanitize
|
||||
from .seal import SealError, seal, unseal
|
||||
from .serve import ServeError, serve
|
||||
from .walk import walk
|
||||
|
||||
|
||||
def _show(m, path):
|
||||
print(f"{path}: {m.id} (owner {m.owner_host})")
|
||||
print(f" roots: {', '.join(m.scope.roots)}")
|
||||
if m.scope.exclude:
|
||||
print(f" exclude: {', '.join(m.scope.exclude)}")
|
||||
print(f" mask: {' '.join(m.mask)}")
|
||||
print(f" stages: {' -> '.join(m.stages)}")
|
||||
if m.sanitize:
|
||||
print(f" sanitize: {m.sanitize.mode} (scope {m.sanitize.scope}"
|
||||
f"{', verify' if m.sanitize.verify else ''})")
|
||||
if m.embed:
|
||||
print(f" embed: {m.embed.model} {m.embed.dim}d, chunk "
|
||||
f"{m.embed.chunk_size_tokens}/{m.embed.chunk_overlap_tokens}")
|
||||
if m.summarize:
|
||||
print(f" summarize: {m.summarize.model} per {', '.join(m.summarize.per)}")
|
||||
print(f" sweep: {m.sweep}")
|
||||
print(f" sink: catalog={m.sink.catalog} vectors={m.sink.vectors} "
|
||||
f"plaintext={m.sink.plaintext}"
|
||||
f"{' seal_to=' + ','.join(m.sink.seal_to) if m.sink.seal_to else ''}")
|
||||
if m.grants:
|
||||
for g in m.grants:
|
||||
print(f" grant: {g.host} may {', '.join(g.ops)} until {g.until.isoformat()}")
|
||||
else:
|
||||
print(" grants: none")
|
||||
print(f" signed: {'yes (' + m.signature.alg + ')' if m.signed else 'NO — runtime will refuse'}")
|
||||
print(f" digest: sha256:{signing_digest(m.raw)}")
|
||||
|
||||
|
||||
def _version():
|
||||
# pkgutil reads from a zipapp as well as a checkout; a plain open() cannot
|
||||
import pkgutil
|
||||
try:
|
||||
data = pkgutil.get_data(__package__ or "bastion_agent", "VERSION")
|
||||
return data.decode("utf-8").strip() if data else "0.0.0"
|
||||
except (OSError, ValueError):
|
||||
return "0.0.0"
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(prog="bastion-agent")
|
||||
parser.add_argument("--version", action="version", version="bastion-agent %s" % _version())
|
||||
sub = parser.add_subparsers(dest="cmd", required=True)
|
||||
p_doc = sub.add_parser("doctor", help="preflight: python, tools, endpoints, environment")
|
||||
p_doc.add_argument("--json", action="store_true")
|
||||
p_val = sub.add_parser("validate", help="validate manifest(s), exit 1 on any error")
|
||||
p_val.add_argument("--require-signed", action="store_true",
|
||||
help="also fail when the signature block is absent")
|
||||
p_val.add_argument("paths", nargs="+")
|
||||
p_show = sub.add_parser("show", help="print a manifest's effective plan")
|
||||
p_show.add_argument("paths", nargs="+")
|
||||
p_pay = sub.add_parser("payload", help="emit the canonical signing payload (bytes) to stdout")
|
||||
p_pay.add_argument("path")
|
||||
p_walk = sub.add_parser("walk", help="run the walk stage: filesystem + metadata map")
|
||||
p_walk.add_argument("path")
|
||||
p_walk.add_argument("--out", help="output directory (default <root>/.agent, or build/freeagent/<id> for several roots)")
|
||||
p_walk.add_argument("--no-hash", action="store_true",
|
||||
help="metadata only; skip SHA-256 (fast survey)")
|
||||
p_walk.add_argument("--require-signed", action="store_true",
|
||||
help="refuse an unsigned manifest, as the runtime will")
|
||||
p_ext = sub.add_parser("extract", help="run the extract stage: text out of walked files")
|
||||
p_ext.add_argument("path")
|
||||
p_ext.add_argument("--out", help="output directory (default <root>/.agent, or build/freeagent/<id> for several roots)")
|
||||
p_ext.add_argument("--no-external", action="store_true",
|
||||
help="built-in extractors only; ignore pdftotext/tesseract on PATH")
|
||||
p_ext.add_argument("--require-signed", action="store_true")
|
||||
p_san = sub.add_parser("sanitize", help="run the sanitize stage: PII out of the extracted text")
|
||||
p_san.add_argument("path")
|
||||
p_san.add_argument("--out", help="output directory (default <root>/.agent, or build/freeagent/<id> for several roots)")
|
||||
p_san.add_argument("--require-signed", action="store_true")
|
||||
for name, helptext in (("chunk", "run the chunk stage: sanitized text -> overlapping windows"),
|
||||
("embed", "run the embed stage: chunks -> local vectors.db")):
|
||||
p = sub.add_parser(name, help=helptext)
|
||||
p.add_argument("path")
|
||||
p.add_argument("--out")
|
||||
p.add_argument("--require-signed", action="store_true")
|
||||
p_srch = sub.add_parser("search", help="query the local vectors (proves the lane end to end)")
|
||||
p_srch.add_argument("path")
|
||||
p_srch.add_argument("query")
|
||||
p_srch.add_argument("--out")
|
||||
p_srch.add_argument("-k", type=int, default=5)
|
||||
p_cat = sub.add_parser("catalog", help="run the catalog stage: file + directory cards")
|
||||
p_cat.add_argument("path")
|
||||
p_cat.add_argument("--out")
|
||||
p_cat.add_argument("--require-signed", action="store_true")
|
||||
p_seal = sub.add_parser("seal", help="run the seal stage: catalog -> age envelope for the owner")
|
||||
p_seal.add_argument("path")
|
||||
p_seal.add_argument("--out")
|
||||
p_seal.add_argument("--require-signed", action="store_true")
|
||||
p_unseal = sub.add_parser("unseal", help="owner-side: open a sealed bundle with an age identity")
|
||||
p_unseal.add_argument("sealed")
|
||||
p_unseal.add_argument("--identity", required=True, help="age identity file (age-keygen output)")
|
||||
p_unseal.add_argument("--dest", required=True)
|
||||
p_sw = sub.add_parser("sweep", help="the runtime: gate the manifest, run its stages on its cadence")
|
||||
p_sw.add_argument("path")
|
||||
p_sw.add_argument("--out", help="output directory (default <root>/.agent, or build/freeagent/<id> "
|
||||
"for several roots; the same folder serve uses)")
|
||||
p_sw.add_argument("--once", action="store_true", help="one run, then exit (any cadence)")
|
||||
p_sw.add_argument("--allow-unsigned", action="store_true",
|
||||
help="DEVELOPMENT ONLY: run an unsigned or unverifiable manifest; recorded in run.jsonl")
|
||||
p_sign = sub.add_parser("sign", help="write the signature block into a manifest (host-side helper)")
|
||||
p_sign.add_argument("path")
|
||||
p_sign_by = p_sign.add_mutually_exclusive_group(required=True)
|
||||
p_sign_by.add_argument("--key", help="PEM private key (signs with openssl)")
|
||||
p_sign_by.add_argument("--value", help="a base64 DER signature over `payload`, made "
|
||||
"elsewhere (the bastion signs with its own key)")
|
||||
p_sign.add_argument("--cert", required=True, help="cert path to record in the manifest")
|
||||
p_srv = sub.add_parser("serve", help="the query endpoint: search/read_range over mTLS for owner + grantees")
|
||||
p_srv.add_argument("path")
|
||||
p_srv.add_argument("--out")
|
||||
p_srv.add_argument("--bind", help="host:port or [v6]:port (default FREEAGENT_BIND or [::]:8470)")
|
||||
p_srv.add_argument("--allow-unsigned", action="store_true",
|
||||
help="DEVELOPMENT ONLY: serve an unsigned/unverifiable manifest")
|
||||
p_srv.add_argument("--port-file", help="write the bound port here (for a supervisor binding port 0)")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.cmd == "doctor":
|
||||
return _doctor.main(["--json"] if args.json else [])
|
||||
|
||||
if args.cmd == "serve":
|
||||
try:
|
||||
m = load_manifest(args.path)
|
||||
except ManifestError as exc:
|
||||
print(exc, file=sys.stderr)
|
||||
return 1
|
||||
out = args.out or default_out(m)
|
||||
try:
|
||||
serve(args.path, out, allow_unsigned=args.allow_unsigned, bind=args.bind,
|
||||
port_file=args.port_file)
|
||||
except ServeError as exc:
|
||||
print("serve: %s" % exc, file=sys.stderr)
|
||||
return 1
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
return 0
|
||||
|
||||
if args.cmd == "sign":
|
||||
try:
|
||||
if args.value:
|
||||
import base64
|
||||
try:
|
||||
der = base64.b64decode(args.value, validate=True)
|
||||
except Exception:
|
||||
print("sign: --value is not base64", file=sys.stderr)
|
||||
return 1
|
||||
write_signature(args.path, der, args.cert)
|
||||
else:
|
||||
sign_manifest_file(args.path, args.key, args.cert)
|
||||
except (GateError, ManifestError) as exc:
|
||||
print("sign: %s" % exc, file=sys.stderr)
|
||||
return 1
|
||||
print("%s: signed (cert %s)" % (args.path, args.cert))
|
||||
return 0
|
||||
|
||||
if args.cmd == "sweep":
|
||||
try:
|
||||
last = sweep(args.path, args.out, allow_unsigned=args.allow_unsigned, once=args.once)
|
||||
except RunError as exc:
|
||||
print("sweep: %s" % exc, file=sys.stderr)
|
||||
return 1
|
||||
return 0 if (last and last.get("ok")) else 1
|
||||
|
||||
if args.cmd == "unseal":
|
||||
try:
|
||||
inv = unseal(args.sealed, args.identity, args.dest)
|
||||
except SealError as exc:
|
||||
print("unseal: %s" % exc, file=sys.stderr)
|
||||
return 1
|
||||
print("%s: %d members verified -> %s (manifest %s, sealed %s)"
|
||||
% (args.sealed, len(inv["members"]), args.dest, inv["manifest_id"], inv["sealed_at"]))
|
||||
return 0
|
||||
|
||||
if args.cmd == "seal":
|
||||
try:
|
||||
m = load_manifest(args.path, require_signed=args.require_signed)
|
||||
except ManifestError as exc:
|
||||
print(exc, file=sys.stderr)
|
||||
return 1
|
||||
out = args.out or default_out(m)
|
||||
try:
|
||||
s = seal(m, out)
|
||||
except SealError as exc:
|
||||
print("seal: %s" % exc, file=sys.stderr)
|
||||
return 1
|
||||
print("%s: sealed %d members (%d bytes -> %d) to %d recipient(s)%s -> %s"
|
||||
% (m.id, len(s["members"]), s["bytes_plain"], s["bytes_sealed"],
|
||||
len(s["recipients"]), " + signature" if s["signature"] else "",
|
||||
os.path.join(out, s["file"])))
|
||||
return 0
|
||||
|
||||
if args.cmd == "catalog":
|
||||
try:
|
||||
m = load_manifest(args.path, require_signed=args.require_signed)
|
||||
except ManifestError as exc:
|
||||
print(exc, file=sys.stderr)
|
||||
return 1
|
||||
if "catalog" not in m.stages:
|
||||
print("%s: 'catalog' is not in this manifest's stages" % args.path, file=sys.stderr)
|
||||
return 1
|
||||
out = args.out or default_out(m)
|
||||
try:
|
||||
s = catalog(m, out)
|
||||
except (FileNotFoundError, CatalogError) as exc:
|
||||
print("catalog: %s" % exc, file=sys.stderr)
|
||||
return 1
|
||||
c = s["counts"]
|
||||
print("%s: %d file cards (%d with text), %d directory cards, %d names masked; "
|
||||
"summaries %d new, %d reused [%s]; catalog %d bytes -> %s"
|
||||
% (m.id, c["files"], c["with_text"], c["dirs"], c["names_masked"],
|
||||
c["summaries"], c["summaries_reused"], s["summarizer"] or "none",
|
||||
s["bytes_catalog"], out))
|
||||
if s["stopped"]:
|
||||
print("catalog: summarizer STOPPED: %s" % s["stopped"], file=sys.stderr)
|
||||
return 1
|
||||
return 0
|
||||
|
||||
if args.cmd in ("chunk", "embed", "search"):
|
||||
try:
|
||||
m = load_manifest(args.path, require_signed=getattr(args, "require_signed", False))
|
||||
except ManifestError as exc:
|
||||
print(exc, file=sys.stderr)
|
||||
return 1
|
||||
stage = "embed" if args.cmd == "search" else args.cmd
|
||||
if stage not in m.stages:
|
||||
print("%s: %r is not in this manifest's stages" % (args.path, stage), file=sys.stderr)
|
||||
return 1
|
||||
out = args.out or default_out(m)
|
||||
try:
|
||||
if args.cmd == "chunk":
|
||||
s = chunk(m, out)
|
||||
c = s["counts"]
|
||||
print("%s: %d docs -> %d chunks (%d/%d tokens), %d skipped -> %s"
|
||||
% (m.id, c["docs"], c["chunks"], s["size_tokens"], s["overlap_tokens"],
|
||||
c["skipped"], out))
|
||||
return 0
|
||||
if args.cmd == "embed":
|
||||
s = embed(m, out)
|
||||
c = s["counts"]
|
||||
print("%s: %d chunks in %d docs -> %d embedded, %d reused [%s, %dd]; "
|
||||
"%d stale vectors kept; store %s"
|
||||
% (m.id, c["chunks"], c["docs"], c["embedded"], c["reused"], s["endpoint"],
|
||||
s["dim"], s["stale_vectors"], s["store"]))
|
||||
return 0
|
||||
hits = search(m, out, args.query, k=args.k)
|
||||
for h in hits:
|
||||
snippet = h["text"].replace("\n", " ")[:160]
|
||||
print("%.4f %s#%d %s" % (h["score"], h["doc"][:12], h["idx"], snippet))
|
||||
if not hits:
|
||||
print("no vectors for model %s" % m.embed.model, file=sys.stderr)
|
||||
return 0
|
||||
except (FileNotFoundError, EmbedError) as exc:
|
||||
print("%s: %s" % (args.cmd, exc), file=sys.stderr)
|
||||
return 1
|
||||
|
||||
if args.cmd == "sanitize":
|
||||
try:
|
||||
m = load_manifest(args.path, require_signed=args.require_signed)
|
||||
except ManifestError as exc:
|
||||
print(exc, file=sys.stderr)
|
||||
return 1
|
||||
if "sanitize" not in m.stages:
|
||||
print("%s: 'sanitize' is not in this manifest's stages" % args.path, file=sys.stderr)
|
||||
return 1
|
||||
out = args.out or default_out(m)
|
||||
try:
|
||||
s = sanitize(m, out)
|
||||
except (FileNotFoundError, SanitizeError) as exc:
|
||||
print("sanitize: %s" % exc, file=sys.stderr)
|
||||
return 1
|
||||
c = s["counts"]
|
||||
types = ", ".join("%s=%d" % kv for kv in sorted(s["types"].items())) or "none"
|
||||
print("%s: %d files -> %d sanitized, %d reused, %d skipped, %d errors, %d not run "
|
||||
"[%s, scope %s]; tokens: %s -> %s"
|
||||
% (m.id, c["files"], c["ok"], c["reused"], c["skipped"], c["error"], c["not_run"],
|
||||
s["sterilizer"], s["scope"], types, out))
|
||||
if s["stopped"]:
|
||||
print("sanitize STOPPED (fail closed): %s" % s["stopped"], file=sys.stderr)
|
||||
return 1
|
||||
return 2 if c["error"] else 0
|
||||
|
||||
if args.cmd == "extract":
|
||||
try:
|
||||
m = load_manifest(args.path, require_signed=args.require_signed)
|
||||
except ManifestError as exc:
|
||||
print(exc, file=sys.stderr)
|
||||
return 1
|
||||
if "extract" not in m.stages:
|
||||
print("%s: 'extract' is not in this manifest's stages" % args.path, file=sys.stderr)
|
||||
return 1
|
||||
out = args.out or default_out(m)
|
||||
try:
|
||||
s = extract(m, out, use_external=not args.no_external)
|
||||
except FileNotFoundError as exc:
|
||||
print(exc, file=sys.stderr)
|
||||
return 1
|
||||
c = s["counts"]
|
||||
tools = ", ".join(k for k, v in s["tools"].items() if v) or "none"
|
||||
print("%s: %d files -> %d extracted (%d chars), %d reused, %d empty, %d unsupported, "
|
||||
"%d skipped, %d errors; tools: %s -> %s"
|
||||
% (m.id, c["files"], c["ok"], s["chars"], c["reused"], c["empty"],
|
||||
c["unsupported"], c["skipped"], c["error"], tools, out))
|
||||
return 2 if c["error"] else 0
|
||||
|
||||
if args.cmd == "walk":
|
||||
try:
|
||||
m = load_manifest(args.path, require_signed=args.require_signed)
|
||||
except ManifestError as exc:
|
||||
print(exc, file=sys.stderr)
|
||||
return 1
|
||||
out = args.out or default_out(m)
|
||||
s = walk(m, out, do_hash=not args.no_hash)
|
||||
c = s["counts"]
|
||||
print("%s: walked %d entries, %d files (%d bytes), %d dirs, %d links; "
|
||||
"%d hashed, %d reused, %d too large, %d excluded, %d errors -> %s"
|
||||
% (m.id, s["entries"], c["files"], s["bytes"], c["dirs"], c["symlinks"],
|
||||
c["hashed"], c["hash_reused"], c["too_large"], c["excluded"],
|
||||
c["errors"], out))
|
||||
for r in s["roots"]:
|
||||
if r.get("error"):
|
||||
print(" root %d %s: %s" % (r["index"], r["path"], r["error"]), file=sys.stderr)
|
||||
return 2 if c["errors"] else 0
|
||||
|
||||
if args.cmd == "payload":
|
||||
try:
|
||||
m = load_manifest(args.path)
|
||||
except ManifestError as exc:
|
||||
print(exc, file=sys.stderr)
|
||||
return 1
|
||||
sys.stdout.buffer.write(signing_payload(m.raw))
|
||||
return 0
|
||||
|
||||
rc = 0
|
||||
for path in args.paths:
|
||||
try:
|
||||
m = load_manifest(path, require_signed=getattr(args, "require_signed", False))
|
||||
except ManifestError as exc:
|
||||
print(f"{path}:", file=sys.stderr)
|
||||
print(exc, file=sys.stderr)
|
||||
rc = 1
|
||||
continue
|
||||
if args.cmd == "validate":
|
||||
print(f"{path}: OK — {m.id}, {len(m.stages)} stages, {len(m.grants)} grants, "
|
||||
f"{'signed' if m.signed else 'unsigned'}")
|
||||
else:
|
||||
_show(m, path)
|
||||
return rc
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,34 @@
|
||||
# bastion agent sweep — systemd USER unit (Linux). Runs the manifest's stages on
|
||||
# its cadence under the owner's account; the corpus never leaves that user's
|
||||
# permissions.
|
||||
#
|
||||
# mkdir -p ~/.config/systemd/user ~/.config/bastion-agent
|
||||
# cp bastion-agent/bastion-agent.service ~/.config/systemd/user/
|
||||
# # environment (never in the manifest): recipients, endpoints, tools
|
||||
# printf 'FREEAGENT_OWNER_RECIPIENT=age1...\nEMBED_URL=http://127.0.0.1:11434\n' > ~/.config/bastion-agent/env
|
||||
# chmod 600 ~/.config/bastion-agent/env
|
||||
# systemctl --user daemon-reload && systemctl --user enable --now bastion-agent
|
||||
# journalctl --user -u bastion-agent -f
|
||||
#
|
||||
# Mac: a launchd agent with the same ExecStart (KeepAlive true); Windows: a
|
||||
# Scheduled Task "At log on" running the same command. Neither file is in
|
||||
# this repository.
|
||||
[Unit]
|
||||
Description=bastion agent sweep (%h/.config/bastion-agent/manifest.yaml)
|
||||
After=network-online.target
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
WorkingDirectory=%h/.local/share/bastion-agent
|
||||
EnvironmentFile=%h/.config/bastion-agent/env
|
||||
# the repo checkout (or the packaged binary later); adjust to where it lives
|
||||
ExecStart=/usr/bin/python3 -m bastion_agent sweep %h/.config/bastion-agent/manifest.yaml --out %h/.local/share/bastion-agent/out
|
||||
Restart=on-failure
|
||||
RestartSec=60
|
||||
# the vault, stores and envelope are 0600 already; keep the unit from widening them
|
||||
UMask=0077
|
||||
NoNewPrivileges=true
|
||||
PrivateTmp=true
|
||||
|
||||
[Install]
|
||||
WantedBy=default.target
|
||||
@@ -0,0 +1,142 @@
|
||||
"""Build the bastion agent distributable: one reproducible zipapp.
|
||||
|
||||
python3 -m bastion_agent.build [--out dist]
|
||||
|
||||
Produces dist/bastion-agent-<version>.pyz and dist/bastion-agent-<version>.pyz.sha256.
|
||||
The .pyz runs on any Python >= 3.8 on Mac, Windows or Linux:
|
||||
|
||||
python3 bastion-agent-0.1.0.pyz doctor
|
||||
python3 bastion-agent-0.1.0.pyz sweep manifest.yaml
|
||||
|
||||
Inside: the `bastion_agent` package (no tests), the `pii` package (no tests,
|
||||
no Dockerfile), and a vendored pure-Python PyYAML (MIT; the one import
|
||||
the manifest loader needs that is not stdlib — the C accelerator is left
|
||||
out on purpose so the same file runs everywhere). age and openssl stay
|
||||
what they are: single external binaries `doctor` checks for. The repo's
|
||||
`LICENSE` (GNU GPL v3 or later) is archived at the top level, so the
|
||||
license travels with the program as GPL-3 sections 4 and 5 require.
|
||||
|
||||
Reproducible: fixed timestamps, sorted members, no bytecode, and the
|
||||
vendored PyYAML pinned to one release (PYYAML_VERSION). Two builds of the
|
||||
same tree are byte-identical on any machine, so the .sha256 means something
|
||||
and anyone can rebuild a release and compare.
|
||||
A native per-OS binary (PyInstaller/Nuitka) needs a build machine per OS
|
||||
and is not produced here; this is the artifact until then.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
import zipfile
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
TOOLS = os.path.dirname(HERE)
|
||||
FIXED_TIME = (2026, 1, 1, 0, 0, 0)
|
||||
SHEBANG = b"#!/usr/bin/env python3\n"
|
||||
|
||||
MAIN_SHIM = '''"""bastion agent zipapp entry point."""
|
||||
import sys
|
||||
from bastion_agent.__main__ import main
|
||||
sys.exit(main())
|
||||
'''
|
||||
|
||||
YAML_NOTICE = """PyYAML (vendored, pure-Python modules only)
|
||||
License: MIT — Copyright (c) 2017-2021 Ingy döt Net, Copyright (c) 2006-2016 Kirill Simonov
|
||||
https://github.com/yaml/pyyaml — the C accelerator (_yaml) is deliberately not included.
|
||||
"""
|
||||
|
||||
|
||||
def version():
|
||||
with open(os.path.join(HERE, "VERSION"), encoding="utf-8") as fh:
|
||||
return fh.read().strip()
|
||||
|
||||
|
||||
def _package_files(src_dir, exclude_prefixes=("test_",), exclude_names=()):
|
||||
for name in sorted(os.listdir(src_dir)):
|
||||
if not name.endswith(".py") or name.startswith(exclude_prefixes) or name in exclude_names:
|
||||
continue
|
||||
with open(os.path.join(src_dir, name), "rb") as fh:
|
||||
yield name, fh.read()
|
||||
|
||||
|
||||
# The PyYAML release that is vendored. Its source is copied in from whatever
|
||||
# `import yaml` finds, so without a pin two machines with different PyYAML
|
||||
# builds would make different .pyz files from the same tree. Move it on
|
||||
# purpose, in its own commit.
|
||||
PYYAML_VERSION = "6.0.3"
|
||||
|
||||
|
||||
def _yaml_files():
|
||||
import yaml
|
||||
if getattr(yaml, "__version__", "") != PYYAML_VERSION:
|
||||
raise SystemExit(
|
||||
"build needs PyYAML %s to vendor (found %s at %s): pip install pyyaml==%s"
|
||||
% (PYYAML_VERSION, getattr(yaml, "__version__", "?"), yaml.__file__, PYYAML_VERSION))
|
||||
d = os.path.dirname(yaml.__file__)
|
||||
for name in sorted(os.listdir(d)):
|
||||
if name.endswith(".py"):
|
||||
with open(os.path.join(d, name), "rb") as fh:
|
||||
yield name, fh.read()
|
||||
|
||||
|
||||
def members():
|
||||
"""(arcname, bytes) for everything in the zipapp, deterministic order."""
|
||||
out = [("__main__.py", MAIN_SHIM.encode("utf-8"))]
|
||||
# build.py and verify_release.py are tools for making and checking a
|
||||
# release; neither is part of the program a bastion runs.
|
||||
for name, data in _package_files(HERE, exclude_names=("build.py", "verify_release.py")):
|
||||
out.append(("bastion_agent/" + name, data))
|
||||
with open(os.path.join(HERE, "VERSION"), "rb") as fh:
|
||||
out.append(("bastion_agent/VERSION", fh.read()))
|
||||
for name, data in _package_files(os.path.join(TOOLS, "pii")):
|
||||
out.append(("pii/" + name, data))
|
||||
for name, data in _yaml_files():
|
||||
out.append(("yaml/" + name, data))
|
||||
out.append(("yaml/LICENSE-NOTICE.txt", YAML_NOTICE.encode("utf-8")))
|
||||
with open(os.path.join(TOOLS, "LICENSE"), "rb") as fh:
|
||||
out.append(("LICENSE", fh.read()))
|
||||
return sorted(out, key=lambda m: m[0])
|
||||
|
||||
|
||||
def build_bytes():
|
||||
buf = io.BytesIO()
|
||||
buf.write(SHEBANG)
|
||||
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
|
||||
for arcname, data in members():
|
||||
info = zipfile.ZipInfo(arcname, date_time=FIXED_TIME)
|
||||
info.compress_type = zipfile.ZIP_DEFLATED
|
||||
info.external_attr = 0o644 << 16
|
||||
zf.writestr(info, data)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def build(out_dir="dist"):
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
data = build_bytes()
|
||||
ver = version()
|
||||
path = os.path.join(out_dir, "bastion-agent-%s.pyz" % ver)
|
||||
with open(path, "wb") as fh:
|
||||
fh.write(data)
|
||||
try:
|
||||
os.chmod(path, 0o755)
|
||||
except OSError:
|
||||
pass
|
||||
digest = hashlib.sha256(data).hexdigest()
|
||||
with open(path + ".sha256", "w", encoding="utf-8") as fh:
|
||||
fh.write("%s %s\n" % (digest, os.path.basename(path)))
|
||||
return path, digest, len(data)
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
ap = argparse.ArgumentParser(prog="bastion_agent.build")
|
||||
ap.add_argument("--out", default="dist")
|
||||
args = ap.parse_args(argv)
|
||||
path, digest, size = build(args.out)
|
||||
print("%s %d bytes sha256 %s" % (path, size, digest))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,444 @@
|
||||
"""bastion agent `catalog` stage — the curated dataset.
|
||||
|
||||
The catalog is what a host reasons over and what `seal` ships. One card
|
||||
per file, one per directory, small and structured, built from:
|
||||
|
||||
walk.jsonl metadata (path, size, mtime, mime, hash)
|
||||
extract.jsonl whether text came out, and how
|
||||
sanitize.jsonl entity counts (token types minted per file)
|
||||
sanitize/<sha>.txt the SANITIZED text — the only text this stage reads
|
||||
chunks.jsonl chunk count per document (optional)
|
||||
|
||||
Outputs:
|
||||
<out>/catalog.jsonl file cards, sorted by (root, path)
|
||||
<out>/catalog.dirs.jsonl directory cards, sorted by (root, path)
|
||||
<out>/catalog.summary.json
|
||||
|
||||
Every text-derived field (title, head, keywords, summary) comes from the
|
||||
sanitized store, so the catalog carries PII tokens, never values. File
|
||||
and directory NAMES are plaintext by nature; they go through the rules
|
||||
pass and any precision anchor in a name (an email, a card number) is
|
||||
masked as [PII_TYPE] in the card. Contextual PII in a filename
|
||||
("Marisol_medical.pdf") is not caught by rules — the honest limit of the
|
||||
sanitize stage applies here too, and the doc says so.
|
||||
|
||||
The optional summarizer (manifest `summarize`) is a local chat model — a
|
||||
sensor: it fills the `summary` string and nothing else, is reached via
|
||||
SUMMARIZE_URL (never the manifest), sees only sanitized text, and is
|
||||
cached per (document, model, version). The first failure stops
|
||||
summarizing (cards still get written, `summary_error` set, exit 1).
|
||||
"""
|
||||
|
||||
import datetime as _dt
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
from .manifest import signing_digest
|
||||
from .sanitize import store_path as sanitized_store_path
|
||||
|
||||
|
||||
def _rules_detect():
|
||||
"""pii.detect.detect — a sibling package here and in the .pyz (see
|
||||
sanitize.pii_modules for why the old `tools.pii` branch is gone)."""
|
||||
from pii.detect import detect
|
||||
return detect
|
||||
|
||||
CATALOG_VERSION = 1
|
||||
HEAD_CHARS = 300
|
||||
TITLE_CHARS = 120
|
||||
KEYWORDS_PER_CARD = 8
|
||||
KEYWORDS_PER_DIR = 12
|
||||
SUMMARY_INPUT_CHARS = 12000
|
||||
SUMMARY_MAX_CHARS = 600
|
||||
DEFAULT_SUMMARIZE_URL = "http://127.0.0.1:11434"
|
||||
SUMMARIZE_TIMEOUT = 300.0
|
||||
|
||||
WORD_RE = re.compile(r"[^\W\d_]{3,}", re.UNICODE)
|
||||
TOKEN_MARK_RE = re.compile(r"\[PII_[A-Z_]+_\d+\]")
|
||||
STOPWORDS = set("""
|
||||
the and for that with this from are was were will have has had not but you your our their they them
|
||||
its his her she him who what which when where why how all any can could would should may might must
|
||||
into onto over under than then there here also just more most much many some such very only own same
|
||||
other another each both few between about after before during while above below off out too again
|
||||
further once here per via use used using does did done being been each same so if or as at by on in
|
||||
of to is it be an a we us me my i no yes do go get got let new old one two three first last next see
|
||||
set run make made take took like well also back down up still even ever never always often sometimes
|
||||
""".split())
|
||||
|
||||
|
||||
class CatalogError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
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 _read_jsonl(out_dir, name, required=True):
|
||||
path = os.path.join(out_dir, name)
|
||||
if not os.path.exists(path):
|
||||
if required:
|
||||
raise FileNotFoundError("no %s at %s — run the earlier stages first" % (name, path))
|
||||
return []
|
||||
with open(path, encoding="utf-8") as fh:
|
||||
return [json.loads(l) for l in fh]
|
||||
|
||||
|
||||
# ── text facts ──────────────────────────────────────────────────────────────
|
||||
|
||||
def title_of(text):
|
||||
for line in text.splitlines():
|
||||
s = line.strip().lstrip("#*-= \t").strip()
|
||||
if len(s) >= 3:
|
||||
return s[:TITLE_CHARS]
|
||||
return ""
|
||||
|
||||
|
||||
def head_of(text):
|
||||
s = re.sub(r"\s+", " ", text).strip()
|
||||
return s[:HEAD_CHARS]
|
||||
|
||||
|
||||
def term_counts(text):
|
||||
counts = {}
|
||||
text = TOKEN_MARK_RE.sub(" ", text) # [PII_*_n] placeholders are not keywords
|
||||
for w in WORD_RE.findall(text.lower()):
|
||||
if w in STOPWORDS:
|
||||
continue
|
||||
counts[w] = counts.get(w, 0) + 1
|
||||
return counts
|
||||
|
||||
|
||||
def keywords(counts, df, ndocs, k):
|
||||
"""Top-k terms by TF-IDF; ties broken alphabetically for determinism."""
|
||||
scored = []
|
||||
for w, tf in counts.items():
|
||||
idf = math.log((1 + ndocs) / (1 + df.get(w, 0))) + 1.0
|
||||
scored.append((-(tf * idf), w))
|
||||
scored.sort()
|
||||
return [w for _, w in scored[:k]]
|
||||
|
||||
|
||||
def mask_name(name, detect):
|
||||
"""Rules pass over a file/dir name: precision anchors -> [PII_TYPE]."""
|
||||
findings = detect(name)
|
||||
if not findings:
|
||||
return name
|
||||
out = name
|
||||
for f in sorted(findings, key=lambda f: -len(f.value)):
|
||||
out = out.replace(f.value, "[PII_%s]" % f.type)
|
||||
return out
|
||||
|
||||
|
||||
def mask_path(path, detect):
|
||||
"""Every component of a relative path through mask_name — a parent
|
||||
directory can carry an address as easily as a filename, and the file
|
||||
card's path is what seal ships. Returns (masked
|
||||
path, number of components changed)."""
|
||||
if not path:
|
||||
return path, 0
|
||||
out, changed = [], 0
|
||||
for seg in path.split("/"):
|
||||
m = mask_name(seg, detect) if seg else seg
|
||||
if m != seg:
|
||||
changed += 1
|
||||
out.append(m)
|
||||
return "/".join(out), changed
|
||||
|
||||
|
||||
# ── the summarizer (a sensor) ───────────────────────────────────────────────
|
||||
|
||||
SUMMARY_PROMPT = (
|
||||
"Summarize the following document in at most three plain sentences for a catalog card: "
|
||||
"what it is, what it covers, and who or what it concerns. Placeholders like [PII_PERSON_1] "
|
||||
"stand for redacted personal data — keep them as they are, never guess what they hide. "
|
||||
"Reply with the summary only.\n\n"
|
||||
)
|
||||
DIR_SUMMARY_PROMPT = (
|
||||
"These are catalog entries for the files in one folder. In at most three plain sentences "
|
||||
"say what the folder is for and what kinds of material it holds. Placeholders like "
|
||||
"[PII_PERSON_1] are redacted personal data — keep them as they are. Reply with the "
|
||||
"summary only.\n\n"
|
||||
)
|
||||
|
||||
|
||||
class Summarizer:
|
||||
def __init__(self, url, model, api=None, timeout=SUMMARIZE_TIMEOUT):
|
||||
self.url = url.rstrip("/")
|
||||
self.model = model
|
||||
self.api = api
|
||||
self.timeout = timeout
|
||||
|
||||
def _post(self, path, body):
|
||||
req = urllib.request.Request(
|
||||
self.url + path, data=json.dumps(body).encode("utf-8"),
|
||||
headers={"Content-Type": "application/json"})
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=self.timeout) as resp:
|
||||
return json.load(resp), None
|
||||
except urllib.error.HTTPError as e:
|
||||
return None, "HTTP %d" % e.code
|
||||
except (OSError, ValueError) as e:
|
||||
return None, str(e)
|
||||
|
||||
def _ollama(self, prompt):
|
||||
data, err = self._post("/api/chat", {
|
||||
"model": self.model, "stream": False, "options": {"temperature": 0},
|
||||
"messages": [{"role": "user", "content": prompt}]})
|
||||
if err:
|
||||
return None, err
|
||||
try:
|
||||
return data["message"]["content"], None
|
||||
except (KeyError, TypeError) as e:
|
||||
return None, "malformed ollama response: %s" % e
|
||||
|
||||
def _openai(self, prompt):
|
||||
data, err = self._post("/v1/chat/completions", {
|
||||
"model": self.model, "temperature": 0,
|
||||
"messages": [{"role": "user", "content": prompt}]})
|
||||
if err:
|
||||
return None, err
|
||||
try:
|
||||
return data["choices"][0]["message"]["content"], None
|
||||
except (KeyError, IndexError, TypeError) as e:
|
||||
return None, "malformed chat response: %s" % e
|
||||
|
||||
def complete(self, prompt):
|
||||
if self.api == "ollama":
|
||||
text, err = self._ollama(prompt)
|
||||
elif self.api == "openai":
|
||||
text, err = self._openai(prompt)
|
||||
else:
|
||||
text, err = self._ollama(prompt)
|
||||
if text is not None:
|
||||
self.api = "ollama"
|
||||
else:
|
||||
text, err2 = self._openai(prompt)
|
||||
if text is not None:
|
||||
self.api = "openai"
|
||||
else:
|
||||
err = "ollama: %s; openai: %s" % (err, err2)
|
||||
if text is None:
|
||||
raise CatalogError("summarizer %s (model %s): %s" % (self.url, self.model, err))
|
||||
return clean_summary(text)
|
||||
|
||||
def describe(self):
|
||||
return "%s@%s(%s)" % (self.model, self.url, self.api or "auto")
|
||||
|
||||
|
||||
def clean_summary(text):
|
||||
text = re.sub(r"<think>.*?</think>", "", text, flags=re.S) # thinking models
|
||||
text = re.sub(r"\s+", " ", text).strip().strip('"')
|
||||
return text[:SUMMARY_MAX_CHARS]
|
||||
|
||||
|
||||
def build_summarizer(manifest, env):
|
||||
if manifest.summarize is None:
|
||||
return None
|
||||
url = env.get("SUMMARIZE_URL", DEFAULT_SUMMARIZE_URL).strip() or DEFAULT_SUMMARIZE_URL
|
||||
api = env.get("SUMMARIZE_API", "").strip().lower() or None
|
||||
if api not in (None, "ollama", "openai"):
|
||||
raise CatalogError("SUMMARIZE_API=%r is not ollama|openai" % api)
|
||||
return Summarizer(url, manifest.summarize.model, api=api,
|
||||
timeout=float(env.get("SUMMARIZE_TIMEOUT", SUMMARIZE_TIMEOUT)))
|
||||
|
||||
|
||||
def _cache_path(out_dir):
|
||||
return os.path.join(out_dir, "catalog.cache.json")
|
||||
|
||||
|
||||
def _load_cache(out_dir):
|
||||
try:
|
||||
with open(_cache_path(out_dir), encoding="utf-8") as fh:
|
||||
return json.load(fh)
|
||||
except (OSError, ValueError):
|
||||
return {}
|
||||
|
||||
|
||||
# ── the stage ───────────────────────────────────────────────────────────────
|
||||
|
||||
def catalog(manifest, out_dir, env=None, summarizer=None, now=None):
|
||||
env = dict(os.environ) if env is None else env
|
||||
summ = summarizer if summarizer is not None else build_summarizer(manifest, env)
|
||||
per = set(manifest.summarize.per) if manifest.summarize else set()
|
||||
walk_rows = _read_jsonl(out_dir, "walk.jsonl")
|
||||
extract_rows = {(r["root"], r["path"]): r for r in _read_jsonl(out_dir, "extract.jsonl")}
|
||||
sanitize_rows = {(r["root"], r["path"]): r for r in _read_jsonl(out_dir, "sanitize.jsonl")}
|
||||
chunk_counts = {}
|
||||
for c in _read_jsonl(out_dir, "chunks.jsonl", required=False):
|
||||
chunk_counts[c["doc"]] = chunk_counts.get(c["doc"], 0) + 1
|
||||
detect = _rules_detect()
|
||||
cache = _load_cache(out_dir)
|
||||
stamp = now or _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
|
||||
# pass 1: text facts per distinct sanitized document
|
||||
texts = {}
|
||||
doc_terms = {}
|
||||
for key, s in sanitize_rows.items():
|
||||
if s.get("status") != "ok":
|
||||
continue
|
||||
sha = s["sha256"]
|
||||
if sha in texts:
|
||||
continue
|
||||
with open(sanitized_store_path(out_dir, sha), encoding="utf-8") as fh:
|
||||
texts[sha] = fh.read()
|
||||
doc_terms[sha] = term_counts(texts[sha])
|
||||
df = {}
|
||||
for counts in doc_terms.values():
|
||||
for w in counts:
|
||||
df[w] = df.get(w, 0) + 1
|
||||
ndocs = len(doc_terms)
|
||||
|
||||
counts = {"files": 0, "with_text": 0, "dirs": 0, "summaries": 0, "summaries_reused": 0,
|
||||
"summary_errors": 0, "names_masked": 0}
|
||||
stopped = None
|
||||
summary_key = lambda sha, kind: "%s|%s|%s|v%d" % (
|
||||
kind, sha, manifest.summarize.model if manifest.summarize else "", CATALOG_VERSION)
|
||||
|
||||
def summarize(kind, sha, prompt_body):
|
||||
nonlocal stopped
|
||||
if summ is None or stopped is not None:
|
||||
return None, None
|
||||
k = summary_key(sha, kind)
|
||||
if k in cache:
|
||||
counts["summaries_reused"] += 1
|
||||
return cache[k], None
|
||||
try:
|
||||
text = summ.complete((SUMMARY_PROMPT if kind == "file" else DIR_SUMMARY_PROMPT)
|
||||
+ prompt_body[:SUMMARY_INPUT_CHARS])
|
||||
except CatalogError as e:
|
||||
stopped = str(e)
|
||||
counts["summary_errors"] += 1
|
||||
return None, stopped
|
||||
cache[k] = text
|
||||
counts["summaries"] += 1
|
||||
return text, None
|
||||
|
||||
# pass 2: file cards
|
||||
cards = []
|
||||
dir_acc = {}
|
||||
for w in walk_rows:
|
||||
if w.get("kind") != "file":
|
||||
continue
|
||||
key = (w["root"], w["path"])
|
||||
counts["files"] += 1
|
||||
name = w["path"].rsplit("/", 1)[-1]
|
||||
shown_path, changed = mask_path(w["path"], detect)
|
||||
masked = shown_path.rsplit("/", 1)[-1]
|
||||
counts["names_masked"] += changed
|
||||
e = extract_rows.get(key, {})
|
||||
s = sanitize_rows.get(key, {})
|
||||
card = {
|
||||
"root": w["root"], "path": shown_path, "name": masked, "kind": "file",
|
||||
"size": w.get("size"), "mtime": w.get("mtime"), "mime": w.get("mime"),
|
||||
"sha256": w.get("sha256"), "v": CATALOG_VERSION,
|
||||
"text": e.get("status") if e else "not_extracted",
|
||||
"extractor": e.get("extractor"),
|
||||
"sanitized": s.get("status") if s else None,
|
||||
}
|
||||
sha = s.get("sha256") if s.get("status") == "ok" else None
|
||||
if sha and sha in texts:
|
||||
text = texts[sha]
|
||||
counts["with_text"] += 1
|
||||
card.update({
|
||||
"title": title_of(text),
|
||||
"head": head_of(text),
|
||||
"chars": len(text),
|
||||
"keywords": keywords(doc_terms[sha], df, ndocs, KEYWORDS_PER_CARD),
|
||||
"entities": s.get("types", {}),
|
||||
"chunks": chunk_counts.get(sha, 0),
|
||||
})
|
||||
if "file" in per:
|
||||
summary, err = summarize("file", sha, text)
|
||||
card["summary"] = summary
|
||||
if err:
|
||||
card["summary_error"] = err
|
||||
elif "file" in per:
|
||||
card["summary"] = None # nothing to summarize
|
||||
cards.append(card)
|
||||
# directory aggregation: every ancestor, root included as ""
|
||||
parts = w["path"].split("/")[:-1]
|
||||
for depth in range(len(parts) + 1):
|
||||
dpath = "/".join(parts[:depth])
|
||||
acc = dir_acc.setdefault((w["root"], dpath), {
|
||||
"files": 0, "bytes": 0, "mimes": {}, "latest": "", "terms": {}, "entities": {},
|
||||
"children": [], "with_text": 0})
|
||||
acc["files"] += 1
|
||||
acc["bytes"] += w.get("size") or 0
|
||||
m = w.get("mime") or "?"
|
||||
acc["mimes"][m] = acc["mimes"].get(m, 0) + 1
|
||||
if (w.get("mtime") or "") > acc["latest"]:
|
||||
acc["latest"] = w["mtime"]
|
||||
if sha and sha in doc_terms:
|
||||
acc["with_text"] += 1
|
||||
for t, n in doc_terms[sha].items():
|
||||
acc["terms"][t] = acc["terms"].get(t, 0) + n
|
||||
for t, n in (s.get("types") or {}).items():
|
||||
acc["entities"][t] = acc["entities"].get(t, 0) + n
|
||||
if depth == len(parts):
|
||||
acc["children"].append(card)
|
||||
cards.sort(key=lambda c: (c["root"], c["path"]))
|
||||
|
||||
# pass 3: directory cards
|
||||
dcards = []
|
||||
for (root, dpath), acc in dir_acc.items():
|
||||
shown, _ = mask_path(dpath, detect)
|
||||
masked = shown.rsplit("/", 1)[-1] if shown else ""
|
||||
mimes = sorted(acc["mimes"].items(), key=lambda kv: (-kv[1], kv[0]))[:6]
|
||||
card = {
|
||||
"root": root, "path": shown, "name": masked, "kind": "dir", "v": CATALOG_VERSION,
|
||||
"files": acc["files"], "with_text": acc["with_text"], "bytes": acc["bytes"],
|
||||
"latest": acc["latest"], "types": dict(mimes),
|
||||
"keywords": keywords(acc["terms"], df, ndocs, KEYWORDS_PER_DIR),
|
||||
"entities": acc["entities"],
|
||||
}
|
||||
if "directory" in per:
|
||||
body = "\n".join(
|
||||
"- %s: %s" % (c["name"], (c.get("summary") or c.get("head") or "")[:200])
|
||||
for c in sorted(acc["children"], key=lambda c: c["path"])[:60])
|
||||
sha = hashlib.sha256(
|
||||
("%s|%s|" % (root, dpath)).encode("utf-8") + body.encode("utf-8")).hexdigest()
|
||||
summary, err = summarize("dir", sha, body)
|
||||
card["summary"] = summary
|
||||
if err:
|
||||
card["summary_error"] = err
|
||||
dcards.append(card)
|
||||
dcards.sort(key=lambda c: (c["root"], c["path"]))
|
||||
counts["dirs"] = len(dcards)
|
||||
|
||||
_atomic_write_text(os.path.join(out_dir, "catalog.jsonl"), "".join(
|
||||
json.dumps(c, ensure_ascii=False, sort_keys=True) + "\n" for c in cards))
|
||||
_atomic_write_text(os.path.join(out_dir, "catalog.dirs.jsonl"), "".join(
|
||||
json.dumps(c, ensure_ascii=False, sort_keys=True) + "\n" for c in dcards))
|
||||
_atomic_write_text(_cache_path(out_dir), json.dumps(cache, ensure_ascii=False, sort_keys=True))
|
||||
summary = {
|
||||
"freeagent": 1, "stage": "catalog", "catalog_version": CATALOG_VERSION,
|
||||
"manifest_id": manifest.id,
|
||||
"manifest_digest": "sha256:" + signing_digest(manifest.raw),
|
||||
"cataloged_at": stamp,
|
||||
"summarizer": summ.describe() if summ else None,
|
||||
"summarize_per": sorted(per),
|
||||
"counts": counts, "stopped": stopped,
|
||||
"bytes_catalog": sum(len(json.dumps(c, ensure_ascii=False)) + 1 for c in cards + dcards),
|
||||
}
|
||||
_atomic_write_text(os.path.join(out_dir, "catalog.summary.json"),
|
||||
json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True) + "\n")
|
||||
return summary
|
||||
@@ -0,0 +1,119 @@
|
||||
"""`bastion-agent doctor` — the preflight a member runs before the first sweep.
|
||||
|
||||
Reports what this machine has and what it lacks, per stage, and never
|
||||
changes anything. Required items are what a stage refuses to run
|
||||
without; optional ones are the sensors that widen coverage (pdftotext,
|
||||
tesseract, the gliner package) — their absence is a recorded gap in the
|
||||
stage output, not a failure here.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import shutil
|
||||
import sys
|
||||
import urllib.request
|
||||
|
||||
PY_FLOOR = (3, 8)
|
||||
|
||||
|
||||
def _which(name):
|
||||
return shutil.which(name)
|
||||
|
||||
|
||||
def _probe_http(url, path, timeout=5):
|
||||
"""Reachable means the server answered — any HTTP status counts (a
|
||||
proxy that only permits the embed path answers 403 on '/'). Only a
|
||||
connection failure is 'not reachable'."""
|
||||
import urllib.error
|
||||
try:
|
||||
with urllib.request.urlopen(url.rstrip("/") + path, timeout=timeout) as r:
|
||||
return True, "HTTP %d" % r.status
|
||||
except urllib.error.HTTPError as e:
|
||||
return True, "reachable (HTTP %d on %s)" % (e.code, path)
|
||||
except Exception as e:
|
||||
return False, str(e)[:120]
|
||||
|
||||
|
||||
def checks(env=None):
|
||||
env = dict(os.environ) if env is None else env
|
||||
out = []
|
||||
|
||||
def add(name, ok, detail, required=True, stage=None):
|
||||
out.append({"check": name, "ok": bool(ok), "detail": detail, "required": required,
|
||||
"stage": stage})
|
||||
|
||||
v = sys.version_info[:3]
|
||||
add("python >= %d.%d" % PY_FLOOR, v >= PY_FLOOR, "%s %s (%s)" % (platform.python_implementation(),
|
||||
".".join(map(str, v)), sys.executable))
|
||||
add("platform", True, "%s %s %s" % (platform.system(), platform.release(), platform.machine()),
|
||||
required=False)
|
||||
try:
|
||||
import yaml # noqa: F401
|
||||
add("yaml (manifest loader)", True, "PyYAML %s" % getattr(yaml, "__version__", "?"))
|
||||
except ImportError:
|
||||
add("yaml (manifest loader)", False, "PyYAML not importable")
|
||||
|
||||
for tool, stage, req, why in (
|
||||
("age", "seal", True, "the envelope; https://age-encryption.org"),
|
||||
# Not needed to RUN the agent: the gate verifies in the standard
|
||||
# library (x509gate) and the bastion signs with its own key.
|
||||
("openssl", "sign/seal", False, "only for `sign --key` and SEAL_SIGN_KEY seal signatures"),
|
||||
("pdftotext", "extract", False, "PDF text (poppler); without it simple PDFs only"),
|
||||
("tesseract", "extract", False, "image OCR; without it images are recorded unsupported")):
|
||||
p = _which(tool)
|
||||
add(tool, bool(p), p or ("not on PATH — " + why), required=req, stage=stage)
|
||||
|
||||
try:
|
||||
import gliner # noqa: F401
|
||||
add("gliner package", True, "importable (PII_RECALL=gliner available)", required=False, stage="sanitize")
|
||||
except ImportError:
|
||||
add("gliner package", False, "not installed — sanitize sterilize on a laptop needs "
|
||||
"PII_RECALL=gliner (pip install gliner) or a reachable LLM_URL / PII_ADDR",
|
||||
required=False, stage="sanitize")
|
||||
|
||||
embed_url = env.get("EMBED_URL", "http://127.0.0.1:11434")
|
||||
ok, detail = _probe_http(embed_url, "/")
|
||||
add("embedding endpoint %s" % embed_url, ok, detail, required=False, stage="embed")
|
||||
if env.get("PII_ADDR"):
|
||||
url = env["PII_ADDR"] if "://" in env["PII_ADDR"] else "http://" + env["PII_ADDR"]
|
||||
ok, detail = _probe_http(url, "/health")
|
||||
add("pii service %s" % url, ok, detail, required=False, stage="sanitize")
|
||||
if env.get("SUMMARIZE_URL"):
|
||||
ok, detail = _probe_http(env["SUMMARIZE_URL"], "/")
|
||||
add("summarizer %s" % env["SUMMARIZE_URL"], ok, detail, required=False, stage="catalog")
|
||||
|
||||
for var, stage, why in (
|
||||
("FREEAGENT_OWNER_RECIPIENT", "seal", "owner's age public key"),
|
||||
("FREEAGENT_CA_ROOT", "sweep/serve", "CA root for the manifest gate and client certs"),
|
||||
("FREEAGENT_TLS_CERT", "serve", "this machine's cert"),
|
||||
("FREEAGENT_TLS_KEY", "serve", "this machine's key")):
|
||||
val = env.get(var, "")
|
||||
if var.endswith(("_ROOT", "_CERT", "_KEY")) and val:
|
||||
ok, detail = os.path.exists(val), (val if os.path.exists(val) else val + " (missing file)")
|
||||
else:
|
||||
ok, detail = bool(val), (("set" if val else "unset — " + why))
|
||||
add(var, ok, detail, required=False, stage=stage)
|
||||
return out
|
||||
|
||||
|
||||
def render(results):
|
||||
lines = []
|
||||
for r in results:
|
||||
mark = "ok " if r["ok"] else ("MISS" if r["required"] else "opt ")
|
||||
stage = (" [%s]" % r["stage"]) if r["stage"] else ""
|
||||
lines.append("%s %-34s %s%s" % (mark, r["check"], r["detail"], stage))
|
||||
missing = [r for r in results if r["required"] and not r["ok"]]
|
||||
lines.append("")
|
||||
lines.append("required missing: %d" % len(missing) if missing else "all required checks pass")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def main(argv=None, env=None):
|
||||
argv = sys.argv[1:] if argv is None else argv
|
||||
results = checks(env)
|
||||
if "--json" in argv:
|
||||
print(json.dumps(results, indent=2))
|
||||
else:
|
||||
print(render(results))
|
||||
return 1 if any(r["required"] and not r["ok"] for r in results) else 0
|
||||
@@ -0,0 +1,417 @@
|
||||
"""bastion agent `chunk` + `embed` stages, and the local `search` that proves them.
|
||||
|
||||
chunk sanitized store -> <out>/chunks.jsonl
|
||||
Overlapping windows sized in approximate tokens (manifest
|
||||
embed.chunk.size_tokens / overlap_tokens). "Token" here is a
|
||||
deterministic regex unit (a word or a punctuation mark), not a
|
||||
model tokenizer: the stdlib has none, and the same text must chunk
|
||||
identically on every machine. Chunk id = sha256(doc sha, window
|
||||
text) — content-addressed, so an unchanged document keeps its
|
||||
vectors; a changed document has a new sha and is embedded again.
|
||||
|
||||
embed chunks -> <out>/vectors.db (SQLite, stdlib)
|
||||
The endpoint is resolved from the environment at run time
|
||||
(EMBED_URL, default the local Ollama port), never the manifest;
|
||||
the manifest pins the MODEL and DIM. Before the first real call
|
||||
the stage embeds a probe and refuses to run if the served
|
||||
dimension is not the manifest's — a wrong model never writes a
|
||||
vector. Two wire formats, auto-detected: Ollama /api/embed and
|
||||
OpenAI-style /v1/embeddings (llama.cpp server, TEI).
|
||||
Rows are keyed (chunk_id, model): switching models ADDS rows and
|
||||
never deletes the old ones, so going back to an earlier model
|
||||
costs nothing and no vector is ever thrown away.
|
||||
|
||||
search embed the query the same way, cosine over every vector of the
|
||||
manifest's model, return the top-k sanitized chunks. Pure Python;
|
||||
fine to tens of thousands of chunks, which is the point where a
|
||||
real index is a follow-up.
|
||||
"""
|
||||
|
||||
import array
|
||||
import datetime as _dt
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
from .manifest import signing_digest
|
||||
from .sanitize import store_path as sanitized_store_path
|
||||
|
||||
CHUNK_VERSION = 1
|
||||
DEFAULT_EMBED_URL = "http://127.0.0.1:11434" # local Ollama; a sensor, resolved at run time
|
||||
EMBED_BATCH = 32
|
||||
EMBED_TIMEOUT = 300.0
|
||||
PROBE_TEXT = "freeagent embedding probe"
|
||||
|
||||
TOKEN_RE = re.compile(r"\w+|[^\w\s]", re.UNICODE)
|
||||
|
||||
|
||||
class EmbedError(Exception):
|
||||
"""Fail-closed stop: the endpoint is unreachable or serves the wrong model."""
|
||||
|
||||
|
||||
# ── chunking ────────────────────────────────────────────────────────────────
|
||||
|
||||
def count_tokens(text):
|
||||
return len(TOKEN_RE.findall(text))
|
||||
|
||||
|
||||
def chunk_text(text, size_tokens, overlap_tokens):
|
||||
"""Yield (start, end, window_text). Windows are cut at token
|
||||
boundaries; the step is size - overlap tokens; the last window may be
|
||||
shorter. Every character of text is inside at least one window."""
|
||||
spans = [(m.start(), m.end()) for m in TOKEN_RE.finditer(text)]
|
||||
if not spans:
|
||||
return
|
||||
step = max(1, size_tokens - overlap_tokens)
|
||||
i = 0
|
||||
n = len(spans)
|
||||
while True:
|
||||
j = min(n, i + size_tokens)
|
||||
start = 0 if i == 0 else spans[i][0]
|
||||
end = len(text) if j == n else spans[j - 1][1]
|
||||
yield start, end, text[start:end]
|
||||
if j >= n:
|
||||
return
|
||||
i += step
|
||||
|
||||
|
||||
def chunk_id(doc_sha, window_text):
|
||||
h = hashlib.sha256()
|
||||
h.update(doc_sha.encode("ascii"))
|
||||
h.update(b"\x00")
|
||||
h.update(window_text.encode("utf-8"))
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
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 load_sanitized_index(out_dir):
|
||||
path = os.path.join(out_dir, "sanitize.jsonl")
|
||||
if not os.path.exists(path):
|
||||
raise FileNotFoundError("no sanitize index at %s — run the sanitize stage first" % path)
|
||||
with open(path, encoding="utf-8") as fh:
|
||||
return [json.loads(l) for l in fh]
|
||||
|
||||
|
||||
def chunk(manifest, out_dir, now=None):
|
||||
cfg = manifest.embed
|
||||
if cfg is None:
|
||||
raise EmbedError("manifest has no embed section / stage")
|
||||
records = load_sanitized_index(out_dir)
|
||||
counts = {"docs": 0, "chunks": 0, "skipped": 0}
|
||||
seen_docs = set()
|
||||
rows = []
|
||||
for rec in records:
|
||||
if rec.get("status") != "ok":
|
||||
counts["skipped"] += 1
|
||||
continue
|
||||
sha = rec["sha256"]
|
||||
if sha in seen_docs:
|
||||
continue # duplicate content chunks once; paths resolve through sanitize.jsonl
|
||||
seen_docs.add(sha)
|
||||
counts["docs"] += 1
|
||||
with open(sanitized_store_path(out_dir, sha), encoding="utf-8") as fh:
|
||||
text = fh.read()
|
||||
for idx, (start, end, window) in enumerate(
|
||||
chunk_text(text, cfg.chunk_size_tokens, cfg.chunk_overlap_tokens)):
|
||||
rows.append({"chunk_id": chunk_id(sha, window), "doc": sha, "idx": idx,
|
||||
"start": start, "end": end, "chars": len(window),
|
||||
"tokens": count_tokens(window), "v": CHUNK_VERSION})
|
||||
counts["chunks"] += 1
|
||||
rows.sort(key=lambda r: (r["doc"], r["idx"]))
|
||||
_atomic_write_text(os.path.join(out_dir, "chunks.jsonl"), "".join(
|
||||
json.dumps(r, sort_keys=True) + "\n" for r in rows))
|
||||
summary = {
|
||||
"freeagent": 1, "stage": "chunk", "chunk_version": CHUNK_VERSION,
|
||||
"manifest_id": manifest.id,
|
||||
"manifest_digest": "sha256:" + signing_digest(manifest.raw),
|
||||
"chunked_at": now or _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
|
||||
"size_tokens": cfg.chunk_size_tokens, "overlap_tokens": cfg.chunk_overlap_tokens,
|
||||
"counts": counts,
|
||||
}
|
||||
_atomic_write_text(os.path.join(out_dir, "chunk.summary.json"),
|
||||
json.dumps(summary, indent=2, sort_keys=True) + "\n")
|
||||
return summary
|
||||
|
||||
|
||||
# ── the vector store ────────────────────────────────────────────────────────
|
||||
|
||||
_SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS chunks (
|
||||
chunk_id TEXT PRIMARY KEY,
|
||||
doc TEXT NOT NULL,
|
||||
idx INTEGER NOT NULL,
|
||||
start INTEGER NOT NULL,
|
||||
end INTEGER NOT NULL,
|
||||
text TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS chunks_doc ON chunks(doc, idx);
|
||||
CREATE TABLE IF NOT EXISTS vectors (
|
||||
chunk_id TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
dim INTEGER NOT NULL,
|
||||
vec BLOB NOT NULL,
|
||||
created TEXT NOT NULL,
|
||||
PRIMARY KEY (chunk_id, model)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS vectors_model ON vectors(model);
|
||||
"""
|
||||
|
||||
|
||||
def vectors_path(out_dir):
|
||||
return os.path.join(out_dir, "vectors.db")
|
||||
|
||||
|
||||
def open_store(out_dir):
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
conn = sqlite3.connect(vectors_path(out_dir), timeout=30)
|
||||
conn.executescript(_SCHEMA)
|
||||
return conn
|
||||
|
||||
|
||||
def pack(vec):
|
||||
a = array.array("f", vec)
|
||||
return a.tobytes()
|
||||
|
||||
|
||||
def unpack(blob):
|
||||
a = array.array("f")
|
||||
a.frombytes(blob)
|
||||
return a
|
||||
|
||||
|
||||
# ── the endpoint (a sensor) ─────────────────────────────────────────────────
|
||||
|
||||
class Endpoint:
|
||||
"""Local embedding endpoint. `api` is 'ollama' or 'openai'; None
|
||||
auto-detects on the probe."""
|
||||
|
||||
def __init__(self, url, model, api=None, timeout=EMBED_TIMEOUT):
|
||||
self.url = url.rstrip("/")
|
||||
self.model = model
|
||||
self.api = api
|
||||
self.timeout = timeout
|
||||
|
||||
def _post(self, path, body):
|
||||
req = urllib.request.Request(
|
||||
self.url + path, data=json.dumps(body).encode("utf-8"),
|
||||
headers={"Content-Type": "application/json"})
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=self.timeout) as resp:
|
||||
return json.load(resp), None
|
||||
except urllib.error.HTTPError as e:
|
||||
detail = ""
|
||||
try:
|
||||
detail = e.read().decode("utf-8", "replace")[:200]
|
||||
except Exception:
|
||||
pass
|
||||
return None, "HTTP %d %s" % (e.code, detail)
|
||||
except (OSError, ValueError) as e:
|
||||
return None, str(e)
|
||||
|
||||
def _ollama(self, texts):
|
||||
data, err = self._post("/api/embed", {"model": self.model, "input": texts})
|
||||
if err:
|
||||
return None, err
|
||||
try:
|
||||
return [list(map(float, v)) for v in data["embeddings"]], None
|
||||
except (KeyError, TypeError, ValueError) as e:
|
||||
return None, "malformed ollama response: %s" % e
|
||||
|
||||
def _openai(self, texts):
|
||||
data, err = self._post("/v1/embeddings", {"model": self.model, "input": texts})
|
||||
if err:
|
||||
return None, err
|
||||
try:
|
||||
items = sorted(data["data"], key=lambda d: d.get("index", 0))
|
||||
return [list(map(float, d["embedding"])) for d in items], None
|
||||
except (KeyError, TypeError, ValueError) as e:
|
||||
return None, "malformed embeddings response: %s" % e
|
||||
|
||||
def embed(self, texts):
|
||||
if self.api == "ollama":
|
||||
vecs, err = self._ollama(texts)
|
||||
elif self.api == "openai":
|
||||
vecs, err = self._openai(texts)
|
||||
else:
|
||||
vecs, err = self._ollama(texts)
|
||||
if vecs is not None:
|
||||
self.api = "ollama"
|
||||
else:
|
||||
vecs, err2 = self._openai(texts)
|
||||
if vecs is not None:
|
||||
self.api = "openai"
|
||||
else:
|
||||
err = "ollama: %s; openai: %s" % (err, err2)
|
||||
if vecs is None:
|
||||
raise EmbedError("embedding endpoint %s (model %s): %s" % (self.url, self.model, err))
|
||||
if len(vecs) != len(texts):
|
||||
raise EmbedError("endpoint returned %d vectors for %d inputs" % (len(vecs), len(texts)))
|
||||
return vecs
|
||||
|
||||
def probe(self, want_dim):
|
||||
vec = self.embed([PROBE_TEXT])[0]
|
||||
if len(vec) != want_dim:
|
||||
raise EmbedError("model %s serves %d-dim vectors; manifest pins %d — refusing to write"
|
||||
% (self.model, len(vec), want_dim))
|
||||
return len(vec)
|
||||
|
||||
def describe(self):
|
||||
return "%s@%s(%s)" % (self.model, self.url, self.api or "auto")
|
||||
|
||||
|
||||
def build_endpoint(manifest, env):
|
||||
url = env.get("EMBED_URL", DEFAULT_EMBED_URL).strip() or DEFAULT_EMBED_URL
|
||||
api = env.get("EMBED_API", "").strip().lower() or None
|
||||
if api not in (None, "ollama", "openai"):
|
||||
raise EmbedError("EMBED_API=%r is not ollama|openai" % api)
|
||||
return Endpoint(url, manifest.embed.model, api=api,
|
||||
timeout=float(env.get("EMBED_TIMEOUT", EMBED_TIMEOUT)))
|
||||
|
||||
|
||||
# ── embed stage ─────────────────────────────────────────────────────────────
|
||||
|
||||
def load_chunks(out_dir):
|
||||
path = os.path.join(out_dir, "chunks.jsonl")
|
||||
if not os.path.exists(path):
|
||||
raise FileNotFoundError("no chunks at %s — run the chunk stage first" % path)
|
||||
with open(path, encoding="utf-8") as fh:
|
||||
return [json.loads(l) for l in fh]
|
||||
|
||||
|
||||
def embed(manifest, out_dir, env=None, endpoint=None, now=None, batch=None):
|
||||
cfg = manifest.embed
|
||||
if cfg is None:
|
||||
raise EmbedError("manifest has no embed section / stage")
|
||||
env = dict(os.environ) if env is None else env
|
||||
ep = endpoint or build_endpoint(manifest, env)
|
||||
batch = batch or int(env.get("EMBED_BATCH", EMBED_BATCH))
|
||||
chunks = load_chunks(out_dir)
|
||||
ep.probe(cfg.dim)
|
||||
|
||||
conn = open_store(out_dir)
|
||||
counts = {"chunks": len(chunks), "unique": 0, "embedded": 0, "reused": 0, "docs": 0}
|
||||
try:
|
||||
have = {r[0] for r in conn.execute("SELECT chunk_id FROM vectors WHERE model=?", (cfg.model,))}
|
||||
# chunk texts come from the sanitized store, sliced by recorded offsets
|
||||
texts = {}
|
||||
docs = {}
|
||||
for c in chunks:
|
||||
docs.setdefault(c["doc"], []).append(c)
|
||||
counts["docs"] = len(docs)
|
||||
for doc, cs in docs.items():
|
||||
with open(sanitized_store_path(out_dir, doc), encoding="utf-8") as fh:
|
||||
text = fh.read()
|
||||
for c in cs:
|
||||
texts[c["chunk_id"]] = text[c["start"]:c["end"]]
|
||||
# `chunks` is derived from THIS sweep (rebuilt each run); `vectors`
|
||||
# is the durable store and is only ever added to.
|
||||
with conn:
|
||||
conn.execute("DELETE FROM chunks")
|
||||
conn.executemany(
|
||||
"INSERT OR REPLACE INTO chunks(chunk_id, doc, idx, start, end, text) VALUES (?,?,?,?,?,?)",
|
||||
[(c["chunk_id"], c["doc"], c["idx"], c["start"], c["end"], texts[c["chunk_id"]])
|
||||
for c in chunks])
|
||||
# identical windows (a repeated paragraph) share a chunk id: embed once
|
||||
unique = {}
|
||||
for c in chunks:
|
||||
unique.setdefault(c["chunk_id"], c)
|
||||
counts["unique"] = len(unique)
|
||||
todo = [c for c in unique.values() if c["chunk_id"] not in have]
|
||||
counts["reused"] = len(unique) - len(todo)
|
||||
stamp = now or _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
for i in range(0, len(todo), batch):
|
||||
group = todo[i:i + batch]
|
||||
vecs = ep.embed([texts[c["chunk_id"]] for c in group])
|
||||
rows = []
|
||||
for c, v in zip(group, vecs):
|
||||
if len(v) != cfg.dim:
|
||||
raise EmbedError("vector of dim %d for chunk %s; manifest pins %d"
|
||||
% (len(v), c["chunk_id"][:12], cfg.dim))
|
||||
rows.append((c["chunk_id"], cfg.model, cfg.dim, pack(v), stamp))
|
||||
with conn:
|
||||
conn.executemany(
|
||||
"INSERT OR REPLACE INTO vectors(chunk_id, model, dim, vec, created) VALUES (?,?,?,?,?)",
|
||||
rows)
|
||||
counts["embedded"] += len(rows)
|
||||
# vectors for chunks that no longer exist in this sweep are kept
|
||||
# (never dump); report how many so the host can see the drift
|
||||
stale = conn.execute(
|
||||
"SELECT COUNT(*) FROM vectors v WHERE v.model=? AND NOT EXISTS "
|
||||
"(SELECT 1 FROM chunks c WHERE c.chunk_id=v.chunk_id)", (cfg.model,)).fetchone()[0]
|
||||
models = [dict(model=m, dim=d, vectors=n) for m, d, n in conn.execute(
|
||||
"SELECT model, dim, COUNT(*) FROM vectors GROUP BY model, dim ORDER BY model")]
|
||||
finally:
|
||||
conn.close()
|
||||
summary = {
|
||||
"freeagent": 1, "stage": "embed",
|
||||
"manifest_id": manifest.id,
|
||||
"manifest_digest": "sha256:" + signing_digest(manifest.raw),
|
||||
"embedded_at": now or _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
|
||||
"endpoint": ep.describe(), "model": cfg.model, "dim": cfg.dim,
|
||||
"counts": counts, "stale_vectors": stale, "models_in_store": models,
|
||||
"store": vectors_path(out_dir),
|
||||
}
|
||||
_atomic_write_text(os.path.join(out_dir, "embed.summary.json"),
|
||||
json.dumps(summary, indent=2, sort_keys=True) + "\n")
|
||||
return summary
|
||||
|
||||
|
||||
# ── search ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def cosine(a, b):
|
||||
dot = na = nb = 0.0
|
||||
for x, y in zip(a, b):
|
||||
dot += x * y
|
||||
na += x * x
|
||||
nb += y * y
|
||||
if na == 0.0 or nb == 0.0:
|
||||
return 0.0
|
||||
return dot / math.sqrt(na * nb)
|
||||
|
||||
|
||||
def search(manifest, out_dir, query, k=5, env=None, endpoint=None):
|
||||
cfg = manifest.embed
|
||||
if cfg is None:
|
||||
raise EmbedError("manifest has no embed section")
|
||||
env = dict(os.environ) if env is None else env
|
||||
ep = endpoint or build_endpoint(manifest, env)
|
||||
qv = ep.embed([query])[0]
|
||||
if len(qv) != cfg.dim:
|
||||
raise EmbedError("query vector dim %d != manifest dim %d" % (len(qv), cfg.dim))
|
||||
conn = open_store(out_dir)
|
||||
try:
|
||||
scored = []
|
||||
for chunk_id_, blob in conn.execute("SELECT chunk_id, vec FROM vectors WHERE model=?", (cfg.model,)):
|
||||
scored.append((cosine(qv, unpack(blob)), chunk_id_))
|
||||
scored.sort(reverse=True)
|
||||
out = []
|
||||
for score, cid in scored[:k]:
|
||||
row = conn.execute("SELECT doc, idx, start, end, text FROM chunks WHERE chunk_id=?", (cid,)).fetchone()
|
||||
if row is None:
|
||||
continue
|
||||
out.append({"score": round(score, 6), "chunk_id": cid, "doc": row[0], "idx": row[1],
|
||||
"start": row[2], "end": row[3], "text": row[4]})
|
||||
return out
|
||||
finally:
|
||||
conn.close()
|
||||
@@ -0,0 +1,768 @@
|
||||
"""bastion agent `extract` stage — text out of the files the walk found.
|
||||
|
||||
Input is the walk map (walk.jsonl); output is a content-addressed text
|
||||
store plus an index:
|
||||
|
||||
<out>/extract/<aa>/<sha256>.txt the extracted text, one file per
|
||||
distinct content hash
|
||||
<out>/extract.jsonl one record per walked file:
|
||||
status, extractor, chars, error
|
||||
<out>/extract.summary.json
|
||||
|
||||
Content-addressing is the cache: a file whose hash already has a text
|
||||
entry is not opened again, on this sweep or the next, and two copies of
|
||||
the same document are extracted once.
|
||||
|
||||
Portability rules are the walker's: stdlib only, Python 3.8 syntax floor,
|
||||
runs on Mac/Windows/Linux. Office formats are zip + XML and need no
|
||||
library. Two OUTSIDE tools are used when — and only when — they are on
|
||||
the machine's PATH: `pdftotext` (poppler) for PDFs and `tesseract` for
|
||||
images. They are sensors, resolved at run time and named in every record
|
||||
they produced; nothing pins them and nothing fails for their absence —
|
||||
the file is recorded `unsupported` with the tool that would handle it,
|
||||
so the host can see exactly what it cannot read.
|
||||
|
||||
The built-in PDF path handles simple PDFs (uncompressed or Flate content
|
||||
streams with simple fonts). It refuses to emit text that does not look
|
||||
like text (CID/Identity-encoded fonts come out as byte salad), recording
|
||||
`unsupported` instead — garbage in the catalog is worse than a gap.
|
||||
"""
|
||||
|
||||
import codecs
|
||||
import csv
|
||||
import datetime as _dt
|
||||
import email
|
||||
import email.policy
|
||||
import html
|
||||
import html.parser
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import unicodedata
|
||||
import zipfile
|
||||
import zlib
|
||||
|
||||
from . import pdftext as _pdftext
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
from .manifest import signing_digest
|
||||
|
||||
STATUS_OK, STATUS_UNSUPPORTED, STATUS_ERROR, STATUS_SKIPPED, STATUS_EMPTY = (
|
||||
"ok", "unsupported", "error", "skipped", "empty")
|
||||
|
||||
# Bump when an extractor's output changes: every cached verdict and text
|
||||
# carries the version it was produced by and is redone under a new one.
|
||||
EXTRACT_VERSION = 2 # 2: the built-in PDF reader (pdftext) reads real PDFs
|
||||
|
||||
ZIP_MEMBER_CAP = 64 * 1024 * 1024 # decompressed bytes we will read from one member
|
||||
TOOL_TIMEOUT = 300 # seconds per external-tool call
|
||||
PDF_TEXT_MIN_RATIO = 0.9 # printable share below which builtin PDF text is refused
|
||||
|
||||
TEXT_MIMES_PREFIX = ("text/",)
|
||||
TEXT_MIMES = {
|
||||
"application/json", "application/x-ndjson", "application/yaml", "application/toml",
|
||||
"application/xml", "application/sql", "image/svg+xml", "application/x-pem-file",
|
||||
}
|
||||
DOCX = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
||||
XLSX = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
|
||||
PPTX = "application/vnd.openxmlformats-officedocument.presentationml.presentation"
|
||||
ODT, ODS, ODP = ("application/vnd.oasis.opendocument.text",
|
||||
"application/vnd.oasis.opendocument.spreadsheet",
|
||||
"application/vnd.oasis.opendocument.presentation")
|
||||
IMAGE_OCR_MIMES = ("image/png", "image/jpeg", "image/tiff", "image/bmp", "image/gif", "image/webp")
|
||||
|
||||
|
||||
# ── text normalisation ──────────────────────────────────────────────────────
|
||||
|
||||
def _normalise(text):
|
||||
text = text.replace("\r\n", "\n").replace("\r", "\n").replace("\x00", "")
|
||||
text = unicodedata.normalize("NFC", text)
|
||||
return text.strip("\n") + ("\n" if text.strip() else "")
|
||||
|
||||
|
||||
def decode_bytes(data):
|
||||
"""Best-effort decode with the encoding recorded. BOMs first, strict
|
||||
UTF-8 next, UTF-16 by NUL pattern, then cp1252 with replacement."""
|
||||
if data.startswith(b"\xef\xbb\xbf"):
|
||||
return data[3:].decode("utf-8", "replace"), "utf-8-sig"
|
||||
if data.startswith(b"\xff\xfe") or data.startswith(b"\xfe\xff"):
|
||||
return data.decode("utf-16", "replace"), "utf-16"
|
||||
# BOM-less UTF-16 is byte-valid UTF-8 (NUL is a code point), so the
|
||||
# NUL-density test has to come before the strict UTF-8 attempt.
|
||||
sample = data[:4096]
|
||||
if len(sample) >= 4 and sample.count(b"\x00") > len(sample) // 4:
|
||||
enc = "utf-16-le" if sample[1:2] == b"\x00" else "utf-16-be"
|
||||
return data.decode(enc, "replace"), enc
|
||||
try:
|
||||
return data.decode("utf-8"), "utf-8"
|
||||
except UnicodeDecodeError:
|
||||
pass
|
||||
return data.decode("cp1252", "replace"), "cp1252"
|
||||
|
||||
|
||||
# ── html ────────────────────────────────────────────────────────────────────
|
||||
|
||||
_BLOCK_TAGS = {"p", "div", "br", "li", "tr", "h1", "h2", "h3", "h4", "h5", "h6",
|
||||
"section", "article", "header", "footer", "table", "ul", "ol",
|
||||
"blockquote", "pre", "hr", "td", "th", "dd", "dt", "title"}
|
||||
|
||||
|
||||
class _HtmlText(html.parser.HTMLParser):
|
||||
def __init__(self):
|
||||
super().__init__(convert_charrefs=True)
|
||||
self.parts = []
|
||||
self._skip = 0
|
||||
|
||||
def handle_starttag(self, tag, attrs):
|
||||
if tag in ("script", "style", "noscript", "template"):
|
||||
self._skip += 1
|
||||
elif tag in _BLOCK_TAGS:
|
||||
self.parts.append("\n")
|
||||
|
||||
def handle_endtag(self, tag):
|
||||
if tag in ("script", "style", "noscript", "template"):
|
||||
self._skip = max(0, self._skip - 1)
|
||||
elif tag in _BLOCK_TAGS:
|
||||
self.parts.append("\n")
|
||||
|
||||
def handle_data(self, data):
|
||||
if not self._skip:
|
||||
self.parts.append(data)
|
||||
|
||||
def text(self):
|
||||
out = "".join(self.parts)
|
||||
out = re.sub(r"[ \t\f\v]+", " ", out)
|
||||
out = re.sub(r" *\n *", "\n", out)
|
||||
return re.sub(r"\n{3,}", "\n\n", out)
|
||||
|
||||
|
||||
def html_to_text(markup):
|
||||
p = _HtmlText()
|
||||
p.feed(markup)
|
||||
p.close()
|
||||
return p.text()
|
||||
|
||||
|
||||
# ── office (zip + xml) ──────────────────────────────────────────────────────
|
||||
#
|
||||
# STREAMED, ON A BUDGET. A 200 KB spreadsheet can hold a gigabyte of XML, and
|
||||
# parsing a member whole (ElementTree.fromstring) built the full tree in
|
||||
# memory — two of them at once across sheets (measured 2.2 GB peak for a
|
||||
# 1.3 MB xlsx). Members are fed to a pull parser in chunks; each wanted element
|
||||
# (a paragraph, a row) is read and then dropped from its parent, so memory is
|
||||
# one element, not one document. A document's decompressed XML, its text and
|
||||
# its wall time are capped; reaching a cap stops the read with the text so far
|
||||
# and says so on the last line. A DTD is refused outright: Office XML never
|
||||
# carries one, and one is how entity expansion is smuggled in.
|
||||
|
||||
OFFICE_DOC_CAP = 256 * 1024 * 1024 # decompressed XML bytes, one document
|
||||
OFFICE_TEXT_CAP = 16 * 1024 * 1024 # characters of text, one document
|
||||
OFFICE_SECONDS = 60.0 # wall time, one document
|
||||
# Parts one paragraph / row may hold while it is being read. A real paragraph
|
||||
# has tens of runs and an Excel row at most 16,384 cells; a crafted one with
|
||||
# millions is refused rather than held in memory.
|
||||
OFFICE_PARTS_PER_ITEM = 200_000
|
||||
_OFFICE_CHUNK = 64 * 1024
|
||||
_DTD = re.compile(rb"<!(?:DOCTYPE|ENTITY)", re.I)
|
||||
_DROPPED = "\x00dropped"
|
||||
_XML_ENCODING = re.compile(r"""(<\?xml[^>]*?encoding\s*=\s*["'])[^"']*(["'])""")
|
||||
|
||||
|
||||
class _OfficeStop(Exception):
|
||||
"""A per-document cap was reached; the text so far stands."""
|
||||
|
||||
|
||||
class _OfficeBudget:
|
||||
def __init__(self):
|
||||
self.bytes = self.chars = 0
|
||||
self.deadline = time.monotonic() + OFFICE_SECONDS
|
||||
self.cut_items = 0 # paragraphs/rows cut at OFFICE_PARTS_PER_ITEM
|
||||
|
||||
def take_bytes(self, n):
|
||||
self.bytes += n
|
||||
if self.bytes > OFFICE_DOC_CAP:
|
||||
raise _OfficeStop("the document expands past %d MB of XML" % (OFFICE_DOC_CAP >> 20))
|
||||
if time.monotonic() > self.deadline:
|
||||
raise _OfficeStop("the document took longer than %ds to read" % OFFICE_SECONDS)
|
||||
|
||||
def take_text(self, out, line):
|
||||
self.chars += len(line) + 1
|
||||
if self.chars > OFFICE_TEXT_CAP:
|
||||
raise _OfficeStop("the document holds more than %d MB of text" % (OFFICE_TEXT_CAP >> 20))
|
||||
out.append(line)
|
||||
|
||||
|
||||
def _zip_elements(zf, name, budget, wanted, skip_within=()):
|
||||
"""Yield each `wanted` element of one member as it closes, then drop it.
|
||||
A wanted tag inside a `skip_within` ancestor is left in place for that
|
||||
ancestor (an ODF table row reads its own cells' paragraphs). Every other
|
||||
finished element that no open wanted element needs is dropped as it
|
||||
closes too, so siblings never pile up — memory is one item, and removal
|
||||
from a parent stays O(1)."""
|
||||
info = zf.getinfo(name)
|
||||
if info.file_size > ZIP_MEMBER_CAP:
|
||||
raise ValueError("zip member %s too large (%d bytes)" % (name, info.file_size))
|
||||
keep = set(wanted) | set(skip_within)
|
||||
parser = ET.XMLPullParser(("start", "end"))
|
||||
stack = []
|
||||
state = {"skipping": 0, "keeping": 0, "parts": 0}
|
||||
|
||||
# Finished elements are cleared at once (their subtree is freed) and
|
||||
# detached from their parents in one batch per fed chunk. The pull parser
|
||||
# reads the whole chunk ahead, so a closing element is NOT its parent's
|
||||
# last child when its event is handled, and remove() scans every kept
|
||||
# sibling from the start: per element that was quadratic.
|
||||
detach = {}
|
||||
|
||||
def drop(el):
|
||||
el.clear()
|
||||
# Until flush() detaches it, a dropped element is still a child of
|
||||
# its parent; renamed so no reader's tag match sees it (a cut row
|
||||
# would otherwise end in empty cells).
|
||||
el.tag = _DROPPED
|
||||
if stack:
|
||||
parent = stack[-1]
|
||||
detach.setdefault(id(parent), (parent, set()))[1].add(id(el))
|
||||
|
||||
def flush():
|
||||
for parent, gone in detach.values():
|
||||
parent[:] = [c for c in parent if id(c) not in gone]
|
||||
detach.clear()
|
||||
|
||||
def drain():
|
||||
for ev, el in parser.read_events():
|
||||
if ev == "start":
|
||||
stack.append(el)
|
||||
if el.tag in skip_within:
|
||||
state["skipping"] += 1
|
||||
if el.tag in keep:
|
||||
state["keeping"] += 1
|
||||
elif state["keeping"]:
|
||||
state["parts"] += 1
|
||||
if state["parts"] == OFFICE_PARTS_PER_ITEM + 1:
|
||||
budget.cut_items += 1
|
||||
continue
|
||||
stack.pop()
|
||||
if el.tag in skip_within:
|
||||
state["skipping"] -= 1
|
||||
if el.tag in keep:
|
||||
state["keeping"] -= 1
|
||||
if not state["keeping"]:
|
||||
state["parts"] = 0
|
||||
if el.tag in wanted and not (state["skipping"] and el.tag not in skip_within):
|
||||
# How many wanted items are still open around this one: a
|
||||
# docx text box is a paragraph inside a paragraph.
|
||||
yield el, state["keeping"]
|
||||
drop(el)
|
||||
elif not state["keeping"] or state["parts"] > OFFICE_PARTS_PER_ITEM:
|
||||
# Past the cap, an item's further parts are dropped as they
|
||||
# close: that one paragraph or row is cut, the rest of the
|
||||
# document still reads.
|
||||
drop(el)
|
||||
|
||||
seen, tail, utf16 = 0, b"", None
|
||||
with zf.open(name) as fh:
|
||||
while True:
|
||||
chunk = fh.read(_OFFICE_CHUNK)
|
||||
if not chunk:
|
||||
break
|
||||
if not seen:
|
||||
# OOXML allows UTF-8 or UTF-16. UTF-16 is decoded to UTF-8
|
||||
# here, BEFORE the DTD check, so a DTD cannot ride past it in
|
||||
# a wide encoding; the declaration is rewritten to match.
|
||||
# Anything else wide (UTF-32, no BOM) is refused.
|
||||
if chunk[:2] in (b"\xff\xfe", b"\xfe\xff"):
|
||||
utf16 = codecs.getincrementaldecoder("utf-16")("strict")
|
||||
elif b"\x00" in chunk[:4]:
|
||||
raise ValueError("zip member %s is not UTF-8 or UTF-16 XML" % name)
|
||||
if utf16 is not None:
|
||||
text = utf16.decode(chunk)
|
||||
if not seen:
|
||||
text = _XML_ENCODING.sub(r"\1UTF-8\2", text, count=1)
|
||||
raw = chunk
|
||||
chunk = text.encode("utf-8")
|
||||
seen += len(raw)
|
||||
else:
|
||||
seen += len(chunk)
|
||||
if seen > ZIP_MEMBER_CAP:
|
||||
raise ValueError("zip member %s too large" % name)
|
||||
if _DTD.search(tail + chunk):
|
||||
raise ValueError("zip member %s declares a DTD; refused" % name)
|
||||
tail = chunk[-16:]
|
||||
budget.take_bytes(len(chunk))
|
||||
parser.feed(chunk)
|
||||
for item in drain():
|
||||
yield item
|
||||
flush()
|
||||
parser.close()
|
||||
for item in drain():
|
||||
yield item
|
||||
flush()
|
||||
|
||||
|
||||
def _office(read):
|
||||
"""Run one reader with a fresh budget; a cap reached keeps the text so
|
||||
far and names the cap on the last line."""
|
||||
budget, out = _OfficeBudget(), []
|
||||
try:
|
||||
read(budget, out)
|
||||
except _OfficeStop as e:
|
||||
out.append("[stopped reading here: %s]" % e)
|
||||
if budget.cut_items:
|
||||
out.append("[%d paragraph(s) or row(s) held more than %d parts; each was cut there]"
|
||||
% (budget.cut_items, OFFICE_PARTS_PER_ITEM))
|
||||
return "\n".join(out)
|
||||
|
||||
|
||||
def _local(tag):
|
||||
return tag.rsplit("}", 1)[-1]
|
||||
|
||||
|
||||
W = "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}"
|
||||
A = "{http://schemas.openxmlformats.org/drawingml/2006/main}"
|
||||
S = "{http://schemas.openxmlformats.org/spreadsheetml/2006/main}"
|
||||
TEXT_NS = "{urn:oasis:names:tc:opendocument:xmlns:text:1.0}"
|
||||
TABLE_NS = "{urn:oasis:names:tc:opendocument:xmlns:table:1.0}"
|
||||
|
||||
|
||||
def _docx_paragraph(p):
|
||||
parts = []
|
||||
for node in p.iter():
|
||||
t = _local(node.tag)
|
||||
if t == "t" and node.text:
|
||||
parts.append(node.text)
|
||||
elif t == "tab":
|
||||
parts.append("\t")
|
||||
elif t in ("br", "cr"):
|
||||
parts.append("\n")
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
def extract_docx(data):
|
||||
def read(budget, out):
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as zf:
|
||||
names = set(zf.namelist())
|
||||
order = ["word/document.xml"] + sorted(
|
||||
n for n in names if re.match(r"word/(header|footer)\d*\.xml$", n))
|
||||
for name in order:
|
||||
if name not in names:
|
||||
continue
|
||||
boxed = [] # text-box paragraphs, read after the paragraph holding them
|
||||
for p, depth in _zip_elements(zf, name, budget, {W + "p"}):
|
||||
if depth:
|
||||
boxed.append(_docx_paragraph(p))
|
||||
continue
|
||||
budget.take_text(out, _docx_paragraph(p))
|
||||
for line in boxed:
|
||||
budget.take_text(out, line)
|
||||
boxed = []
|
||||
for line in boxed:
|
||||
budget.take_text(out, line)
|
||||
out.append("")
|
||||
return _office(read)
|
||||
|
||||
|
||||
def extract_xlsx(data):
|
||||
def read(budget, out):
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as zf:
|
||||
names = set(zf.namelist())
|
||||
shared = []
|
||||
if "xl/sharedStrings.xml" in names:
|
||||
for si, _ in _zip_elements(zf, "xl/sharedStrings.xml", budget, {S + "si"}):
|
||||
text = "".join(t.text or "" for t in si.iter(S + "t"))
|
||||
budget.take_text([], text) # shared strings count against the cap too
|
||||
shared.append(text)
|
||||
sheets = sorted(n for n in names if re.match(r"xl/worksheets/sheet\d+\.xml$", n))
|
||||
for name in sheets:
|
||||
out.append("# " + name.rsplit("/", 1)[-1][:-4])
|
||||
for row, _ in _zip_elements(zf, name, budget, {S + "row"}):
|
||||
cells = []
|
||||
for c in row.iter(S + "c"):
|
||||
v = c.find(S + "v")
|
||||
is_node = c.find(S + "is")
|
||||
val = ""
|
||||
if c.get("t") == "s" and v is not None and v.text is not None:
|
||||
try:
|
||||
val = shared[int(v.text)]
|
||||
except (ValueError, IndexError):
|
||||
val = ""
|
||||
elif c.get("t") == "inlineStr" and is_node is not None:
|
||||
val = "".join(t.text or "" for t in is_node.iter(S + "t"))
|
||||
elif v is not None and v.text is not None:
|
||||
val = v.text
|
||||
cells.append(val)
|
||||
if any(cells):
|
||||
budget.take_text(out, "\t".join(cells))
|
||||
out.append("")
|
||||
return _office(read)
|
||||
|
||||
|
||||
def extract_pptx(data):
|
||||
def read(budget, out):
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as zf:
|
||||
names = zf.namelist()
|
||||
slides = sorted((n for n in names if re.match(r"ppt/slides/slide\d+\.xml$", n)),
|
||||
key=lambda n: int(re.search(r"(\d+)\.xml$", n).group(1)))
|
||||
for name in slides:
|
||||
out.append("# " + name.rsplit("/", 1)[-1][:-4])
|
||||
for p, _ in _zip_elements(zf, name, budget, {A + "p"}):
|
||||
line = "".join(t.text or "" for t in p.iter(A + "t"))
|
||||
if line:
|
||||
budget.take_text(out, line)
|
||||
out.append("")
|
||||
return _office(read)
|
||||
|
||||
|
||||
def extract_odf(data):
|
||||
row_tag, cell_tag = TABLE_NS + "table-row", TABLE_NS + "table-cell"
|
||||
|
||||
def read(budget, out):
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as zf:
|
||||
wanted = {TEXT_NS + "p", TEXT_NS + "h", row_tag}
|
||||
for el, _ in _zip_elements(zf, "content.xml", budget, wanted, skip_within={row_tag}):
|
||||
if el.tag == row_tag:
|
||||
# A row's cells hold paragraphs; the row is one line.
|
||||
cells = ["".join(c.itertext()) for c in el if c.tag == cell_tag]
|
||||
if any(cells):
|
||||
budget.take_text(out, "\t".join(cells))
|
||||
else:
|
||||
budget.take_text(out, "".join(el.itertext()))
|
||||
return _office(read)
|
||||
|
||||
|
||||
# ── email, rtf ──────────────────────────────────────────────────────────────
|
||||
|
||||
def extract_eml(data):
|
||||
msg = email.message_from_bytes(data, policy=email.policy.default)
|
||||
head = []
|
||||
for h in ("From", "To", "Cc", "Date", "Subject"):
|
||||
if msg[h]:
|
||||
head.append("%s: %s" % (h, msg[h]))
|
||||
body = msg.get_body(preferencelist=("plain", "html"))
|
||||
text = ""
|
||||
if body is not None:
|
||||
content = body.get_content()
|
||||
text = html_to_text(content) if body.get_content_type() == "text/html" else content
|
||||
return "\n".join(head) + "\n\n" + text
|
||||
|
||||
|
||||
_RTF_SKIP_GROUPS = {"fonttbl", "colortbl", "stylesheet", "info", "pict", "object",
|
||||
"header", "footer", "xmlnstbl", "themedata", "datastore", "generator"}
|
||||
_RTF_TOKEN = re.compile(r"\\([a-zA-Z]+)(-?\d+)? ?|\\'([0-9a-fA-F]{2})|\\([^a-zA-Z])|([{}])|([^\\{}]+)")
|
||||
|
||||
|
||||
def extract_rtf(data):
|
||||
text = data.decode("latin-1", "replace")
|
||||
out = []
|
||||
stack = []
|
||||
skip_depth = None
|
||||
uc_skip = 0
|
||||
for m in _RTF_TOKEN.finditer(text):
|
||||
word, arg, hexch, esc, brace, plain = m.groups()
|
||||
if brace == "{":
|
||||
stack.append(len(out))
|
||||
continue
|
||||
if brace == "}":
|
||||
if stack:
|
||||
stack.pop()
|
||||
# the skipped group closes when the stack drops below the depth
|
||||
# it was opened at
|
||||
if skip_depth is not None and len(stack) < skip_depth:
|
||||
skip_depth = None
|
||||
continue
|
||||
if skip_depth is not None:
|
||||
continue
|
||||
if word:
|
||||
if word in _RTF_SKIP_GROUPS:
|
||||
skip_depth = len(stack)
|
||||
elif word in ("par", "line", "sect", "page"):
|
||||
out.append("\n")
|
||||
elif word == "tab":
|
||||
out.append("\t")
|
||||
elif word == "u" and arg is not None:
|
||||
cp = int(arg)
|
||||
out.append(chr(cp + 65536 if cp < 0 else cp))
|
||||
uc_skip = 1
|
||||
continue
|
||||
if hexch:
|
||||
if uc_skip:
|
||||
uc_skip -= 1
|
||||
continue
|
||||
out.append(bytes([int(hexch, 16)]).decode("cp1252", "replace"))
|
||||
continue
|
||||
if esc:
|
||||
out.append({"~": "\u00a0", "-": "", "_": "-", "\n": "\n", "\r": ""}.get(esc, esc))
|
||||
continue
|
||||
if plain:
|
||||
if uc_skip:
|
||||
plain = plain[uc_skip:]
|
||||
uc_skip = 0
|
||||
out.append(plain.replace("\r", "").replace("\n", ""))
|
||||
return "".join(out)
|
||||
|
||||
|
||||
# ── pdf (built-in best effort) ──────────────────────────────────────────────
|
||||
|
||||
def extract_pdf_builtin(data):
|
||||
"""Text from the document's pages with the standard library (pdftext):
|
||||
fonts' /ToUnicode maps, compressed streams and object streams, so a
|
||||
browser-printed PDF reads the same as an old literal-string one. Returns
|
||||
None when the result does not look like text (encrypted files, pages that
|
||||
are only pictures). pdftext never raises and reads on a per-document
|
||||
budget, so a hostile file costs bounded time and is at worst no text."""
|
||||
text = _pdftext.extract_text(data)
|
||||
if not text or not text.strip():
|
||||
return None
|
||||
printable = sum(1 for ch in text if ch.isprintable() or ch in "\n\t\f")
|
||||
if printable / max(1, len(text)) < PDF_TEXT_MIN_RATIO:
|
||||
return None
|
||||
return text
|
||||
|
||||
|
||||
# ── external tools (sensors) ────────────────────────────────────────────────
|
||||
|
||||
def tool_path(name):
|
||||
return shutil.which(name)
|
||||
|
||||
|
||||
def run_tool(argv, timeout=TOOL_TIMEOUT):
|
||||
"""Run an external tool; returns stdout bytes or raises RuntimeError."""
|
||||
try:
|
||||
proc = subprocess.run(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
||||
timeout=timeout, check=False)
|
||||
except (OSError, subprocess.TimeoutExpired) as exc:
|
||||
raise RuntimeError("%s: %s" % (argv[0], exc))
|
||||
if proc.returncode != 0:
|
||||
err = proc.stderr.decode("utf-8", "replace").strip().splitlines()
|
||||
raise RuntimeError("%s exited %d: %s" % (argv[0], proc.returncode, err[-1] if err else ""))
|
||||
return proc.stdout
|
||||
|
||||
|
||||
def extract_pdf_tool(path, pdftotext):
|
||||
return decode_bytes(run_tool([pdftotext, "-layout", "-enc", "UTF-8", path, "-"]))[0]
|
||||
|
||||
|
||||
def extract_image_tool(path, tesseract):
|
||||
return decode_bytes(run_tool([tesseract, path, "stdout"]))[0]
|
||||
|
||||
|
||||
# ── dispatch ────────────────────────────────────────────────────────────────
|
||||
|
||||
class Extractor:
|
||||
"""Maps a walk record to (status, extractor_name, text_or_None, error)."""
|
||||
|
||||
def __init__(self, use_external=True):
|
||||
self.pdftotext = tool_path("pdftotext") if use_external else None
|
||||
self.tesseract = tool_path("tesseract") if use_external else None
|
||||
|
||||
def tools(self):
|
||||
return {"pdftotext": self.pdftotext, "tesseract": self.tesseract}
|
||||
|
||||
def _read(self, path):
|
||||
with open(path, "rb") as fh:
|
||||
return fh.read()
|
||||
|
||||
def extract(self, path, mime):
|
||||
try:
|
||||
return self._extract(path, mime)
|
||||
except (OSError, ValueError, zipfile.BadZipFile, ET.ParseError, zlib.error,
|
||||
RuntimeError, UnicodeError, KeyError, csv.Error) as exc:
|
||||
# csv.Error is NOT a ValueError: one >128 KiB CSV field escaped this
|
||||
# tuple and crash-looped the whole sweep
|
||||
return STATUS_ERROR, None, None, "%s: %s" % (type(exc).__name__, exc)
|
||||
|
||||
def _extract(self, path, mime):
|
||||
if mime is None:
|
||||
return STATUS_UNSUPPORTED, None, None, "no mime"
|
||||
if mime == "inode/x-empty":
|
||||
return STATUS_OK, "empty", "", None # recorded as `empty` by the stage
|
||||
if mime.startswith(TEXT_MIMES_PREFIX) or mime in TEXT_MIMES:
|
||||
text, enc = decode_bytes(self._read(path))
|
||||
if mime == "text/html":
|
||||
return STATUS_OK, "html", html_to_text(text), None
|
||||
if mime == "message/rfc822":
|
||||
return STATUS_OK, "eml", extract_eml(self._read(path)), None
|
||||
if mime in ("text/csv", "text/tab-separated-values"):
|
||||
rows = csv.reader(io.StringIO(text), delimiter="\t" if "tab" in mime else ",")
|
||||
return STATUS_OK, "csv", "\n".join("\t".join(r) for r in rows), None
|
||||
return STATUS_OK, "text:" + enc, text, None
|
||||
if mime == "message/rfc822":
|
||||
return STATUS_OK, "eml", extract_eml(self._read(path)), None
|
||||
if mime == "application/rtf":
|
||||
return STATUS_OK, "rtf", extract_rtf(self._read(path)), None
|
||||
if mime == DOCX:
|
||||
return STATUS_OK, "docx", extract_docx(self._read(path)), None
|
||||
if mime == XLSX:
|
||||
return STATUS_OK, "xlsx", extract_xlsx(self._read(path)), None
|
||||
if mime == PPTX:
|
||||
return STATUS_OK, "pptx", extract_pptx(self._read(path)), None
|
||||
if mime in (ODT, ODS, ODP):
|
||||
return STATUS_OK, "odf", extract_odf(self._read(path)), None
|
||||
if mime == "application/pdf":
|
||||
if self.pdftotext:
|
||||
return STATUS_OK, "pdftotext", extract_pdf_tool(path, self.pdftotext), None
|
||||
text = extract_pdf_builtin(self._read(path))
|
||||
if text is None:
|
||||
return (STATUS_UNSUPPORTED, "pdf-builtin", None,
|
||||
"needs pdftotext (poppler) on this machine")
|
||||
return STATUS_OK, "pdf-builtin", text, None
|
||||
if mime in IMAGE_OCR_MIMES:
|
||||
if self.tesseract:
|
||||
return STATUS_OK, "tesseract", extract_image_tool(path, self.tesseract), None
|
||||
return STATUS_UNSUPPORTED, None, None, "needs tesseract on this machine"
|
||||
return STATUS_UNSUPPORTED, None, None, "no extractor for %s" % mime
|
||||
|
||||
|
||||
# ── the stage ───────────────────────────────────────────────────────────────
|
||||
|
||||
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, "extract", sha[:2], sha + ".txt")
|
||||
|
||||
|
||||
def load_walk(out_dir):
|
||||
path = os.path.join(out_dir, "walk.jsonl")
|
||||
if not os.path.exists(path):
|
||||
raise FileNotFoundError("no walk map at %s — run the walk stage first" % path)
|
||||
summary_path = os.path.join(out_dir, "walk.summary.json")
|
||||
with open(summary_path, encoding="utf-8") as fh:
|
||||
summary = json.load(fh)
|
||||
roots = {r["index"]: r["resolved"] for r in summary["roots"]}
|
||||
records = []
|
||||
with open(path, encoding="utf-8") as fh:
|
||||
for line in fh:
|
||||
rec = json.loads(line)
|
||||
if rec.get("kind") == "file":
|
||||
records.append(rec)
|
||||
return roots, records, summary
|
||||
|
||||
|
||||
def load_index(out_dir):
|
||||
"""Previous extract.jsonl: sha256 -> record, so unchanged content is
|
||||
not re-extracted and prior unsupported/error verdicts are re-tried
|
||||
only when the machine's tools changed."""
|
||||
idx = {}
|
||||
path = os.path.join(out_dir, "extract.jsonl")
|
||||
try:
|
||||
with open(path, 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 extract(manifest, out_dir, use_external=True, now=None):
|
||||
roots, files, walk_summary = load_walk(out_dir)
|
||||
ex = Extractor(use_external=use_external)
|
||||
tools = ex.tools()
|
||||
prev = load_index(out_dir)
|
||||
counts = {k: 0 for k in ("files", "ok", "reused", "unsupported", "error", "skipped", "empty")}
|
||||
chars = 0
|
||||
index = []
|
||||
done_sha = {}
|
||||
|
||||
for rec in files:
|
||||
counts["files"] += 1
|
||||
sha = rec.get("sha256")
|
||||
base = {"root": rec["root"], "path": rec["path"], "sha256": sha, "mime": rec.get("mime"),
|
||||
"v": EXTRACT_VERSION}
|
||||
if rec.get("skip") or rec.get("error") or not sha:
|
||||
counts["skipped"] += 1
|
||||
index.append(dict(base, status=STATUS_SKIPPED,
|
||||
error=rec.get("skip") or rec.get("error") or "no hash"))
|
||||
continue
|
||||
if sha in done_sha: # duplicate content on this sweep
|
||||
dup = done_sha[sha]
|
||||
index.append(dict(dup, root=rec["root"], path=rec["path"]))
|
||||
counts["reused" if dup.get("reused") else dup["status"]] += 1
|
||||
if dup.get("reused") is None and dup["status"] == STATUS_OK:
|
||||
chars += dup.get("chars", 0)
|
||||
continue
|
||||
spath = store_path(out_dir, sha)
|
||||
old = prev.get(sha)
|
||||
if old and old.get("v") != EXTRACT_VERSION:
|
||||
old = None # produced by an older extractor: redo
|
||||
if old and old.get("status") in (STATUS_OK, STATUS_EMPTY) and (
|
||||
old["status"] == STATUS_EMPTY or os.path.exists(spath)):
|
||||
out = dict(base, status=old["status"], extractor=old.get("extractor"),
|
||||
chars=old.get("chars", 0), reused=True)
|
||||
counts["reused"] += 1
|
||||
chars += out["chars"]
|
||||
index.append(out)
|
||||
done_sha[sha] = out
|
||||
continue
|
||||
if old and old.get("status") == STATUS_UNSUPPORTED and old.get("tools") == tools:
|
||||
out = dict(base, status=STATUS_UNSUPPORTED, extractor=old.get("extractor"),
|
||||
error=old.get("error"), tools=tools, reused=True)
|
||||
counts["unsupported"] += 1
|
||||
index.append(out)
|
||||
done_sha[sha] = out
|
||||
continue
|
||||
|
||||
full = os.path.join(roots[rec["root"]], *rec["path"].split("/"))
|
||||
status, name, text, err = ex.extract(full, rec.get("mime"))
|
||||
out = dict(base, status=status, extractor=name)
|
||||
if status == STATUS_OK:
|
||||
text = _normalise(text or "")
|
||||
if not text.strip():
|
||||
out["status"] = STATUS_EMPTY
|
||||
out["chars"] = 0
|
||||
else:
|
||||
_atomic_write_text(spath, text)
|
||||
out["chars"] = len(text)
|
||||
chars += len(text)
|
||||
else:
|
||||
out["error"] = err
|
||||
if status == STATUS_UNSUPPORTED:
|
||||
out["tools"] = tools
|
||||
counts[out["status"]] += 1
|
||||
index.append(out)
|
||||
done_sha[sha] = out
|
||||
|
||||
index.sort(key=lambda r: (r["root"], r["path"]))
|
||||
idx_path = os.path.join(out_dir, "extract.jsonl")
|
||||
_atomic_write_text(idx_path, "".join(
|
||||
json.dumps(r, ensure_ascii=False, sort_keys=True) + "\n" for r in index))
|
||||
summary = {
|
||||
"freeagent": 1,
|
||||
"stage": "extract",
|
||||
"extractor_version": EXTRACT_VERSION,
|
||||
"manifest_id": manifest.id,
|
||||
"manifest_digest": "sha256:" + signing_digest(manifest.raw),
|
||||
"walk_digest": walk_summary.get("manifest_digest"),
|
||||
"extracted_at": now or _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
|
||||
"tools": tools,
|
||||
"counts": counts,
|
||||
"chars": chars,
|
||||
"entries": len(index),
|
||||
}
|
||||
_atomic_write_text(os.path.join(out_dir, "extract.summary.json"),
|
||||
json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True) + "\n")
|
||||
return summary
|
||||
@@ -0,0 +1,32 @@
|
||||
# bastion agent installer — Windows (PowerShell 5+). No admin, no downloads.
|
||||
# powershell -ExecutionPolicy Bypass -File install.ps1 dist\bastion-agent-0.1.0.pyz
|
||||
# Puts bastion-agent.cmd in %LOCALAPPDATA%\bastion-agent and runs the preflight.
|
||||
param([Parameter(Mandatory = $true)][string]$Pyz)
|
||||
$ErrorActionPreference = "Stop"
|
||||
$py = Get-Command python -ErrorAction SilentlyContinue
|
||||
if (-not $py) { $py = Get-Command py -ErrorAction SilentlyContinue }
|
||||
if (-not $py) { Write-Error "Python 3.8+ is required (python.org or the Store); none on PATH"; exit 1 }
|
||||
& $py.Source -c "import sys; sys.exit(0 if sys.version_info >= (3, 8) else 1)"
|
||||
if ($LASTEXITCODE -ne 0) { Write-Error "Python is older than 3.8"; exit 1 }
|
||||
if (Test-Path "$Pyz.sha256") {
|
||||
$want = (Get-Content "$Pyz.sha256").Split(" ")[0].ToLower()
|
||||
$have = (Get-FileHash -Algorithm SHA256 $Pyz).Hash.ToLower()
|
||||
if ($want -ne $have) { Write-Error "checksum MISMATCH - refusing to install"; exit 1 }
|
||||
Write-Host "checksum ok"
|
||||
}
|
||||
$dir = Join-Path $env:LOCALAPPDATA "bastion-agent"
|
||||
New-Item -ItemType Directory -Force -Path $dir | Out-Null
|
||||
Copy-Item $Pyz (Join-Path $dir "bastion-agent.pyz") -Force
|
||||
$cmd = Join-Path $dir "bastion-agent.cmd"
|
||||
"@echo off`r`n`"$($py.Source)`" `"$dir\bastion-agent.pyz`" %*`r`n" | Set-Content -Path $cmd -Encoding ASCII
|
||||
Write-Host "installed $cmd"
|
||||
$userPath = [Environment]::GetEnvironmentVariable("Path", "User")
|
||||
if ($userPath -notlike "*$dir*") {
|
||||
[Environment]::SetEnvironmentVariable("Path", "$userPath;$dir", "User")
|
||||
Write-Host "added $dir to your user PATH (open a new terminal)"
|
||||
}
|
||||
& $cmd --version
|
||||
& $cmd doctor
|
||||
Write-Host ""
|
||||
Write-Host "next: put your manifest at %APPDATA%\bastion-agent\manifest.yaml, set FREEAGENT_OWNER_RECIPIENT"
|
||||
Write-Host " (age-keygen), then: bastion-agent sweep %APPDATA%\bastion-agent\manifest.yaml --once"
|
||||
@@ -0,0 +1,43 @@
|
||||
#!/usr/bin/env sh
|
||||
# bastion agent installer — Linux and macOS. POSIX sh, no root, no downloads.
|
||||
# sh install.sh dist/bastion-agent-0.1.7.pyz
|
||||
# Puts `bastion-agent` on your PATH (~/.local/bin) and runs the preflight.
|
||||
#
|
||||
# This is the HAND install. The managed path — the bastion deploying the agent
|
||||
# from a grant — does not use this script.
|
||||
set -eu
|
||||
PYZ="${1:?usage: install.sh <bastion-agent-X.Y.Z.pyz>}"
|
||||
# FREEAGENT_BIN_DIR keeps its name deliberately: the six FREEAGENT_* variables
|
||||
# are the contract the bastion sets when it spawns the agent, and renaming them
|
||||
# here would break every bastion that still sets the old ones. They move in
|
||||
# their own compatibility window, not this one.
|
||||
BIN="${BASTION_AGENT_BIN_DIR:-${FREEAGENT_BIN_DIR:-$HOME/.local/bin}}"
|
||||
PY="$(command -v python3 || true)"
|
||||
[ -n "$PY" ] || { echo "python3 is required (3.8 or newer); none on PATH" >&2; exit 1; }
|
||||
"$PY" -c 'import sys; sys.exit(0 if sys.version_info >= (3, 8) else 1)' \
|
||||
|| { echo "python3 is older than 3.8: $("$PY" --version)" >&2; exit 1; }
|
||||
if [ -f "${PYZ}.sha256" ]; then
|
||||
(cd "$(dirname "$PYZ")" && { command -v sha256sum >/dev/null && sha256sum -c "$(basename "$PYZ").sha256" \
|
||||
|| shasum -a 256 -c "$(basename "$PYZ").sha256"; }) >/dev/null \
|
||||
&& echo "checksum ok" || { echo "checksum MISMATCH — refusing to install" >&2; exit 1; }
|
||||
fi
|
||||
mkdir -p "$BIN"
|
||||
cp "$PYZ" "$BIN/bastion-agent.pyz"
|
||||
printf '#!/usr/bin/env sh\nexec python3 "%s/bastion-agent.pyz" "$@"\n' "$BIN" > "$BIN/bastion-agent"
|
||||
chmod 755 "$BIN/bastion-agent" "$BIN/bastion-agent.pyz"
|
||||
# An earlier hand install left a `freeagent` wrapper pointing at a file this
|
||||
# install replaces. Repoint it rather than leaving it broken on someone's PATH;
|
||||
# it is the same program under its old name.
|
||||
if [ -f "$BIN/freeagent" ]; then
|
||||
printf '#!/usr/bin/env sh\n# renamed to bastion-agent; this wrapper kept so an old habit still works\nexec python3 "%s/bastion-agent.pyz" "$@"\n' "$BIN" > "$BIN/freeagent"
|
||||
chmod 755 "$BIN/freeagent"
|
||||
rm -f "$BIN/freeagent.pyz" "$BIN/bastion_agent.pyz"
|
||||
echo "note: kept $BIN/freeagent working — it now runs bastion-agent"
|
||||
fi
|
||||
echo "installed $BIN/bastion-agent ($("$BIN/bastion-agent" --version))"
|
||||
case ":$PATH:" in *":$BIN:"*) ;; *) echo "note: add $BIN to PATH";; esac
|
||||
echo
|
||||
"$BIN/bastion-agent" doctor || true
|
||||
echo
|
||||
echo "next: put your manifest at ~/.config/bastion-agent/manifest.yaml, set FREEAGENT_OWNER_RECIPIENT"
|
||||
echo " (age-keygen), then: bastion-agent sweep ~/.config/bastion-agent/manifest.yaml --once"
|
||||
@@ -0,0 +1,576 @@
|
||||
"""bastion agent job-manifest loader — the only way tools read a manifest.
|
||||
|
||||
Validation collects every problem before raising, so one run reports the whole
|
||||
file. See docs/manifest.md for the format contract. The loader fixes
|
||||
the *shape* and the signing payload; cryptographic verification of the
|
||||
signature belongs to the runtime (run.py, x509gate.py), never here.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import datetime as _dt
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
import yaml
|
||||
|
||||
SCHEMA_VERSION = 1
|
||||
|
||||
# The agent lives WITH its folder: <root>/.agent/ holds the manifest and every
|
||||
# prepared output (catalog, sanitized index, vectors, run log). The walker
|
||||
# always excludes it, so an agent never indexes its own output. Secrets (the
|
||||
# PII vault, the identity key) are NOT kept there — a Documents folder is
|
||||
# often cloud-synced — see FREEAGENT_VAULT_DIR and the bastion's state dir.
|
||||
AGENT_DIR = ".agent"
|
||||
|
||||
SLUG_RE = re.compile(r"^[a-z0-9-]+$")
|
||||
DATE_RE = re.compile(r"^\d{4}-\d{2}-\d{2}$")
|
||||
# a root must mean the same place wherever the manifest lands: POSIX absolute,
|
||||
# home-relative, Windows drive-absolute (C:\ or C:/), or UNC (\\server\share)
|
||||
ROOT_RE = re.compile(r"^(/|~$|~/|[A-Za-z]:[\\/]|\\\\[^\\/]+[\\/][^\\/]+)")
|
||||
|
||||
TOP_KEYS_REQUIRED = ("freeagent", "id", "owner", "scope", "mask", "stages", "sweep", "sink")
|
||||
TOP_KEYS_OPTIONAL = ("sanitize", "embed", "summarize", "grants", "signature")
|
||||
|
||||
MASK_OPS = ("read", "search", "index", "summarize")
|
||||
MASK_MANDATORY = "read"
|
||||
|
||||
# canonical stage order; a manifest's stages must be a subset in this order
|
||||
STAGES = ("walk", "extract", "sanitize", "chunk", "embed", "catalog", "seal")
|
||||
# chunk and catalog read the SANITIZED store, so they need sanitize, not
|
||||
# merely extract (the runtime failed closed, the schema should
|
||||
# have refused first)
|
||||
STAGE_NEEDS_STAGE = {
|
||||
"extract": "walk",
|
||||
"sanitize": "extract",
|
||||
"chunk": "sanitize",
|
||||
"embed": "chunk",
|
||||
"catalog": "sanitize",
|
||||
"seal": "catalog",
|
||||
}
|
||||
STAGE_NEEDS_OP = {"embed": "index", "catalog": "index"}
|
||||
|
||||
SANITIZE_MODES = ("sterilize", "flag")
|
||||
SUMMARIZE_PER = ("file", "directory")
|
||||
SWEEPS = ("15m", "hourly", "daily", "manual") # same vocabulary as the platform context store
|
||||
SINK_CATALOG = ("local", "host")
|
||||
SINK_VECTORS = ("local",)
|
||||
SINK_PLAINTEXT = "never" # a constant, not a setting
|
||||
GRANT_OPS = ("search", "read_range")
|
||||
GRANT_OP_NEEDS_MASK = {"search": "search", "read_range": "read"}
|
||||
SIGNATURE_ALGS = ("ecdsa-p256-sha256",)
|
||||
|
||||
DEFAULT_MAX_FILE_BYTES = 256 * 1024 * 1024
|
||||
|
||||
|
||||
class ManifestError(Exception):
|
||||
"""Raised with every validation error joined into one message."""
|
||||
|
||||
def __init__(self, errors):
|
||||
self.errors = list(errors)
|
||||
super().__init__(
|
||||
"manifest validation failed:\n" + "\n".join(f" - {e}" for e in self.errors)
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Scope:
|
||||
roots: tuple
|
||||
exclude: tuple = ()
|
||||
follow_symlinks: bool = False
|
||||
max_file_bytes: int = DEFAULT_MAX_FILE_BYTES
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Sanitize:
|
||||
mode: str
|
||||
scope: str
|
||||
verify: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Embed:
|
||||
model: str
|
||||
dim: int
|
||||
chunk_size_tokens: int
|
||||
chunk_overlap_tokens: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Summarize:
|
||||
model: str
|
||||
per: tuple
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Sink:
|
||||
catalog: str
|
||||
vectors: str = "local"
|
||||
plaintext: str = SINK_PLAINTEXT
|
||||
seal_to: tuple = ()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Grant:
|
||||
host: str
|
||||
ops: tuple
|
||||
until: _dt.date
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Signature:
|
||||
alg: str
|
||||
cert: str
|
||||
value: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Manifest:
|
||||
id: str
|
||||
owner_host: str
|
||||
scope: Scope
|
||||
mask: tuple
|
||||
stages: tuple
|
||||
sweep: str
|
||||
sink: Sink
|
||||
sanitize: Sanitize = None
|
||||
embed: Embed = None
|
||||
summarize: Summarize = None
|
||||
grants: tuple = ()
|
||||
signature: Signature = None
|
||||
raw: dict = None # the parsed document, for signing_payload()
|
||||
|
||||
@property
|
||||
def signed(self):
|
||||
return self.signature is not None
|
||||
|
||||
def has_stage(self, name):
|
||||
return name in self.stages
|
||||
|
||||
|
||||
# ── signing payload ─────────────────────────────────────────────────────────
|
||||
|
||||
def signing_payload(doc):
|
||||
"""Bytes the signature covers: the document minus `signature`, as
|
||||
canonical JSON (sorted keys, no whitespace, UTF-8). Signer and runtime
|
||||
must agree byte-for-byte, so this is the single definition."""
|
||||
body = {k: v for k, v in doc.items() if k != "signature"}
|
||||
return json.dumps(
|
||||
body, sort_keys=True, separators=(",", ":"), ensure_ascii=False,
|
||||
default=_canonical_scalar,
|
||||
).encode("utf-8")
|
||||
|
||||
|
||||
def _canonical_scalar(obj):
|
||||
# YAML parses a bare 2026-12-31 as a date; a quoted one stays a string.
|
||||
# Both must sign to the same bytes, so dates canonicalise to ISO text.
|
||||
if isinstance(obj, _dt.date) and not isinstance(obj, _dt.datetime):
|
||||
return obj.isoformat()
|
||||
raise TypeError(f"signing payload cannot contain {type(obj).__name__}")
|
||||
|
||||
|
||||
def signing_digest(doc):
|
||||
return hashlib.sha256(signing_payload(doc)).hexdigest()
|
||||
|
||||
|
||||
# ── helpers ─────────────────────────────────────────────────────────────────
|
||||
|
||||
def _is_int(v):
|
||||
return isinstance(v, int) and not isinstance(v, bool)
|
||||
|
||||
|
||||
def _is_str(v, lo=1, hi=None):
|
||||
return isinstance(v, str) and len(v) >= lo and (hi is None or len(v) <= hi)
|
||||
|
||||
|
||||
def _check_keys(where, raw, required, optional, errors):
|
||||
"""Unknown/missing key check. Returns True when the key set is clean."""
|
||||
if not isinstance(raw, dict):
|
||||
errors.append(f"{where}: must be a map")
|
||||
return False
|
||||
unknown = set(raw) - set(required) - set(optional)
|
||||
for key in sorted(unknown, key=str):
|
||||
errors.append(f"{where}: unknown key {key!r}")
|
||||
missing = [k for k in required if k not in raw]
|
||||
if missing:
|
||||
errors.append(f"{where}: missing required key(s): {', '.join(missing)}")
|
||||
return not unknown and not missing
|
||||
|
||||
|
||||
def _str_list(where, key, val, errors, allowed=None, nonempty=True, hi=None):
|
||||
"""Validate a list of strings (optionally from a closed vocabulary, no
|
||||
duplicates). Returns the tuple, or None when invalid."""
|
||||
if not isinstance(val, list) or not all(isinstance(s, str) for s in val):
|
||||
errors.append(f"{where}: {key} must be a list of strings")
|
||||
return None
|
||||
if nonempty and not val:
|
||||
errors.append(f"{where}: {key} must not be empty")
|
||||
return None
|
||||
ok = True
|
||||
if allowed is not None:
|
||||
bad = [s for s in val if s not in allowed]
|
||||
if bad:
|
||||
errors.append(f"{where}: {key} has value(s) not in {allowed}: {bad}")
|
||||
ok = False
|
||||
if hi is not None:
|
||||
long = [s for s in val if not _is_str(s, 1, hi)]
|
||||
if long:
|
||||
errors.append(f"{where}: {key} entries must be 1-{hi} chars: {long}")
|
||||
ok = False
|
||||
dupes = sorted({s for s in val if val.count(s) > 1})
|
||||
if dupes:
|
||||
errors.append(f"{where}: {key} has duplicate(s): {dupes}")
|
||||
ok = False
|
||||
return tuple(val) if ok else None
|
||||
|
||||
|
||||
# ── sections ────────────────────────────────────────────────────────────────
|
||||
|
||||
def _validate_scope(raw, errors):
|
||||
where = "scope"
|
||||
if not _check_keys(where, raw, ("roots",), ("exclude", "follow_symlinks", "max_file_bytes"), errors):
|
||||
return None
|
||||
ok = True
|
||||
roots = _str_list(where, "roots", raw.get("roots"), errors)
|
||||
if roots is not None:
|
||||
bad = [r for r in roots if not ROOT_RE.match(r)]
|
||||
if bad:
|
||||
errors.append(
|
||||
f"{where}: roots must be absolute (/, C:\\, \\\\server\\share) or start with '~/': {bad}"
|
||||
)
|
||||
ok = False
|
||||
else:
|
||||
ok = False
|
||||
exclude = _str_list(where, "exclude", raw.get("exclude", []), errors, nonempty=False)
|
||||
if exclude is None:
|
||||
ok = False
|
||||
follow = raw.get("follow_symlinks", False)
|
||||
if not isinstance(follow, bool):
|
||||
errors.append(f"{where}: follow_symlinks must be a boolean")
|
||||
ok = False
|
||||
max_bytes = raw.get("max_file_bytes", DEFAULT_MAX_FILE_BYTES)
|
||||
if not _is_int(max_bytes) or max_bytes < 1:
|
||||
errors.append(f"{where}: max_file_bytes must be a positive integer")
|
||||
ok = False
|
||||
if not ok:
|
||||
return None
|
||||
return Scope(roots=roots, exclude=exclude, follow_symlinks=follow, max_file_bytes=max_bytes)
|
||||
|
||||
|
||||
def _validate_stages(raw, errors):
|
||||
stages = _str_list("stages", "stages", raw, errors, allowed=STAGES)
|
||||
if stages is None:
|
||||
return None
|
||||
ok = True
|
||||
if stages[0] != "walk":
|
||||
errors.append("stages: 'walk' must be present and first")
|
||||
ok = False
|
||||
order = [STAGES.index(s) for s in stages]
|
||||
if order != sorted(order):
|
||||
errors.append(f"stages: must follow the canonical order {STAGES}")
|
||||
ok = False
|
||||
for s in stages:
|
||||
need = STAGE_NEEDS_STAGE.get(s)
|
||||
if need and need not in stages:
|
||||
errors.append(f"stages: {s!r} requires {need!r}")
|
||||
ok = False
|
||||
return stages if ok else None
|
||||
|
||||
|
||||
def _validate_sanitize(raw, manifest_id, errors):
|
||||
where = "sanitize"
|
||||
if not _check_keys(where, raw, ("mode",), ("scope", "verify"), errors):
|
||||
return None
|
||||
ok = True
|
||||
mode = raw.get("mode")
|
||||
if mode not in SANITIZE_MODES:
|
||||
errors.append(f"{where}: mode {mode!r} is not one of {SANITIZE_MODES}")
|
||||
ok = False
|
||||
scope = raw.get("scope", manifest_id)
|
||||
if not _is_str(scope, 1, 63):
|
||||
errors.append(f"{where}: scope must be a string of 1-63 chars")
|
||||
ok = False
|
||||
verify = raw.get("verify", False)
|
||||
if not isinstance(verify, bool):
|
||||
errors.append(f"{where}: verify must be a boolean")
|
||||
ok = False
|
||||
elif verify and mode == "flag":
|
||||
errors.append(f"{where}: verify is only valid with mode 'sterilize' (flag is rules-only)")
|
||||
ok = False
|
||||
if not ok:
|
||||
return None
|
||||
return Sanitize(mode=mode, scope=scope, verify=verify)
|
||||
|
||||
|
||||
def _validate_embed(raw, errors):
|
||||
where = "embed"
|
||||
if not _check_keys(where, raw, ("model", "dim", "chunk"), (), errors):
|
||||
return None
|
||||
ok = True
|
||||
model = raw.get("model")
|
||||
if not _is_str(model, 1, 128):
|
||||
errors.append(f"{where}: model must be a string of 1-128 chars")
|
||||
ok = False
|
||||
dim = raw.get("dim")
|
||||
if not _is_int(dim) or not 1 <= dim <= 8192:
|
||||
errors.append(f"{where}: dim must be an integer 1-8192")
|
||||
ok = False
|
||||
chunk = raw.get("chunk")
|
||||
size = overlap = None
|
||||
if _check_keys(f"{where}.chunk", chunk, ("size_tokens", "overlap_tokens"), (), errors):
|
||||
size = chunk.get("size_tokens")
|
||||
overlap = chunk.get("overlap_tokens")
|
||||
if not _is_int(size) or not 1 <= size <= 8192:
|
||||
errors.append(f"{where}.chunk: size_tokens must be an integer 1-8192")
|
||||
ok = False
|
||||
if not _is_int(overlap) or overlap < 0:
|
||||
errors.append(f"{where}.chunk: overlap_tokens must be a non-negative integer")
|
||||
ok = False
|
||||
elif _is_int(size) and overlap >= size:
|
||||
errors.append(f"{where}.chunk: overlap_tokens must be smaller than size_tokens")
|
||||
ok = False
|
||||
else:
|
||||
ok = False
|
||||
if not ok:
|
||||
return None
|
||||
return Embed(model=model, dim=dim, chunk_size_tokens=size, chunk_overlap_tokens=overlap)
|
||||
|
||||
|
||||
def _validate_summarize(raw, errors):
|
||||
where = "summarize"
|
||||
if not _check_keys(where, raw, ("model", "per"), (), errors):
|
||||
return None
|
||||
model = raw.get("model")
|
||||
ok = True
|
||||
if not _is_str(model, 1, 128):
|
||||
errors.append(f"{where}: model must be a string of 1-128 chars")
|
||||
ok = False
|
||||
per = _str_list(where, "per", raw.get("per"), errors, allowed=SUMMARIZE_PER)
|
||||
if per is None:
|
||||
ok = False
|
||||
return Summarize(model=model, per=per) if ok else None
|
||||
|
||||
|
||||
def _validate_sink(raw, errors):
|
||||
where = "sink"
|
||||
if not _check_keys(where, raw, ("catalog", "vectors", "plaintext"), ("seal_to",), errors):
|
||||
return None
|
||||
ok = True
|
||||
catalog = raw.get("catalog")
|
||||
if catalog not in SINK_CATALOG:
|
||||
errors.append(f"{where}: catalog {catalog!r} is not one of {SINK_CATALOG}")
|
||||
ok = False
|
||||
vectors = raw.get("vectors")
|
||||
if vectors not in SINK_VECTORS:
|
||||
errors.append(f"{where}: vectors {vectors!r} is not one of {SINK_VECTORS} (v1: local only)")
|
||||
ok = False
|
||||
plaintext = raw.get("plaintext")
|
||||
if plaintext != SINK_PLAINTEXT:
|
||||
errors.append(
|
||||
f"{where}: plaintext must be {SINK_PLAINTEXT!r} — it is an invariant, not a setting"
|
||||
)
|
||||
ok = False
|
||||
seal_to = _str_list(where, "seal_to", raw.get("seal_to", []), errors, nonempty=False, hi=128)
|
||||
if seal_to is None:
|
||||
ok = False
|
||||
if not ok:
|
||||
return None
|
||||
return Sink(catalog=catalog, vectors=vectors, plaintext=plaintext, seal_to=seal_to)
|
||||
|
||||
|
||||
def _validate_grants(raw, owner_host, mask, errors):
|
||||
if raw is None:
|
||||
raw = []
|
||||
if not isinstance(raw, list):
|
||||
errors.append("grants: must be a list")
|
||||
return None
|
||||
grants = []
|
||||
seen = set()
|
||||
ok = True
|
||||
for idx, g in enumerate(raw):
|
||||
where = f"grants[{idx}]"
|
||||
if not _check_keys(where, g, ("host", "ops", "until"), (), errors):
|
||||
ok = False
|
||||
continue
|
||||
host = g.get("host")
|
||||
if not _is_str(host, 1, 63) or not SLUG_RE.match(host):
|
||||
errors.append(f"{where}: host must be a lowercase slug [a-z0-9-]+ of 1-63 chars")
|
||||
ok = False
|
||||
elif host == owner_host:
|
||||
errors.append(f"{where}: host {host!r} is the owner — a grant to yourself is a mistake")
|
||||
ok = False
|
||||
elif host in seen:
|
||||
errors.append(f"{where}: duplicate grant for host {host!r}")
|
||||
ok = False
|
||||
seen.add(host)
|
||||
ops = _str_list(where, "ops", g.get("ops"), errors, allowed=GRANT_OPS)
|
||||
if ops is None:
|
||||
ok = False
|
||||
elif mask is not None:
|
||||
for op in ops:
|
||||
need = GRANT_OP_NEEDS_MASK[op]
|
||||
if need not in mask:
|
||||
errors.append(f"{where}: op {op!r} requires {need!r} in mask")
|
||||
ok = False
|
||||
until = g.get("until")
|
||||
until_date = None
|
||||
if isinstance(until, _dt.date) and not isinstance(until, _dt.datetime):
|
||||
until_date = until # PyYAML parses bare YYYY-MM-DD as a date
|
||||
elif isinstance(until, str) and DATE_RE.match(until):
|
||||
try:
|
||||
until_date = _dt.date.fromisoformat(until)
|
||||
except ValueError:
|
||||
pass
|
||||
if until_date is None:
|
||||
errors.append(f"{where}: until {until!r} is not a real YYYY-MM-DD date")
|
||||
ok = False
|
||||
if ok and until_date is not None and ops is not None:
|
||||
grants.append(Grant(host=host, ops=ops, until=until_date))
|
||||
return tuple(grants) if ok else None
|
||||
|
||||
|
||||
def _validate_signature(raw, errors):
|
||||
where = "signature"
|
||||
if not _check_keys(where, raw, ("alg", "cert", "value"), (), errors):
|
||||
return None
|
||||
ok = True
|
||||
alg = raw.get("alg")
|
||||
if alg not in SIGNATURE_ALGS:
|
||||
errors.append(f"{where}: alg {alg!r} is not one of {SIGNATURE_ALGS}")
|
||||
ok = False
|
||||
cert = raw.get("cert")
|
||||
if not _is_str(cert, 1, 256):
|
||||
errors.append(f"{where}: cert must be a path string of 1-256 chars")
|
||||
ok = False
|
||||
value = raw.get("value")
|
||||
if not _is_str(value, 1):
|
||||
errors.append(f"{where}: value must be a non-empty base64 string")
|
||||
ok = False
|
||||
else:
|
||||
try:
|
||||
base64.b64decode(value, validate=True)
|
||||
except Exception:
|
||||
errors.append(f"{where}: value is not valid base64")
|
||||
ok = False
|
||||
return Signature(alg=alg, cert=cert, value=value) if ok else None
|
||||
|
||||
|
||||
# ── entry point ─────────────────────────────────────────────────────────────
|
||||
|
||||
def validate_document(doc, require_signed=False):
|
||||
"""Validate a parsed manifest document. Raises ManifestError listing all
|
||||
problems; returns a Manifest otherwise."""
|
||||
errors = []
|
||||
if not isinstance(doc, dict):
|
||||
raise ManifestError(["top level must be a map"])
|
||||
_check_keys("top level", doc, TOP_KEYS_REQUIRED, TOP_KEYS_OPTIONAL, errors)
|
||||
|
||||
ver = doc.get("freeagent")
|
||||
if "freeagent" in doc and ver != SCHEMA_VERSION:
|
||||
errors.append(f"freeagent: schema version {ver!r} is not {SCHEMA_VERSION}")
|
||||
|
||||
manifest_id = doc.get("id")
|
||||
if "id" in doc and (not _is_str(manifest_id, 1, 63) or not SLUG_RE.match(manifest_id)):
|
||||
errors.append(f"id {manifest_id!r} is not a lowercase slug [a-z0-9-]+ of 1-63 chars")
|
||||
manifest_id = None
|
||||
|
||||
owner_host = None
|
||||
if "owner" in doc and _check_keys("owner", doc.get("owner"), ("host",), (), errors):
|
||||
owner_host = doc["owner"].get("host")
|
||||
if not _is_str(owner_host, 1, 63) or not SLUG_RE.match(owner_host):
|
||||
errors.append(f"owner.host {owner_host!r} is not a lowercase slug [a-z0-9-]+ of 1-63 chars")
|
||||
owner_host = None
|
||||
|
||||
scope = _validate_scope(doc.get("scope"), errors) if "scope" in doc else None
|
||||
|
||||
mask = None
|
||||
if "mask" in doc:
|
||||
mask = _str_list("mask", "mask", doc.get("mask"), errors, allowed=MASK_OPS)
|
||||
if mask is not None and MASK_MANDATORY not in mask:
|
||||
errors.append(f"mask: {MASK_MANDATORY!r} is mandatory")
|
||||
mask = None
|
||||
|
||||
stages = _validate_stages(doc.get("stages"), errors) if "stages" in doc else None
|
||||
|
||||
sweep = doc.get("sweep")
|
||||
if "sweep" in doc and sweep not in SWEEPS:
|
||||
errors.append(f"sweep {sweep!r} is not one of {SWEEPS}")
|
||||
sweep = None
|
||||
|
||||
sink = _validate_sink(doc.get("sink"), errors) if "sink" in doc else None
|
||||
|
||||
# conditional sections — present iff their stage is
|
||||
sanitize = embed = summarize = None
|
||||
if stages is not None:
|
||||
if "sanitize" in stages and "sanitize" not in doc:
|
||||
errors.append("sanitize: section required when 'sanitize' is in stages")
|
||||
if "sanitize" not in stages and "sanitize" in doc:
|
||||
errors.append("sanitize: section present but 'sanitize' is not in stages")
|
||||
if "embed" in stages and "embed" not in doc:
|
||||
errors.append("embed: section required when 'embed' is in stages")
|
||||
if "embed" not in stages and "embed" in doc:
|
||||
errors.append("embed: section present but 'embed' is not in stages")
|
||||
if "sanitize" in doc:
|
||||
sanitize = _validate_sanitize(doc["sanitize"], manifest_id or "", errors)
|
||||
if "embed" in doc:
|
||||
embed = _validate_embed(doc["embed"], errors)
|
||||
if "summarize" in doc:
|
||||
summarize = _validate_summarize(doc["summarize"], errors)
|
||||
|
||||
grants = _validate_grants(doc.get("grants"), owner_host, mask, errors)
|
||||
|
||||
signature = None
|
||||
if "signature" in doc:
|
||||
signature = _validate_signature(doc["signature"], errors)
|
||||
elif require_signed:
|
||||
errors.append("signature: required (--require-signed) but absent")
|
||||
|
||||
# cross-section rules — the doctrine
|
||||
if stages is not None and mask is not None:
|
||||
for s in stages:
|
||||
need = STAGE_NEEDS_OP.get(s)
|
||||
if need and need not in mask:
|
||||
errors.append(f"stages: {s!r} requires {need!r} in mask")
|
||||
if "summarize" in doc:
|
||||
if mask is not None and "summarize" not in mask:
|
||||
errors.append("summarize: section present but 'summarize' is not in mask")
|
||||
if stages is not None and "catalog" not in stages:
|
||||
errors.append("summarize: requires the 'catalog' stage (summaries live in the catalog)")
|
||||
if stages is not None:
|
||||
if sink is not None and sink.catalog == "host" and "seal" not in stages:
|
||||
errors.append("sink: catalog 'host' requires the 'seal' stage — nothing leaves unsealed")
|
||||
if grants and "seal" not in stages:
|
||||
errors.append("grants: any grant requires the 'seal' stage — nothing is served unsealed")
|
||||
if sink is not None and sink.seal_to and "seal" not in stages:
|
||||
errors.append("sink: seal_to given but 'seal' is not in stages")
|
||||
|
||||
if errors:
|
||||
raise ManifestError(errors)
|
||||
return Manifest(
|
||||
id=manifest_id,
|
||||
owner_host=owner_host,
|
||||
scope=scope,
|
||||
mask=mask,
|
||||
stages=stages,
|
||||
sweep=sweep,
|
||||
sink=sink,
|
||||
sanitize=sanitize,
|
||||
embed=embed,
|
||||
summarize=summarize,
|
||||
grants=grants,
|
||||
signature=signature,
|
||||
raw=doc,
|
||||
)
|
||||
|
||||
|
||||
def load_manifest(path, require_signed=False):
|
||||
"""Load and validate a manifest file. Raises ManifestError listing all problems."""
|
||||
try:
|
||||
with open(path) as fh:
|
||||
doc = yaml.safe_load(fh)
|
||||
except FileNotFoundError:
|
||||
raise ManifestError([f"manifest file not found: {path}"])
|
||||
except yaml.YAMLError as exc:
|
||||
raise ManifestError([f"not valid YAML: {exc}"])
|
||||
return validate_document(doc, require_signed=require_signed)
|
||||
File diff suppressed because it is too large
Load Diff
Executable
+299
@@ -0,0 +1,299 @@
|
||||
#!/usr/bin/env bash
|
||||
# Publish the bastion agent distributable to the public edges. It is served at
|
||||
# BOTH addresses, because one of them is compiled into every bastion already
|
||||
# installed:
|
||||
#
|
||||
# new https://privatae.ai/bastion-agent/version.json
|
||||
# legacy https://privatae.ai/freeagent/version.json
|
||||
#
|
||||
# Those are two Caddy paths over ONE directory on each edge (WEBROOT below).
|
||||
# There is a single copy of every file; the new address is an alias, so the two
|
||||
# addresses can never drift or fall out of step at release time.
|
||||
#
|
||||
# THE LEGACY NAME, and when it ends. A bastion older than 0.20.0 knows only the
|
||||
# legacy address, builds the artifact's filename itself and asks the manifest
|
||||
# for `freeagent-<ver>.pyz`. If that address answers with a manifest that lists
|
||||
# only the new name, the bastion gets "no entry" and loses the agent; if the
|
||||
# address does not answer at all, the bastion keeps the agent it has. So:
|
||||
#
|
||||
# - while the legacy address is served, every publish writes the SAME BYTES
|
||||
# under BOTH filenames and lists both, each with the same sha256;
|
||||
# - once it is not served (it answers 404), a publish writes the new name only.
|
||||
#
|
||||
# The script reads which of the two is true on every run, on every edge, so it
|
||||
# cannot publish the one combination that strands a machine. Retiring the legacy address is a
|
||||
# step on the edges, taken once the published bastion is 0.42.0 or later (the
|
||||
# first release that never reads it); docs/packaging.md has it.
|
||||
#
|
||||
# The bastion's own deploy/publish.sh is the model: build
|
||||
# reproducibly, write a manifest with a sha256 per file, upload to EVERY edge,
|
||||
# then verify each edge's bytes against the manifest IT serves. Same key, same
|
||||
# edges, same webroot walls (Caddy refuses backups and dotfiles under both paths).
|
||||
#
|
||||
# Ratchet: refuse to publish a version that is not strictly newer than the one
|
||||
# already published (a rebuilt older version would regress installers that
|
||||
# compare versions). Override deliberately with ALLOW_REPUBLISH=1.
|
||||
set -euo pipefail
|
||||
# Repo root. This script sits in bastion_agent/, which is a top-level package
|
||||
# here, hence one `..`.
|
||||
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
||||
# WHERE THIS PUBLISHES IS NOT BAKED IN. This repository is read by everyone who
|
||||
# runs the agent, so it names no host, no key path and no webroot of ours: the
|
||||
# three below are required, and an unset one stops the script with a message
|
||||
# rather than publishing somewhere wrong. Whoever cuts a release supplies them.
|
||||
: "${KEY:?set KEY to the ssh key that reaches the edges}"
|
||||
# EDGE (one box) satisfies the requirement on its own -- it is the documented
|
||||
# single-box override below, and checking EDGES first would kill it with an error
|
||||
# naming a variable the operator was never told to set.
|
||||
if [ -n "${EDGE:-}" ]; then EDGES="$EDGE"; fi
|
||||
: "${EDGES:?set EDGES to the space-separated edge hosts to publish to (or EDGE=<one> for a single box)}"
|
||||
: "${WEBROOT:?set WEBROOT to the directory on each edge that both URLs serve}"
|
||||
: "${BACKUP_DIR:?set BACKUP_DIR to a directory OUTSIDE the webroot for rotated installers}"
|
||||
# OUTSIDE is load-bearing: a backup inside the webroot is served at a stable public
|
||||
# URL, which is how a superseded installer stays reachable forever. Assert it.
|
||||
case "$BACKUP_DIR" in
|
||||
"$WEBROOT"|"$WEBROOT"/*)
|
||||
echo "BACKUP_DIR ($BACKUP_DIR) is inside WEBROOT ($WEBROOT) -- rotated installers would be served publicly" >&2
|
||||
exit 1;;
|
||||
esac
|
||||
# The manifest is signed with the bastion's release key, which never comes to
|
||||
# this machine: it stays on the key holder's own computer. So a publish runs
|
||||
# twice. The first run builds, writes the manifest and prints the one
|
||||
# release-sign command to run there; the second, with SIGNATURE=<the line it
|
||||
# printed>, rebuilds (the build is reproducible, so the manifest is the same
|
||||
# bytes), checks the signature against the keys the bastion pins, and
|
||||
# publishes. BASTION_TREE is a checkout of privatae-bastion at the release the
|
||||
# fleet is moving to: its deploy/release-verify.sh and pinned keys are what a
|
||||
# bastion will check this manifest against.
|
||||
: "${BASTION_TREE:?set BASTION_TREE to a privatae-bastion checkout (its deploy/release-verify.sh checks the signature)}"
|
||||
[ -x "$BASTION_TREE/deploy/release-verify.sh" ] || { echo "$BASTION_TREE has no deploy/release-verify.sh" >&2; exit 1; }
|
||||
# The bastions that will read this manifest run the published bastion release,
|
||||
# and check it with the keys THAT release pins. So the signature is held to
|
||||
# those keys too, not only to BASTION_TREE's.
|
||||
FLEET_FROM="$(curl -fsS --max-time 15 https://privatae.ai/bastion/version.json | python3 -c 'import sys,json;print(json.load(sys.stdin).get("built_from",""))' 2>/dev/null || true)"
|
||||
git -C "$BASTION_TREE" fetch -q origin 2>/dev/null || true
|
||||
if [ -n "$FLEET_FROM" ] && ! git -C "$BASTION_TREE" cat-file -e "$FLEET_FROM^{commit}" 2>/dev/null; then
|
||||
echo "BASTION_TREE does not have the published bastion commit $FLEET_FROM; fetch it first" >&2
|
||||
exit 1
|
||||
fi
|
||||
BASE_NEW="https://privatae.ai/bastion-agent"
|
||||
BASE_OLD="https://privatae.ai/freeagent"
|
||||
|
||||
VER="$(tr -d '[:space:]' < "$HERE/bastion_agent/VERSION")"
|
||||
# Is the legacy address still served? It is asked of EVERY edge, not only of
|
||||
# the one the public name reaches now: each edge serves from its own disk and
|
||||
# its own Caddy config, and an edge that still has the route would serve a
|
||||
# manifest without the legacy name the day traffic fails over to it. Each edge
|
||||
# is asked for the route in the config Caddy has loaded (its local admin API).
|
||||
# Anything but a clear "no route" (an answer, a timeout, an error) counts as
|
||||
# served, because writing the legacy name when it is not needed costs one extra
|
||||
# file, and leaving it out when it is needed costs every old bastion its agent.
|
||||
# The public address is asked too, and a 404 there is also required.
|
||||
# Counts the routes whose path matcher would take /freeagent/version.json,
|
||||
# however the matcher is spelled (/freeagent/*, /freeagent*, a regexp), which
|
||||
# covers handle, handle_path and redir alike. Anything unreadable is "unknown".
|
||||
ROUTES_PY='
|
||||
import fnmatch, json, re, sys
|
||||
try:
|
||||
cfg = json.load(sys.stdin)
|
||||
except Exception:
|
||||
print("unknown"); sys.exit()
|
||||
if not isinstance(cfg, dict) or "apps" not in cfg:
|
||||
print("unknown"); sys.exit()
|
||||
target = "/freeagent/version.json"
|
||||
n = 0
|
||||
def walk(o):
|
||||
global n
|
||||
if isinstance(o, dict):
|
||||
for k, v in o.items():
|
||||
if k == "path" and isinstance(v, list) and any(isinstance(p, str) and fnmatch.fnmatchcase(target, p) and p not in ("*", "/*") for p in v):
|
||||
n += 1
|
||||
elif k == "path_regexp" and isinstance(v, dict) and isinstance(v.get("pattern"), str):
|
||||
try:
|
||||
if re.search(v["pattern"], target) and "freeagent" in v["pattern"].lower():
|
||||
n += 1
|
||||
except re.error:
|
||||
n += 1
|
||||
walk(v)
|
||||
elif isinstance(o, list):
|
||||
for x in o:
|
||||
walk(x)
|
||||
walk(cfg)
|
||||
print(n)
|
||||
'
|
||||
LEGACY=0
|
||||
for E in $EDGES; do
|
||||
n="$(ssh -n -i "$KEY" "$E" 'curl -fsS --max-time 10 http://localhost:2019/config/' 2>/dev/null \
|
||||
| python3 -c "$ROUTES_PY" 2>/dev/null || true)"
|
||||
if [ "$n" = 0 ]; then echo " $E: no /freeagent route"
|
||||
else LEGACY=1; echo " $E: /freeagent route ${n:-unknown} (served)"; fi
|
||||
done
|
||||
LEGACY_CODE="$(curl -sS -o /dev/null -w '%{http_code}' --max-time 15 "$BASE_OLD/version.json" 2>/dev/null || true)"
|
||||
[ "$LEGACY_CODE" = 404 ] || LEGACY=1
|
||||
if [ "$LEGACY" = 0 ]; then echo "legacy address retired on every edge (public: 404): publishing the new name only"
|
||||
else echo "legacy address still served somewhere (public: HTTP ${LEGACY_CODE:-none}): publishing both names"; fi
|
||||
GITSHA="$(git -C "$HERE" rev-parse --short HEAD 2>/dev/null || echo unknown)"
|
||||
echo "publishing bastion agent v$VER ($GITSHA)"
|
||||
|
||||
# Both addresses serve one directory, so the new one is the record of what is
|
||||
# published, and it stays served after the legacy one is retired.
|
||||
CUR="$(curl -fsS --max-time 15 "$BASE_NEW/version.json" 2>/dev/null | python3 -c 'import sys,json;print(json.load(sys.stdin).get("version",""))' 2>/dev/null || true)"
|
||||
if [ -n "$CUR" ]; then
|
||||
newer="$(python3 -c 'import sys
|
||||
a=[int(x) for x in sys.argv[1].split(".")]; b=[int(x) for x in sys.argv[2].split(".")]
|
||||
print("yes" if a>b else "no")' "$VER" "$CUR")"
|
||||
if [ "$newer" = yes ]; then echo "ratchet ok: $CUR -> $VER"; else
|
||||
echo "RATCHET REFUSED: v$VER is not newer than published v$CUR (ALLOW_REPUBLISH=1 to override)" >&2
|
||||
[ "${ALLOW_REPUBLISH:-}" = "1" ] || exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
DIST="$HERE/dist"
|
||||
(cd "$HERE" && python3 -m bastion_agent.build --out dist >/dev/null)
|
||||
PYZ="$DIST/bastion-agent-$VER.pyz"
|
||||
[ -f "$PYZ" ] || { echo "build did not produce $PYZ" >&2; exit 1; }
|
||||
python3 "$PYZ" --version >/dev/null
|
||||
|
||||
UPLOAD=("$PYZ" "$PYZ.sha256")
|
||||
if [ "$LEGACY" = 1 ]; then
|
||||
# The same bytes under the legacy filename, for bastions that ask for it by name.
|
||||
LEGACY_PYZ="$DIST/freeagent-$VER.pyz"
|
||||
cp -f "$PYZ" "$LEGACY_PYZ"
|
||||
sed "s|bastion-agent-$VER.pyz|freeagent-$VER.pyz|" "$PYZ.sha256" > "$LEGACY_PYZ.sha256"
|
||||
# Prove they are the same artifact rather than assuming the copy worked.
|
||||
a="$(python3 -c 'import hashlib,sys;print(hashlib.sha256(open(sys.argv[1],"rb").read()).hexdigest())' "$PYZ")"
|
||||
b="$(python3 -c 'import hashlib,sys;print(hashlib.sha256(open(sys.argv[1],"rb").read()).hexdigest())' "$LEGACY_PYZ")"
|
||||
[ "$a" = "$b" ] || { echo "the two filenames do not carry the same bytes" >&2; exit 1; }
|
||||
echo "artifact: $a (published under both names)"
|
||||
UPLOAD+=("$LEGACY_PYZ" "$LEGACY_PYZ.sha256")
|
||||
fi
|
||||
|
||||
python3 - "$VER" "$GITSHA" "$DIST" "$HERE/bastion_agent" "$BASE_NEW" "$BASE_OLD" "$LEGACY" > "$DIST/version.json" <<'PY'
|
||||
import sys, os, hashlib, json
|
||||
ver, sha, dist, pkg, base_new, base_old, legacy = sys.argv[1:8]
|
||||
def h(p): return hashlib.sha256(open(p, "rb").read()).hexdigest()
|
||||
def entry(name, base):
|
||||
p = os.path.join(dist, name)
|
||||
return name, {"url": "%s/dl/%s" % (base, name), "sha256": h(p), "bytes": os.path.getsize(p)}
|
||||
|
||||
files = {}
|
||||
names = [("bastion-agent-%s.pyz" % ver, base_new),
|
||||
("bastion-agent-%s.pyz.sha256" % ver, base_new)]
|
||||
# While the legacy address is served, the legacy pair is what a bastion older
|
||||
# than 0.20.0 asks for by name; leave it out and that bastion loses the agent.
|
||||
if legacy == "1":
|
||||
names += [("freeagent-%s.pyz" % ver, base_old),
|
||||
("freeagent-%s.pyz.sha256" % ver, base_old)]
|
||||
for name, base in names:
|
||||
k, v = entry(name, base)
|
||||
files[k] = v
|
||||
|
||||
installers = {n: {"url": "%s/%s" % (base_new, n), "sha256": h(os.path.join(pkg, n))}
|
||||
for n in ("install.sh", "install.ps1")}
|
||||
json.dump({"version": ver, "built_from": sha, "python_floor": "3.8",
|
||||
"files": files, "installers": installers,
|
||||
"notes": "https://git.privatae.ai/privatae/bastion-agent/releases/tag/v%s" % ver},
|
||||
sys.stdout, indent=2)
|
||||
PY
|
||||
|
||||
MAN="$DIST/version.json"
|
||||
MAN_SHA="$(python3 -c 'import hashlib,sys;print(hashlib.sha256(open(sys.argv[1],"rb").read()).hexdigest())' "$MAN")"
|
||||
if [ -z "${SIGNATURE:-}" ]; then
|
||||
cat <<ASK
|
||||
|
||||
Built bastion agent v$VER. Nothing has been published.
|
||||
|
||||
On the computer that holds the release key, run:
|
||||
|
||||
release-sign sign bastion-agent $VER $MAN_SHA
|
||||
|
||||
then run this again with the line it prints:
|
||||
|
||||
SIGNATURE='<that line>' bastion_agent/publish.sh
|
||||
|
||||
ASK
|
||||
exit 3
|
||||
fi
|
||||
# Named by the manifest's digest, the name a bastion asks for; old ones are
|
||||
# never removed, so a bastion that read the previous manifest a moment ago
|
||||
# still finds its signature.
|
||||
SIG_NAME="version.json.${MAN_SHA:0:16}.sig"
|
||||
printf '%s\n' "$SIGNATURE" > "$DIST/$SIG_NAME"
|
||||
SIGNED_BY="$("$BASTION_TREE/deploy/release-verify.sh" bastion-agent "$MAN" "$DIST/$SIG_NAME" ${FLEET_FROM:+"$FLEET_FROM"})" \
|
||||
|| { echo "SIGNATURE does not verify for bastion agent v$VER (manifest $MAN_SHA); nothing was published." >&2
|
||||
echo "If the key holder signed the command the first run printed, compare the 'legacy address' lines of the two runs: a change on the edges between them changes the manifest." >&2
|
||||
exit 1; }
|
||||
echo "manifest signature ok (key $SIGNED_BY)"
|
||||
|
||||
for E in $EDGES; do
|
||||
echo "--- publishing to $E ---"
|
||||
ssh -i "$KEY" "$E" 'rm -rf /tmp/ba && mkdir -p /tmp/ba'
|
||||
scp -q -i "$KEY" "${UPLOAD[@]}" "$DIST/version.json" "$DIST/$SIG_NAME" \
|
||||
"$HERE/bastion_agent/install.sh" "$HERE/bastion_agent/install.ps1" "$E":/tmp/ba/
|
||||
ssh -i "$KEY" "$E" "
|
||||
set -e
|
||||
install -d -m755 \"$WEBROOT/dl\"
|
||||
install -d -m700 \"$BACKUP_DIR\"
|
||||
ts=\$(date +%s)
|
||||
for f in install.sh install.ps1 version.json; do cp -a \"$WEBROOT/\$f\" \"$BACKUP_DIR/\$f.bak-\$ts\" 2>/dev/null || true; done
|
||||
for f in /tmp/ba/bastion-agent-* /tmp/ba/freeagent-*; do [ -e \"\$f\" ] || continue; b=\$(basename \"\$f\"); mv \"\$f\" \"$WEBROOT/dl/\$b\"; chmod 644 \"$WEBROOT/dl/\$b\"; done
|
||||
# The signature lands before the manifest it signs.
|
||||
for f in /tmp/ba/version.json.*.sig; do b=\$(basename \"\$f\"); mv \"\$f\" \"$WEBROOT/\$b\"; chmod 644 \"$WEBROOT/\$b\"; done
|
||||
for f in version.json install.sh install.ps1; do mv /tmp/ba/\$f \"$WEBROOT/\$f\"; chmod 644 \"$WEBROOT/\$f\"; done
|
||||
rm -rf /tmp/ba"
|
||||
done
|
||||
|
||||
# Verify each edge against the manifest IT serves, and, while the legacy address
|
||||
# is served, assert the legacy entries are present: their absence is the failure
|
||||
# that strands old bastions.
|
||||
for E in $EDGES; do
|
||||
printf ' %s: ' "$E"
|
||||
ssh -i "$KEY" "$E" "cd $WEBROOT && LEGACY=$LEGACY python3 - <<VERIFY
|
||||
import hashlib, json, os, sys
|
||||
m = json.load(open('version.json'))
|
||||
bad = []
|
||||
for name, v in m['files'].items():
|
||||
p = os.path.join('dl', name)
|
||||
if not os.path.exists(p): bad.append(name + ' MISSING'); continue
|
||||
if hashlib.sha256(open(p, 'rb').read()).hexdigest() != v['sha256']: bad.append(name + ' DIGEST-MISMATCH')
|
||||
for name, v in m['installers'].items():
|
||||
if not os.path.exists(name): bad.append(name + ' MISSING'); continue
|
||||
if hashlib.sha256(open(name, 'rb').read()).hexdigest() != v['sha256']: bad.append(name + ' DIGEST-MISMATCH')
|
||||
sig = 'version.json.%s.sig' % hashlib.sha256(open('version.json', 'rb').read()).hexdigest()[:16]
|
||||
if not os.path.exists(sig): bad.append(sig + ' MISSING')
|
||||
legacy = [n for n in m['files'] if n.startswith('freeagent-')]
|
||||
if os.environ.get('LEGACY') == '1' and not legacy: bad.append('NO LEGACY ENTRY - old bastions would lose the agent')
|
||||
print(('VERIFY FAILED: ' + '; '.join(bad)) if bad else
|
||||
'ok - %d files + %d installers match the manifest (v%s)%s' % (len(m['files']), len(m['installers']), m['version'],
|
||||
', legacy entries present' if legacy else ''))
|
||||
sys.exit(1 if bad else 0)
|
||||
VERIFY"
|
||||
done
|
||||
|
||||
# From the public internet, at every address still served, the way a bastion asks
|
||||
# (version.json.<first 8 bytes of the manifest's sha256>.sig), checked against
|
||||
# the manifest that address serves and the keys the bastion pins.
|
||||
SIGCHECK="$(mktemp -d)"; trap 'rm -rf "$SIGCHECK"' EXIT
|
||||
# The legacy address is checked from outside only when it answers there: while
|
||||
# edges disagree, the public name may reach one that no longer has the route,
|
||||
# and the per-edge check above has already held every edge's files to the
|
||||
# manifest.
|
||||
ADDRS=("$BASE_NEW"); [ "$LEGACY" = 1 ] && [ "$LEGACY_CODE" != 404 ] && ADDRS+=("$BASE_OLD")
|
||||
for B in "${ADDRS[@]}"; do
|
||||
curl -fsS --max-time 30 "$B/version.json" -o "$SIGCHECK/m"
|
||||
H="$(python3 -c 'import hashlib,sys;print(hashlib.sha256(open(sys.argv[1],"rb").read()).hexdigest()[:16])' "$SIGCHECK/m")"
|
||||
curl -fsS --max-time 30 "$B/version.json.$H.sig" -o "$SIGCHECK/s" \
|
||||
|| { echo " VERIFY FAILED: $B/version.json.$H.sig is not served publicly" >&2; exit 1; }
|
||||
if "$BASTION_TREE/deploy/release-verify.sh" bastion-agent "$SIGCHECK/m" "$SIGCHECK/s" ${FLEET_FROM:+"$FLEET_FROM"} >/dev/null; then
|
||||
echo " $B/version.json.$H.sig ok"
|
||||
else
|
||||
echo " VERIFY FAILED: $B/version.json.$H.sig does not verify against the manifest served there" >&2
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
|
||||
echo "published bastion agent v$VER to: $EDGES"
|
||||
echo " new: $BASE_NEW/version.json · $BASE_NEW/dl/bastion-agent-$VER.pyz"
|
||||
if [ "$LEGACY" = 1 ]; then
|
||||
echo " legacy: $BASE_OLD/version.json · $BASE_OLD/dl/freeagent-$VER.pyz"
|
||||
fi
|
||||
@@ -0,0 +1,447 @@
|
||||
"""bastion agent runtime: the `sweep` scheduler and the manifest gate.
|
||||
|
||||
This is what turns seven commands into an agent. Each cycle:
|
||||
|
||||
1. re-read the manifest (a re-signed manifest takes effect next sweep)
|
||||
2. GATE: verify its signature — ECDSA P-256 over the canonical payload,
|
||||
the leaf cert's CN must equal owner.host, and the cert must chain to
|
||||
FREEAGENT_CA_ROOT, which is required. Unsigned or unverifiable → the
|
||||
sweep does nothing and says so. (--allow-unsigned exists for
|
||||
development and is recorded in every run it touches.)
|
||||
3. run the stages in manifest order, each with the module already
|
||||
proven alone; the first failure stops the run (fail closed — a
|
||||
later stage never consumes a half-finished earlier one)
|
||||
4. append one line to run.jsonl, rewrite run.state.json, sleep until
|
||||
the next due time (manifest `sweep`), or exit for `manual`.
|
||||
|
||||
A lock file (atomic create, pid + time inside) keeps two sweeps off one
|
||||
output directory; a stale lock (dead pid on POSIX, or older than
|
||||
LOCK_STALE_HOURS anywhere) is reclaimed and logged.
|
||||
|
||||
Stdlib only. The gate's certificate and signature checks are in x509gate
|
||||
(the stdlib has no ECDSA, so it carries its own P-256 verifier); openssl is
|
||||
needed only by the developer helper that SIGNS a manifest. Python 3.8 floor.
|
||||
"""
|
||||
|
||||
import datetime as _dt
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
from . import catalog as _catalog
|
||||
from . import embed as _embed
|
||||
from . import extract as _extract
|
||||
from . import sanitize as _sanitize
|
||||
from . import seal as _seal
|
||||
from . import walk as _walk
|
||||
from . import x509gate as _x509
|
||||
from .manifest import AGENT_DIR, ManifestError, load_manifest, signing_payload
|
||||
|
||||
RUN_VERSION = 1
|
||||
SWEEP_SECONDS = {"15m": 15 * 60, "hourly": 60 * 60, "daily": 24 * 60 * 60, "manual": None}
|
||||
LOCK_STALE_HOURS = 12
|
||||
MIN_RETRY_SECONDS = 60
|
||||
|
||||
|
||||
class RunError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class GateError(RunError):
|
||||
"""The manifest did not pass the gate; nothing ran."""
|
||||
|
||||
|
||||
# ── the gate ────────────────────────────────────────────────────────────────
|
||||
|
||||
def _openssl():
|
||||
p = shutil.which("openssl")
|
||||
if not p:
|
||||
raise GateError("openssl is not on PATH — cannot verify the manifest signature")
|
||||
return p
|
||||
|
||||
|
||||
def _run(argv, data=None):
|
||||
try:
|
||||
proc = subprocess.run(argv, input=data, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
||||
timeout=60, check=False)
|
||||
except (OSError, subprocess.TimeoutExpired) as exc:
|
||||
raise GateError("%s: %s" % (argv[0], exc))
|
||||
return proc.returncode, proc.stdout, proc.stderr.decode("utf-8", "replace")
|
||||
|
||||
|
||||
def cert_cn(cert_path):
|
||||
try:
|
||||
return _x509.load_cert_file(cert_path).cn()
|
||||
except (OSError, _x509.CertError, ValueError, TypeError, IndexError, OverflowError) as exc:
|
||||
raise GateError("cannot read signing cert %s: %s" % (cert_path, exc))
|
||||
|
||||
|
||||
def verify_signature(manifest, env, manifest_dir="."):
|
||||
"""Raise GateError unless the manifest carries a valid signature.
|
||||
|
||||
Standard library only (x509gate): the gate used to shell out to openssl,
|
||||
which Windows does not ship, so on a Windows machine every manifest was
|
||||
refused and the agent never ran. The checks
|
||||
are the ones openssl made: exactly one CN, equal to the owner; the cert
|
||||
chains to FREEAGENT_CA_ROOT and is inside its validity window; the
|
||||
ECDSA P-256 / SHA-256 signature verifies over the canonical payload."""
|
||||
sig = manifest.signature
|
||||
if sig is None:
|
||||
raise GateError("manifest is unsigned — the runtime refuses to execute it")
|
||||
cert = sig.cert if os.path.isabs(sig.cert) else os.path.join(manifest_dir, sig.cert)
|
||||
if not os.path.exists(cert):
|
||||
raise GateError("signing cert %s not found" % cert)
|
||||
try:
|
||||
leaf = _x509.load_cert_file(cert)
|
||||
cn = leaf.cn()
|
||||
except (OSError, _x509.CertError, ValueError, TypeError, IndexError, OverflowError) as exc:
|
||||
# Anything a malformed cert raises is a refusal, never a crash.
|
||||
raise GateError("cannot read signing cert %s: %s" % (cert, exc))
|
||||
if cn != manifest.owner_host:
|
||||
raise GateError("signing cert CN %r is not the manifest owner %r" % (cn, manifest.owner_host))
|
||||
ca_root = env.get("FREEAGENT_CA_ROOT", "").strip()
|
||||
if not ca_root:
|
||||
# without a root the gate would only prove "signed by the key in the
|
||||
# cert next to the manifest" — that authenticates nothing
|
||||
raise GateError("FREEAGENT_CA_ROOT is not set — the gate cannot verify who signed this manifest")
|
||||
try:
|
||||
with open(ca_root, "rb") as fh:
|
||||
roots = _x509.load_root_file(fh.read())
|
||||
_x509.verify_chain(leaf, roots)
|
||||
except OSError as exc:
|
||||
raise GateError("cannot read FREEAGENT_CA_ROOT %s: %s" % (ca_root, exc))
|
||||
except (_x509.CertError, ValueError, TypeError, IndexError, OverflowError) as exc:
|
||||
raise GateError("signing cert does not chain to FREEAGENT_CA_ROOT: %s" % exc)
|
||||
import base64
|
||||
try:
|
||||
raw = base64.b64decode(sig.value, validate=True)
|
||||
except Exception:
|
||||
raise GateError("signature value is not base64")
|
||||
if not _x509.ecdsa_p256_sha256_verify(leaf.public_key, signing_payload(manifest.raw), raw):
|
||||
raise GateError("manifest signature does not verify (payload changed or wrong key)")
|
||||
return {"cn": cn, "cert": cert, "chained": True}
|
||||
|
||||
|
||||
def manifest_payload(path):
|
||||
"""The exact bytes a manifest's signature covers (the canonical payload).
|
||||
A signer that holds the key elsewhere — the Go bastion, which mints the
|
||||
managed agent's identity with its own crypto — signs these and hands the
|
||||
result back through write_signature."""
|
||||
return signing_payload(load_manifest(path).raw)
|
||||
|
||||
|
||||
def write_signature(path, der_signature, cert_path):
|
||||
"""Write the signature block (a DER ECDSA signature over manifest_payload)
|
||||
into the manifest, replacing any earlier one. Checked before it is kept:
|
||||
the manifest must still load as signed."""
|
||||
import base64
|
||||
with open(path, encoding="utf-8") as fh:
|
||||
text = fh.read()
|
||||
# simpler and safe: drop an existing top-level signature block, then append
|
||||
body = _strip_signature_block(text)
|
||||
block = "signature:\n alg: ecdsa-p256-sha256\n cert: %s\n value: %s\n" % (
|
||||
cert_path, base64.b64encode(der_signature).decode("ascii"))
|
||||
with open(path, "w", encoding="utf-8", newline="\n") as fh:
|
||||
fh.write(body.rstrip("\n") + "\n" + block)
|
||||
load_manifest(path, require_signed=True) # must still validate
|
||||
return block
|
||||
|
||||
|
||||
def sign_manifest_file(path, key_path, cert_path):
|
||||
"""Developer/host helper: sign the canonical payload with a PEM key via
|
||||
openssl and write the signature block. The managed install does not use
|
||||
this — the bastion signs with its own key and calls write_signature — so
|
||||
no machine needs openssl to run the agent."""
|
||||
rc, out, err = _run([_openssl(), "dgst", "-sha256", "-sign", key_path],
|
||||
data=manifest_payload(path))
|
||||
if rc != 0:
|
||||
raise GateError("signing failed: %s" % err.strip())
|
||||
return write_signature(path, out, cert_path)
|
||||
|
||||
|
||||
def _strip_signature_block(text):
|
||||
out, skipping = [], False
|
||||
for line in text.splitlines(keepends=True):
|
||||
if line.startswith("signature:"):
|
||||
skipping = True
|
||||
continue
|
||||
if skipping and (line.startswith(" ") or line.startswith("\t")):
|
||||
continue
|
||||
skipping = False
|
||||
out.append(line)
|
||||
return "".join(out)
|
||||
|
||||
|
||||
# ── the lock ────────────────────────────────────────────────────────────────
|
||||
|
||||
def _pid_alive(pid):
|
||||
if os.name == "nt":
|
||||
return True # cannot cheaply tell; the age rule decides
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
except ProcessLookupError:
|
||||
return False
|
||||
except PermissionError:
|
||||
return True
|
||||
return True
|
||||
|
||||
|
||||
class Lock:
|
||||
def __init__(self, out_dir):
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
self.path = os.path.join(out_dir, "sweep.lock")
|
||||
self.reclaimed = None
|
||||
|
||||
def __enter__(self):
|
||||
for attempt in range(2):
|
||||
try:
|
||||
fd = os.open(self.path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
|
||||
except FileExistsError:
|
||||
if attempt == 0 and self._stale():
|
||||
self.reclaimed = self._read()
|
||||
try:
|
||||
os.unlink(self.path)
|
||||
except OSError:
|
||||
pass
|
||||
continue
|
||||
raise RunError("another sweep holds %s (%s)" % (self.path, self._read()))
|
||||
with os.fdopen(fd, "w") as fh:
|
||||
json.dump({"pid": os.getpid(), "since": time.time()}, fh)
|
||||
return self
|
||||
raise RunError("could not acquire %s" % self.path)
|
||||
|
||||
def _read(self):
|
||||
try:
|
||||
with open(self.path) as fh:
|
||||
return json.load(fh)
|
||||
except (OSError, ValueError):
|
||||
return {}
|
||||
|
||||
def _stale(self):
|
||||
"""A lock is stale when its holder is provably dead. On POSIX that is
|
||||
the pid check and nothing else — a live pid keeps its lock however
|
||||
long the sweep runs (age alone stole live locks). The
|
||||
age rule applies only where the pid cannot be checked (Windows) or the
|
||||
file is unreadable, and then by the FILE's mtime, so a truncated lock
|
||||
is not instantly stale."""
|
||||
info = self._read()
|
||||
pid = info.get("pid")
|
||||
if os.name != "nt" and isinstance(pid, int):
|
||||
return not _pid_alive(pid)
|
||||
try:
|
||||
age = time.time() - os.stat(self.path).st_mtime
|
||||
except OSError:
|
||||
return False
|
||||
return age > LOCK_STALE_HOURS * 3600
|
||||
|
||||
def __exit__(self, *exc):
|
||||
try:
|
||||
os.unlink(self.path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
# ── one run ─────────────────────────────────────────────────────────────────
|
||||
|
||||
STAGES = {
|
||||
"walk": lambda m, out, env: _walk.walk(m, out),
|
||||
"extract": lambda m, out, env: _extract.extract(m, out, use_external=env.get("FREEAGENT_NO_EXTERNAL") != "1"),
|
||||
"sanitize": lambda m, out, env: _sanitize.sanitize(m, out, env=env),
|
||||
"chunk": lambda m, out, env: _embed.chunk(m, out),
|
||||
"embed": lambda m, out, env: _embed.embed(m, out, env=env),
|
||||
"catalog": lambda m, out, env: _catalog.catalog(m, out, env=env),
|
||||
"seal": lambda m, out, env: _seal.seal(m, out, env=env),
|
||||
}
|
||||
STAGE_ERRORS = (FileNotFoundError, _sanitize.SanitizeError, _embed.EmbedError,
|
||||
_catalog.CatalogError, _seal.SealError, OSError)
|
||||
|
||||
|
||||
def _stage_failed(name, summary):
|
||||
"""A stage that completed but reports it stopped is a failure too."""
|
||||
if summary.get("stopped"):
|
||||
return summary["stopped"]
|
||||
return None
|
||||
|
||||
|
||||
def _brief(name, summary):
|
||||
c = summary.get("counts", {})
|
||||
keys = {"walk": ("files", "errors"), "extract": ("ok", "reused", "error"),
|
||||
"sanitize": ("ok", "reused", "not_run"), "chunk": ("chunks",),
|
||||
"embed": ("embedded", "reused"), "catalog": ("files", "dirs", "summaries"),
|
||||
"seal": ()}.get(name, ())
|
||||
out = {k: c[k] for k in keys if k in c}
|
||||
if name == "seal":
|
||||
out = {"bytes_sealed": summary.get("bytes_sealed"), "recipients": len(summary.get("recipients", []))}
|
||||
return out
|
||||
|
||||
|
||||
def default_out(m):
|
||||
"""Where an agent keeps itself: <root>/.agent for a single-root manifest
|
||||
(files-and-folders style: copy the folder, the agent comes with it);
|
||||
build/freeagent/<id> when there are several roots and no obvious home.
|
||||
Every command uses this default, so `sweep` and `serve` run without
|
||||
--out share one folder.
|
||||
|
||||
The directory keeps the old name deliberately: it is on-disk output a user
|
||||
may already have, and renaming it would orphan it."""
|
||||
roots = m.scope.roots
|
||||
if len(roots) == 1:
|
||||
return os.path.join(os.path.abspath(os.path.expanduser(roots[0])), AGENT_DIR)
|
||||
return os.path.join("build", "freeagent", m.id)
|
||||
|
||||
|
||||
def run_once(manifest_path, out_dir=None, env=None, allow_unsigned=False, now=None, clock=time.time):
|
||||
env = dict(os.environ) if env is None else env
|
||||
started = now or _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
t0 = clock()
|
||||
record = {"v": RUN_VERSION, "started": started, "manifest": manifest_path, "stages": [],
|
||||
"ok": False, "allow_unsigned": bool(allow_unsigned)}
|
||||
try:
|
||||
m = load_manifest(manifest_path)
|
||||
except ManifestError as exc:
|
||||
record.update(error="manifest: %s" % "; ".join(exc.errors), finished=_now(now), seconds=round(clock() - t0, 3))
|
||||
# a refused manifest is a run that happened; the append-only log must
|
||||
# say so. With no id to derive a directory from, the
|
||||
# caller's --out or a fixed 'refused' directory holds the record.
|
||||
_append_run(out_dir or os.path.join("build", "freeagent", "refused"), record)
|
||||
return record, None
|
||||
out_dir = out_dir or default_out(m)
|
||||
record["out"] = out_dir
|
||||
record["manifest_id"] = m.id
|
||||
from .manifest import signing_digest
|
||||
record["manifest_digest"] = "sha256:" + signing_digest(m.raw)
|
||||
try:
|
||||
gate = verify_signature(m, env, manifest_dir=os.path.dirname(os.path.abspath(manifest_path)))
|
||||
record["gate"] = gate
|
||||
except GateError as exc:
|
||||
if not allow_unsigned:
|
||||
record.update(error="gate: %s" % exc, finished=_now(now), seconds=round(clock() - t0, 3))
|
||||
_append_run(out_dir, record) # a gate refusal is on the record too
|
||||
return record, m
|
||||
record["gate"] = {"bypassed": str(exc)}
|
||||
try:
|
||||
with Lock(out_dir) as lock:
|
||||
if lock.reclaimed:
|
||||
record["lock_reclaimed"] = lock.reclaimed
|
||||
for name in m.stages:
|
||||
s0 = clock()
|
||||
entry = {"stage": name}
|
||||
try:
|
||||
summary = STAGES[name](m, out_dir, env)
|
||||
err = _stage_failed(name, summary)
|
||||
entry.update(status="failed" if err else "ok", seconds=round(clock() - s0, 3),
|
||||
**_brief(name, summary))
|
||||
if err:
|
||||
entry["error"] = err
|
||||
except STAGE_ERRORS as exc:
|
||||
entry.update(status="failed", seconds=round(clock() - s0, 3),
|
||||
error="%s: %s" % (type(exc).__name__, exc))
|
||||
record["stages"].append(entry)
|
||||
if entry["status"] != "ok":
|
||||
record["error"] = "stage %s: %s" % (name, entry["error"])
|
||||
break
|
||||
else:
|
||||
record["ok"] = True
|
||||
except RunError as exc:
|
||||
record["error"] = str(exc)
|
||||
record["finished"] = _now(now)
|
||||
record["seconds"] = round(clock() - t0, 3)
|
||||
_append_run(out_dir, record)
|
||||
return record, m
|
||||
|
||||
|
||||
def _now(fixed=None):
|
||||
return fixed or _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
|
||||
|
||||
def _append_run(out_dir, record):
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
with open(os.path.join(out_dir, "run.jsonl"), "a", encoding="utf-8") as fh:
|
||||
fh.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")
|
||||
|
||||
|
||||
def _write_state(out_dir, state):
|
||||
path = os.path.join(out_dir, "run.state.json")
|
||||
fd, tmp = tempfile.mkstemp(prefix=".tmp-", dir=out_dir)
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as fh:
|
||||
json.dump(state, fh, indent=2, sort_keys=True)
|
||||
os.replace(tmp, path)
|
||||
|
||||
|
||||
def load_state(out_dir):
|
||||
try:
|
||||
with open(os.path.join(out_dir, "run.state.json"), encoding="utf-8") as fh:
|
||||
return json.load(fh)
|
||||
except (OSError, ValueError):
|
||||
return {}
|
||||
|
||||
|
||||
def next_due(sweep, last_started_ts, ok, failures):
|
||||
"""Seconds-since-epoch of the next run. Failures retry with backoff
|
||||
(1 min, 2, 4 … capped at the cadence); manual never schedules."""
|
||||
period = SWEEP_SECONDS.get(sweep)
|
||||
if period is None:
|
||||
return None
|
||||
if ok:
|
||||
return last_started_ts + period
|
||||
backoff = min(period, MIN_RETRY_SECONDS * (2 ** min(failures - 1, 10)))
|
||||
return last_started_ts + backoff
|
||||
|
||||
|
||||
# ── the loop ────────────────────────────────────────────────────────────────
|
||||
|
||||
def sweep(manifest_path, out_dir=None, env=None, allow_unsigned=False, once=False,
|
||||
sleep=time.sleep, clock=time.time, max_runs=None, log=None):
|
||||
"""Run on the manifest's cadence until stopped. Returns the last record."""
|
||||
env = dict(os.environ) if env is None else env
|
||||
log = log or (lambda s: sys.stderr.write(s + "\n"))
|
||||
stop = {"flag": False}
|
||||
|
||||
def _stop(signum, frame):
|
||||
stop["flag"] = True
|
||||
log("sweep: stop requested (signal %d); finishing the current run" % signum)
|
||||
|
||||
for sig in (getattr(signal, "SIGTERM", None), getattr(signal, "SIGINT", None)):
|
||||
if sig is not None:
|
||||
try:
|
||||
signal.signal(sig, _stop)
|
||||
except (ValueError, OSError):
|
||||
pass # not the main thread (tests) — run without handlers
|
||||
|
||||
runs = 0
|
||||
last = None
|
||||
while not stop["flag"]:
|
||||
t_start = clock()
|
||||
record, m = run_once(manifest_path, out_dir, env, allow_unsigned=allow_unsigned, clock=clock)
|
||||
runs += 1
|
||||
last = record
|
||||
out = record.get("out") or out_dir or "."
|
||||
state = load_state(out) if os.path.isdir(out) else {}
|
||||
failures = 0 if record["ok"] else state.get("consecutive_failures", 0) + 1
|
||||
sweep_kind = m.sweep if m is not None else "manual"
|
||||
due = next_due(sweep_kind, t_start, record["ok"], failures)
|
||||
state = {"last_started": record["started"], "last_finished": record.get("finished"),
|
||||
"last_ok": record["ok"], "last_error": record.get("error"),
|
||||
"consecutive_failures": failures, "sweep": sweep_kind,
|
||||
"next_due": _dt.datetime.fromtimestamp(due, _dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") if due else None,
|
||||
"pid": os.getpid(), "runs_this_process": runs}
|
||||
if os.path.isdir(out):
|
||||
_write_state(out, state)
|
||||
stages = ", ".join("%s:%s" % (s["stage"], s["status"]) for s in record["stages"]) or "-"
|
||||
log("sweep: %s %s in %.1fs [%s]%s" % (
|
||||
record.get("manifest_id", "?"), "ok" if record["ok"] else "FAILED",
|
||||
record.get("seconds", 0), stages,
|
||||
"" if record["ok"] else " — " + str(record.get("error"))))
|
||||
if once or due is None or (max_runs and runs >= max_runs):
|
||||
break
|
||||
while not stop["flag"]:
|
||||
wait = due - clock()
|
||||
if wait <= 0:
|
||||
break
|
||||
sleep(min(wait, 30))
|
||||
return last
|
||||
@@ -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
|
||||
@@ -0,0 +1,238 @@
|
||||
"""bastion agent `seal` stage — encrypt the catalog to the owner (and named
|
||||
recipients) so it can leave the machine.
|
||||
|
||||
What is sealed — and ONLY this — is an allowlisted bundle:
|
||||
|
||||
manifest.json the manifest as executed (minus signature)
|
||||
catalog.jsonl file cards
|
||||
catalog.dirs.jsonl directory cards
|
||||
*.summary.json every stage's summary — counts, digests,
|
||||
tool names; never a value
|
||||
seal.manifest.json bundle inventory with sha256 per member
|
||||
|
||||
Never in the bundle, by construction (the allowlist is the mechanism, a
|
||||
test is the proof): the extract store (plaintext), the sanitized store,
|
||||
the vault, vectors.db, chunks. `sink.plaintext: never` is enforced here
|
||||
one more time by not having a code path that could add those files.
|
||||
|
||||
The envelope is age (X25519 + ChaCha20-Poly1305): a small, widely reviewed
|
||||
format with a single static binary on Mac/Windows/Linux. It is REQUIRED for this stage
|
||||
— the stdlib cannot produce the envelope, and a catalog that cannot be
|
||||
sealed does not leave. Recipients:
|
||||
|
||||
owner FREEAGENT_OWNER_RECIPIENT (an age1… public key), from the
|
||||
environment, never the manifest. Always present.
|
||||
seal_to manifest sink.seal_to — each entry an age1… recipient (or
|
||||
an ssh-ed25519/ssh-rsa public key line, which age accepts).
|
||||
|
||||
Grantees (manifest `grants`) are NOT recipients: they query the container
|
||||
over the mesh; they never receive the sealed catalog.
|
||||
|
||||
Output: <out>/sealed/<id>-<stamp>.tar.age, <out>/sealed/latest.json (a
|
||||
pointer, not a symlink — Windows), <out>/seal.summary.json. Optional
|
||||
detached signature with the host's key: SEAL_SIGN_KEY (PEM path) →
|
||||
<file>.sig via openssl, so a receiving host can check who sealed it once
|
||||
the runtime's CA verification lands.
|
||||
"""
|
||||
|
||||
import datetime as _dt
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import tarfile
|
||||
import tempfile
|
||||
|
||||
from .manifest import signing_digest
|
||||
|
||||
SEAL_VERSION = 1
|
||||
AGE_RECIPIENT_RE = re.compile(r"^(age1[a-z0-9]{58}|ssh-(ed25519|rsa) [A-Za-z0-9+/=]+( .*)?)$")
|
||||
BUNDLE_ALLOW = ("catalog.jsonl", "catalog.dirs.jsonl")
|
||||
SUMMARY_GLOB_RE = re.compile(r"^[a-z]+\.summary\.json$")
|
||||
TOOL_TIMEOUT = 300
|
||||
|
||||
|
||||
class SealError(Exception):
|
||||
"""Fail closed: nothing leaves."""
|
||||
|
||||
|
||||
def tool_path(name):
|
||||
return shutil.which(name)
|
||||
|
||||
|
||||
def run(argv, data=None, timeout=TOOL_TIMEOUT):
|
||||
try:
|
||||
proc = subprocess.run(argv, input=data, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
||||
timeout=timeout, check=False)
|
||||
except (OSError, subprocess.TimeoutExpired) as exc:
|
||||
raise SealError("%s: %s" % (argv[0], exc))
|
||||
if proc.returncode != 0:
|
||||
err = proc.stderr.decode("utf-8", "replace").strip().splitlines()
|
||||
raise SealError("%s exited %d: %s" % (argv[0], proc.returncode, err[-1] if err else ""))
|
||||
return proc.stdout
|
||||
|
||||
|
||||
def recipients_for(manifest, env):
|
||||
owner = env.get("FREEAGENT_OWNER_RECIPIENT", "").strip()
|
||||
if not owner:
|
||||
raise SealError("FREEAGENT_OWNER_RECIPIENT is not set — the owner's age public key "
|
||||
"is the mandatory recipient; refusing to seal to nobody")
|
||||
recips = [owner] + list(manifest.sink.seal_to)
|
||||
bad = [r for r in recips if not AGE_RECIPIENT_RE.match(r)]
|
||||
if bad:
|
||||
raise SealError("not an age recipient (age1… or ssh-ed25519/ssh-rsa public key): %s"
|
||||
% ", ".join(r[:24] + "…" if len(r) > 24 else r for r in bad))
|
||||
seen = set()
|
||||
out = []
|
||||
for r in recips:
|
||||
if r not in seen:
|
||||
seen.add(r)
|
||||
out.append(r)
|
||||
return out
|
||||
|
||||
|
||||
def _sha256_file(path):
|
||||
h = hashlib.sha256()
|
||||
with open(path, "rb") as fh:
|
||||
for buf in iter(lambda: fh.read(1024 * 1024), b""):
|
||||
h.update(buf)
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def build_bundle(manifest, out_dir, stamp):
|
||||
"""The plaintext tar, in memory. Returns (bytes, inventory)."""
|
||||
members = []
|
||||
for name in BUNDLE_ALLOW:
|
||||
path = os.path.join(out_dir, name)
|
||||
if not os.path.exists(path):
|
||||
raise SealError("no %s in %s — run the catalog stage first" % (name, out_dir))
|
||||
members.append(name)
|
||||
for name in sorted(os.listdir(out_dir)):
|
||||
if SUMMARY_GLOB_RE.match(name) and name != "seal.summary.json":
|
||||
members.append(name)
|
||||
inventory = {"freeagent": 1, "seal_version": SEAL_VERSION, "manifest_id": manifest.id,
|
||||
"manifest_digest": "sha256:" + signing_digest(manifest.raw),
|
||||
"sealed_at": stamp, "members": {}}
|
||||
buf = io.BytesIO()
|
||||
with tarfile.open(fileobj=buf, mode="w") as tar:
|
||||
def add_bytes(name, data):
|
||||
info = tarfile.TarInfo(name)
|
||||
info.size = len(data)
|
||||
info.mtime = 0 # reproducible: the content is what is signed, not the clock
|
||||
info.mode = 0o600
|
||||
tar.addfile(info, io.BytesIO(data))
|
||||
inventory["members"][name] = {"sha256": hashlib.sha256(data).hexdigest(), "bytes": len(data)}
|
||||
|
||||
body = {k: v for k, v in manifest.raw.items() if k != "signature"}
|
||||
add_bytes("manifest.json", json.dumps(body, sort_keys=True, indent=2, default=str).encode("utf-8"))
|
||||
for name in members:
|
||||
with open(os.path.join(out_dir, name), "rb") as fh:
|
||||
add_bytes(name, fh.read())
|
||||
add_bytes("seal.manifest.json", json.dumps(inventory, sort_keys=True, indent=2).encode("utf-8"))
|
||||
return buf.getvalue(), inventory
|
||||
|
||||
|
||||
def seal(manifest, out_dir, env=None, now=None):
|
||||
if "seal" not in manifest.stages:
|
||||
raise SealError("'seal' is not in this manifest's stages")
|
||||
env = dict(os.environ) if env is None else env
|
||||
age = tool_path("age")
|
||||
if not age:
|
||||
raise SealError("`age` is not on PATH — the envelope cannot be produced, nothing leaves "
|
||||
"(install age: https://age-encryption.org, one static binary)")
|
||||
recips = recipients_for(manifest, env)
|
||||
stamp = now or _dt.datetime.now(_dt.timezone.utc).strftime("%Y%m%dT%H%M%SZ")
|
||||
plain, inventory = build_bundle(manifest, out_dir, stamp)
|
||||
|
||||
sealed_dir = os.path.join(out_dir, "sealed")
|
||||
os.makedirs(sealed_dir, exist_ok=True)
|
||||
fname = "%s-%s.tar.age" % (manifest.id, stamp)
|
||||
target = os.path.join(sealed_dir, fname)
|
||||
argv = [age, "--encrypt"]
|
||||
for r in recips:
|
||||
argv += ["-r", r]
|
||||
argv += ["-o", target + ".tmp"]
|
||||
run(argv, data=plain)
|
||||
os.replace(target + ".tmp", target)
|
||||
try:
|
||||
os.chmod(target, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
signature = None
|
||||
sign_key = env.get("SEAL_SIGN_KEY", "").strip()
|
||||
if sign_key:
|
||||
openssl = tool_path("openssl")
|
||||
if not openssl:
|
||||
raise SealError("SEAL_SIGN_KEY set but openssl is not on PATH")
|
||||
sig = run([openssl, "dgst", "-sha256", "-sign", sign_key, target])
|
||||
with open(target + ".sig", "wb") as fh:
|
||||
fh.write(sig)
|
||||
signature = os.path.basename(target) + ".sig"
|
||||
|
||||
summary = {
|
||||
"freeagent": 1, "stage": "seal", "seal_version": SEAL_VERSION,
|
||||
"manifest_id": manifest.id,
|
||||
"manifest_digest": inventory["manifest_digest"],
|
||||
"sealed_at": stamp,
|
||||
"file": os.path.join("sealed", fname),
|
||||
"bytes_plain": len(plain),
|
||||
"bytes_sealed": os.path.getsize(target),
|
||||
"sha256_plain": hashlib.sha256(plain).hexdigest(),
|
||||
"sha256_sealed": _sha256_file(target),
|
||||
"recipients": recips, # public keys — safe to record
|
||||
"members": sorted(inventory["members"]),
|
||||
"signature": signature,
|
||||
"catalog_sink": manifest.sink.catalog,
|
||||
}
|
||||
pointer = {"file": summary["file"], "sha256_sealed": summary["sha256_sealed"],
|
||||
"sealed_at": stamp, "manifest_digest": summary["manifest_digest"]}
|
||||
_atomic_write(os.path.join(sealed_dir, "latest.json"), json.dumps(pointer, indent=2, sort_keys=True) + "\n")
|
||||
_atomic_write(os.path.join(out_dir, "seal.summary.json"), json.dumps(summary, indent=2, sort_keys=True) + "\n")
|
||||
return summary
|
||||
|
||||
|
||||
def unseal(sealed_path, identity_path, dest_dir):
|
||||
"""Owner-side: open a sealed bundle with an age identity file into
|
||||
dest_dir. Verifies every member against the inventory."""
|
||||
age = tool_path("age")
|
||||
if not age:
|
||||
raise SealError("`age` is not on PATH")
|
||||
plain = run([age, "--decrypt", "-i", identity_path, sealed_path])
|
||||
os.makedirs(dest_dir, exist_ok=True)
|
||||
with tarfile.open(fileobj=io.BytesIO(plain), mode="r") as tar:
|
||||
names = tar.getnames()
|
||||
if any(n.startswith("/") or ".." in n.split("/") for n in names):
|
||||
raise SealError("bundle contains an unsafe path")
|
||||
inv_member = tar.extractfile("seal.manifest.json")
|
||||
if inv_member is None:
|
||||
raise SealError("bundle has no seal.manifest.json")
|
||||
inventory = json.loads(inv_member.read().decode("utf-8"))
|
||||
for name in names:
|
||||
data = tar.extractfile(name).read()
|
||||
if name != "seal.manifest.json":
|
||||
want = inventory["members"].get(name, {}).get("sha256")
|
||||
if want != hashlib.sha256(data).hexdigest():
|
||||
raise SealError("member %s does not match the inventory" % name)
|
||||
with open(os.path.join(dest_dir, name), "wb") as fh:
|
||||
fh.write(data)
|
||||
return inventory
|
||||
|
||||
|
||||
def _atomic_write(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
|
||||
@@ -0,0 +1,425 @@
|
||||
"""bastion agent query endpoint — the container on the mesh.
|
||||
|
||||
The other half of "access moves, data does not". Serves the owner's host
|
||||
and the manifest's grantees over MUTUAL TLS; nothing here ever reads the
|
||||
extract store, the vault, or the vectors' raw table for a client — the
|
||||
only bytes that leave are sanitized chunks and catalog cards.
|
||||
|
||||
Identity is the presented client certificate:
|
||||
- the TLS layer requires a client cert that chains to FREEAGENT_CA_ROOT
|
||||
(ssl.CERT_REQUIRED); no cert or a foreign CA never reaches a handler
|
||||
- the cert's CN is the caller. CN == manifest owner.host → owner.
|
||||
CN in manifest.grants (and today <= until) → grantee with that entry's
|
||||
ops. Anyone else who chained → 403, audited.
|
||||
|
||||
Ops (JSON over HTTPS):
|
||||
POST /search {query, k?} search (owner, grantee w/ op)
|
||||
POST /read_range {doc, start, end} read_range (owner, grantee w/ op)
|
||||
from the SANITIZED store only,
|
||||
<= READ_RANGE_MAX bytes
|
||||
GET /catalog[?prefix=&offset=&limit=] owner only: the file cards (paged)
|
||||
GET /catalog/dirs[?...] owner only: the directory cards (paged)
|
||||
GET /status owner only: run.state.json
|
||||
GET /health any chained cert: {ok, id}
|
||||
|
||||
The manifest is re-read and re-gated (signature verified) whenever its
|
||||
file changes: revoke a grant = drop the entry, re-sign, and the next
|
||||
request is refused. A manifest that no longer passes the gate makes the
|
||||
server refuse everything but /health until it does.
|
||||
|
||||
Every request is audited to <out>/serve.audit.jsonl: CN, op, status,
|
||||
bytes, a sha256 of the query (never the query), doc id for read_range.
|
||||
|
||||
Bind: --bind or FREEAGENT_BIND (host:port, default [::]:8470). The bastion
|
||||
binds 127.0.0.1:0 and reads the port from --port-file, so it is the only
|
||||
client; bind another interface only for another machine to reach. Server
|
||||
cert and key: FREEAGENT_TLS_CERT / FREEAGENT_TLS_KEY. Stdlib only; the embedding endpoint is the same sensor `search`
|
||||
already uses.
|
||||
"""
|
||||
|
||||
import datetime as _dt
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import ssl
|
||||
import sys
|
||||
import threading
|
||||
import urllib.parse
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
from . import embed as _embed
|
||||
from .manifest import ManifestError, load_manifest
|
||||
from .run import GateError, verify_signature
|
||||
from .sanitize import store_path as sanitized_store_path
|
||||
|
||||
DEFAULT_BIND = "[::]:8470"
|
||||
MAX_BODY = 64 * 1024
|
||||
READ_RANGE_MAX = 64 * 1024
|
||||
SEARCH_K_MAX = 50
|
||||
CATALOG_PAGE_DEFAULT = 500
|
||||
CATALOG_PAGE_MAX = 2000
|
||||
OWNER_OPS = ("search", "read_range", "catalog", "status")
|
||||
|
||||
|
||||
class ServeError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def parse_bind(bind):
|
||||
bind = bind.strip()
|
||||
if bind.startswith("["):
|
||||
host, _, port = bind[1:].partition("]")
|
||||
port = port.lstrip(":")
|
||||
else:
|
||||
host, _, port = bind.rpartition(":")
|
||||
if not host or not port.isdigit():
|
||||
raise ServeError("FREEAGENT_BIND must be host:port or [v6]:port, got %r" % bind)
|
||||
return host, int(port)
|
||||
|
||||
|
||||
def peer_cn(cert):
|
||||
for rdn in (cert or {}).get("subject", ()):
|
||||
for k, v in rdn:
|
||||
if k == "commonName":
|
||||
return v
|
||||
return None
|
||||
|
||||
|
||||
class Policy:
|
||||
"""Who may do what, from a gated manifest. Reloads on file change."""
|
||||
|
||||
def __init__(self, manifest_path, env, allow_unsigned=False, today=None):
|
||||
self.path = manifest_path
|
||||
self.env = env
|
||||
self.allow_unsigned = allow_unsigned
|
||||
self.today = today
|
||||
self._mtime = None
|
||||
self.manifest = None
|
||||
self.gate_error = None
|
||||
self._lock = threading.Lock()
|
||||
self.reload()
|
||||
|
||||
def reload(self):
|
||||
with self._lock:
|
||||
try:
|
||||
mtime = os.stat(self.path).st_mtime_ns
|
||||
except OSError as e:
|
||||
self.manifest, self.gate_error = None, "manifest unreadable: %s" % e
|
||||
return
|
||||
if mtime == self._mtime and (self.manifest or self.gate_error):
|
||||
return
|
||||
self._mtime = mtime
|
||||
try:
|
||||
m = load_manifest(self.path)
|
||||
except ManifestError as e:
|
||||
self.manifest, self.gate_error = None, "manifest invalid: %s" % "; ".join(e.errors)
|
||||
return
|
||||
except Exception as e: # noqa: BLE001
|
||||
# FAIL CLOSED here too: a manifest that cannot be read at all
|
||||
# (permissions, bad encoding) must not leave the previous one
|
||||
# in force once _mtime has advanced.
|
||||
self.manifest, self.gate_error = None, "manifest unreadable: %s" % e
|
||||
return
|
||||
try:
|
||||
verify_signature(m, self.env, manifest_dir=os.path.dirname(os.path.abspath(self.path)))
|
||||
except GateError as e:
|
||||
if not self.allow_unsigned:
|
||||
self.manifest, self.gate_error = None, "gate: %s" % e
|
||||
return
|
||||
except Exception as e: # noqa: BLE001
|
||||
# FAIL CLOSED: an unexpected error in the gate must never leave
|
||||
# the previous manifest in force (_mtime is already advanced).
|
||||
self.manifest, self.gate_error = None, "gate: %s" % e
|
||||
return
|
||||
self.manifest, self.gate_error = m, None
|
||||
|
||||
def allowed_ops(self, cn):
|
||||
"""Ops for this CN, or an empty tuple. Owner gets everything the
|
||||
mask permits; a grantee gets its entry's ops while unexpired."""
|
||||
self.reload()
|
||||
m = self.manifest
|
||||
if m is None or not cn:
|
||||
return ()
|
||||
if cn == m.owner_host:
|
||||
ops = list(OWNER_OPS)
|
||||
if "search" not in m.mask:
|
||||
ops.remove("search")
|
||||
return tuple(ops)
|
||||
today = self.today or _dt.datetime.now(_dt.timezone.utc).date()
|
||||
for g in m.grants:
|
||||
if g.host == cn:
|
||||
if today > g.until:
|
||||
return ()
|
||||
return tuple(g.ops)
|
||||
return ()
|
||||
|
||||
|
||||
HANDSHAKE_TIMEOUT = 10
|
||||
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
server_version = "freeagent"
|
||||
protocol_version = "HTTP/1.1"
|
||||
# keep-alive is one thread per connection here; an idle client must not
|
||||
# pin that thread forever (a bastion review measured exactly that)
|
||||
timeout = 60
|
||||
|
||||
def setup(self):
|
||||
# The TLS handshake happens HERE, in this connection's own thread —
|
||||
# never in the accept loop. Wrapping the listening socket (the first
|
||||
# cut) ran every handshake synchronously in accept(), so one idle
|
||||
# pre-auth TCP connection stalled every client.
|
||||
# The raw socket arrives with a short timeout from get_request, so a
|
||||
# peer that connects and says nothing is dropped after
|
||||
# HANDSHAKE_TIMEOUT, and only its own thread waited.
|
||||
self.request = self.server.ssl_ctx.wrap_socket(
|
||||
self.request, server_side=True, do_handshake_on_connect=False)
|
||||
self.request.do_handshake()
|
||||
self.connection = self.request
|
||||
super().setup()
|
||||
policy = None # injected by build_server
|
||||
out_dir = None
|
||||
env = None
|
||||
embed_endpoint = None
|
||||
audit_lock = threading.Lock()
|
||||
|
||||
# ── plumbing ────────────────────────────────────────────────────────
|
||||
def _cn(self):
|
||||
try:
|
||||
return peer_cn(self.connection.getpeercert())
|
||||
except (AttributeError, ValueError):
|
||||
return None
|
||||
|
||||
def _send(self, code, obj, cn=None, op=None, extra=None):
|
||||
body = json.dumps(obj, ensure_ascii=False).encode("utf-8")
|
||||
self.send_response(code)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
self._audit(cn, op, code, len(body), extra)
|
||||
|
||||
def _audit(self, cn, op, status, nbytes, extra):
|
||||
rec = {"ts": _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
|
||||
"cn": cn, "op": op, "path": self.path, "status": status, "bytes": nbytes}
|
||||
if extra:
|
||||
rec.update(extra)
|
||||
with self.audit_lock:
|
||||
with open(os.path.join(self.out_dir, "serve.audit.jsonl"), "a", encoding="utf-8") as fh:
|
||||
fh.write(json.dumps(rec, sort_keys=True) + "\n")
|
||||
|
||||
def _body(self):
|
||||
n = int(self.headers.get("Content-Length", "0") or 0)
|
||||
if n > MAX_BODY:
|
||||
raise ValueError("body too large")
|
||||
raw = self.rfile.read(n) if n else b"{}"
|
||||
obj = json.loads(raw.decode("utf-8"))
|
||||
if not isinstance(obj, dict):
|
||||
raise ValueError("body must be a JSON object")
|
||||
return obj
|
||||
|
||||
def _authorize(self, op):
|
||||
cn = self._cn()
|
||||
allowed = self.policy.allowed_ops(cn)
|
||||
if op in allowed:
|
||||
return cn, None
|
||||
if self.policy.gate_error:
|
||||
return cn, (503, {"error": "manifest not serving: %s" % self.policy.gate_error})
|
||||
return cn, (403, {"error": "not permitted for %s: %s" % (cn or "no-cn", op)})
|
||||
|
||||
def log_message(self, fmt, *args):
|
||||
pass # the audit file is the log; no request lines on stderr
|
||||
|
||||
# ── routes ───────────────────────────────────────────────────────────
|
||||
def do_GET(self):
|
||||
cn = self._cn()
|
||||
if self.path == "/health":
|
||||
m = self.policy.manifest
|
||||
self.policy.reload()
|
||||
self._send(200, {"ok": True, "id": m.id if m else None,
|
||||
"serving": self.policy.gate_error is None}, cn, "health")
|
||||
return
|
||||
route, _, qs = self.path.partition("?")
|
||||
if route in ("/catalog", "/catalog/dirs"):
|
||||
# /catalog = file cards; /catalog/dirs = directory cards — the MAP
|
||||
# (sizes, types, newest change, keywords, entity counts per folder)
|
||||
# a host plans a cleanup or an archive from without reading a file.
|
||||
# PAGED: ?prefix=<path prefix>&offset=<n>&limit=<n> — a big store's
|
||||
# file cards are megabytes and a client's reply cap is finite.
|
||||
cn, denied = self._authorize("catalog")
|
||||
if denied:
|
||||
self._send(denied[0], denied[1], cn, "catalog")
|
||||
return
|
||||
q = urllib.parse.parse_qs(qs)
|
||||
prefix = (q.get("prefix") or [""])[0]
|
||||
try:
|
||||
offset = max(0, int((q.get("offset") or ["0"])[0]))
|
||||
limit = max(1, min(CATALOG_PAGE_MAX, int((q.get("limit") or [str(CATALOG_PAGE_DEFAULT)])[0])))
|
||||
except ValueError:
|
||||
self._send(400, {"error": "offset/limit must be integers"}, cn, "catalog")
|
||||
return
|
||||
name = "catalog.dirs.jsonl" if route.endswith("/dirs") else "catalog.jsonl"
|
||||
path = os.path.join(self.out_dir, name)
|
||||
if not os.path.exists(path):
|
||||
self._send(404, {"error": "no catalog yet — the first sweep has not finished"}, cn, "catalog")
|
||||
return
|
||||
with open(path, encoding="utf-8") as fh:
|
||||
cards = [json.loads(l) for l in fh]
|
||||
if prefix:
|
||||
cards = [c for c in cards if str(c.get("path", "")).startswith(prefix)]
|
||||
total = len(cards)
|
||||
page = cards[offset:offset + limit]
|
||||
nxt = offset + limit if offset + limit < total else None
|
||||
self._send(200, {"cards": page, "kind": "dirs" if name.endswith("dirs.jsonl") else "files",
|
||||
"total": total, "offset": offset, "limit": limit, "next_offset": nxt,
|
||||
"prefix": prefix},
|
||||
cn, "catalog", {"kind": name, "n": len(page), "total": total})
|
||||
return
|
||||
if self.path == "/status":
|
||||
cn, denied = self._authorize("status")
|
||||
if denied:
|
||||
self._send(denied[0], denied[1], cn, "status")
|
||||
return
|
||||
path = os.path.join(self.out_dir, "run.state.json")
|
||||
state = {}
|
||||
if os.path.exists(path):
|
||||
with open(path, encoding="utf-8") as fh:
|
||||
state = json.load(fh)
|
||||
self._send(200, {"state": state}, cn, "status")
|
||||
return
|
||||
self._send(404, {"error": "not found"}, cn, None)
|
||||
|
||||
def do_POST(self):
|
||||
if self.path == "/search":
|
||||
cn, denied = self._authorize("search")
|
||||
if denied:
|
||||
self._send(denied[0], denied[1], cn, "search")
|
||||
return
|
||||
try:
|
||||
req = self._body()
|
||||
query = req["query"]
|
||||
k = int(req.get("k", 5))
|
||||
if not isinstance(query, str) or not query.strip():
|
||||
raise ValueError("query must be a non-empty string")
|
||||
k = max(1, min(SEARCH_K_MAX, k))
|
||||
except (ValueError, KeyError, TypeError) as e:
|
||||
self._send(400, {"error": "bad request: %s" % e}, cn, "search")
|
||||
return
|
||||
qhash = hashlib.sha256(query.encode("utf-8")).hexdigest()[:16]
|
||||
try:
|
||||
hits = _embed.search(self.policy.manifest, self.out_dir, query, k=k,
|
||||
env=self.env, endpoint=self.embed_endpoint)
|
||||
except (_embed.EmbedError, OSError) as e:
|
||||
self._send(502, {"error": "search unavailable: %s" % e}, cn, "search", {"q": qhash})
|
||||
return
|
||||
self._send(200, {"hits": hits, "k": k}, cn, "search", {"q": qhash, "hits": len(hits)})
|
||||
return
|
||||
if self.path == "/read_range":
|
||||
cn, denied = self._authorize("read_range")
|
||||
if denied:
|
||||
self._send(denied[0], denied[1], cn, "read_range")
|
||||
return
|
||||
try:
|
||||
req = self._body()
|
||||
doc = req["doc"]
|
||||
start = int(req.get("start", 0))
|
||||
end = int(req["end"])
|
||||
if not isinstance(doc, str) or len(doc) != 64 or not all(c in "0123456789abcdef" for c in doc):
|
||||
raise ValueError("doc must be a sanitized-document sha256")
|
||||
if start < 0 or end <= start:
|
||||
raise ValueError("start/end out of order")
|
||||
end = min(end, start + READ_RANGE_MAX)
|
||||
except (ValueError, KeyError, TypeError) as e:
|
||||
self._send(400, {"error": "bad request: %s" % e}, cn, "read_range")
|
||||
return
|
||||
path = sanitized_store_path(self.out_dir, doc) # the ONLY store with a route
|
||||
if not os.path.exists(path):
|
||||
self._send(404, {"error": "no such document"}, cn, "read_range", {"doc": doc[:12]})
|
||||
return
|
||||
with open(path, encoding="utf-8") as fh:
|
||||
text = fh.read()
|
||||
piece = text[start:end]
|
||||
self._send(200, {"doc": doc, "start": start, "end": start + len(piece),
|
||||
"total": len(text), "text": piece},
|
||||
cn, "read_range", {"doc": doc[:12], "chars": len(piece)})
|
||||
return
|
||||
self._send(404, {"error": "not found"}, self._cn(), None)
|
||||
|
||||
|
||||
class _Server(ThreadingHTTPServer):
|
||||
daemon_threads = True
|
||||
ssl_ctx = None # set by build_server; used by Handler.setup per connection
|
||||
|
||||
def get_request(self):
|
||||
# Plain TCP accept; the handshake is the handler thread's job (Handler.setup).
|
||||
sock, addr = self.socket.accept()
|
||||
sock.settimeout(HANDSHAKE_TIMEOUT)
|
||||
return sock, addr
|
||||
|
||||
def handle_error(self, request, client_address):
|
||||
# a client that fails the TLS handshake (no cert, foreign CA) is the
|
||||
# expected case, not a bug: one line, no traceback, no audit row
|
||||
# (it never became a request)
|
||||
exc = sys.exc_info()[1]
|
||||
if isinstance(exc, (ssl.SSLError, ConnectionError, OSError)):
|
||||
sys.stderr.write("freeagent serve: rejected %s: %s\n" % (client_address[0], exc))
|
||||
return
|
||||
super().handle_error(request, client_address)
|
||||
|
||||
|
||||
def build_ssl_context(cert, key, ca_root):
|
||||
for label, p in (("FREEAGENT_TLS_CERT", cert), ("FREEAGENT_TLS_KEY", key), ("FREEAGENT_CA_ROOT", ca_root)):
|
||||
if not p:
|
||||
raise ServeError("%s is not set — the endpoint is mTLS-only; refusing to serve" % label)
|
||||
if not os.path.exists(p):
|
||||
raise ServeError("%s=%s does not exist" % (label, p))
|
||||
ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
||||
ctx.minimum_version = ssl.TLSVersion.TLSv1_2
|
||||
ctx.verify_mode = ssl.CERT_REQUIRED
|
||||
ctx.load_verify_locations(ca_root)
|
||||
ctx.load_cert_chain(cert, key)
|
||||
return ctx
|
||||
|
||||
|
||||
def build_server(manifest_path, out_dir, env=None, allow_unsigned=False, bind=None,
|
||||
embed_endpoint=None, today=None, port_file=None):
|
||||
env = dict(os.environ) if env is None else env
|
||||
ctx = build_ssl_context(env.get("FREEAGENT_TLS_CERT", ""), env.get("FREEAGENT_TLS_KEY", ""),
|
||||
env.get("FREEAGENT_CA_ROOT", ""))
|
||||
policy = Policy(manifest_path, env, allow_unsigned=allow_unsigned, today=today)
|
||||
if policy.manifest is None and not allow_unsigned:
|
||||
raise ServeError("refusing to start: %s" % policy.gate_error)
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
handler = type("BoundHandler", (Handler,), {
|
||||
"policy": policy, "out_dir": out_dir, "env": env, "embed_endpoint": embed_endpoint})
|
||||
host, port = parse_bind(bind or env.get("FREEAGENT_BIND", DEFAULT_BIND))
|
||||
srv = _Server((host, port), handler, bind_and_activate=False)
|
||||
if ":" in host:
|
||||
import socket
|
||||
srv.address_family = socket.AF_INET6
|
||||
srv.socket = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
|
||||
srv.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
srv.server_bind()
|
||||
srv.server_activate()
|
||||
srv.ssl_ctx = ctx # per-connection wrap in Handler.setup, not on the listener
|
||||
srv.policy = policy
|
||||
if port_file:
|
||||
# a supervisor (the bastion) binds port 0 and reads the real port here
|
||||
tmp = port_file + ".tmp"
|
||||
with open(tmp, "w", encoding="utf-8") as fh:
|
||||
fh.write("%d\n" % srv.server_address[1])
|
||||
os.replace(tmp, port_file)
|
||||
return srv
|
||||
|
||||
|
||||
def serve(manifest_path, out_dir, env=None, allow_unsigned=False, bind=None, log=None, port_file=None):
|
||||
log = log or (lambda s: sys.stderr.write(s + "\n"))
|
||||
srv = build_server(manifest_path, out_dir, env, allow_unsigned=allow_unsigned, bind=bind,
|
||||
port_file=port_file)
|
||||
m = srv.policy.manifest
|
||||
log("freeagent serve: %s on %s (owner %s, %d grant(s), mTLS)" % (
|
||||
m.id if m else "?", srv.server_address[:2], m.owner_host if m else "?", len(m.grants) if m else 0))
|
||||
try:
|
||||
srv.serve_forever()
|
||||
finally:
|
||||
srv.server_close()
|
||||
@@ -0,0 +1,132 @@
|
||||
"""Tests for the zipapp build and the doctor. Run: python3 -m unittest discover ."""
|
||||
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
import zipfile
|
||||
|
||||
from . import build as bd
|
||||
from . import doctor as dr
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
# Repo root: bastion_agent/ is a top-level package.
|
||||
REPO = os.path.dirname(HERE)
|
||||
|
||||
|
||||
class BuildTests(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.tmp = tempfile.mkdtemp(prefix="fa-build-")
|
||||
cls.path, cls.digest, cls.size = bd.build(cls.tmp)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
shutil.rmtree(cls.tmp, ignore_errors=True)
|
||||
|
||||
def test_reproducible(self):
|
||||
self.assertEqual(bd.build_bytes(), bd.build_bytes())
|
||||
p2, d2, _ = bd.build(os.path.join(self.tmp, "again"))
|
||||
self.assertEqual(d2, self.digest)
|
||||
with open(self.path + ".sha256") as fh:
|
||||
self.assertTrue(fh.read().startswith(self.digest))
|
||||
|
||||
def test_contents(self):
|
||||
with open(self.path, "rb") as fh:
|
||||
data = fh.read()
|
||||
self.assertTrue(data.startswith(b"#!/usr/bin/env python3\n"))
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as zf:
|
||||
names = zf.namelist()
|
||||
times = {zf.getinfo(n).date_time for n in names}
|
||||
self.assertIn("__main__.py", names)
|
||||
self.assertIn("bastion_agent/__main__.py", names)
|
||||
self.assertIn("bastion_agent/serve.py", names)
|
||||
self.assertIn("bastion_agent/VERSION", names)
|
||||
self.assertIn("pii/pipeline.py", names)
|
||||
self.assertIn("yaml/__init__.py", names)
|
||||
self.assertIn("yaml/LICENSE-NOTICE.txt", names)
|
||||
self.assertFalse(any("test_" in n for n in names))
|
||||
self.assertFalse(any(n.endswith((".so", ".pyc", ".pyd")) for n in names))
|
||||
self.assertNotIn("bastion_agent/build.py", names)
|
||||
self.assertNotIn("bastion_agent/verify_release.py", names)
|
||||
self.assertNotIn("pii/Dockerfile", names)
|
||||
self.assertEqual(times, {bd.FIXED_TIME})
|
||||
self.assertLess(self.size, 1_000_000) # the ~10MB helper is the bar; this is well under
|
||||
|
||||
def run_pyz(self, *args, env=None, cwd=None):
|
||||
e = dict(os.environ)
|
||||
e.update(env or {})
|
||||
return subprocess.run([sys.executable, self.path] + list(args), capture_output=True,
|
||||
text=True, env=e, cwd=cwd or REPO)
|
||||
|
||||
def test_pyz_runs_version_doctor_validate(self):
|
||||
r = self.run_pyz("--version")
|
||||
self.assertEqual(r.returncode, 0, r.stderr)
|
||||
self.assertEqual(r.stdout.strip(), "bastion-agent " + bd.version())
|
||||
r = self.run_pyz("doctor", "--json")
|
||||
self.assertIn(r.returncode, (0, 1), r.stderr)
|
||||
results = json.loads(r.stdout)
|
||||
byname = {x["check"]: x for x in results}
|
||||
self.assertTrue(byname["python >= 3.8"]["ok"])
|
||||
self.assertTrue(byname["yaml (manifest loader)"]["ok"])
|
||||
r = self.run_pyz("validate", os.path.join(REPO, "manifests", "home-documents.yaml"))
|
||||
self.assertEqual(r.returncode, 0, r.stderr)
|
||||
self.assertIn("OK", r.stdout)
|
||||
|
||||
def test_pyz_runs_a_full_pipeline_outside_the_repo(self):
|
||||
# from a scratch directory, no repo on sys.path: walk -> ... -> catalog via the zipapp
|
||||
work = tempfile.mkdtemp(prefix="fa-pyzrun-")
|
||||
try:
|
||||
corpus = os.path.join(work, "c")
|
||||
os.makedirs(corpus)
|
||||
with open(os.path.join(corpus, "a.txt"), "w") as fh:
|
||||
fh.write("hello m@example.com\n")
|
||||
man = os.path.join(work, "m.yaml")
|
||||
with open(man, "w") as fh:
|
||||
fh.write("freeagent: 1\nid: z\nowner: {host: h}\nscope: {roots: ['%s']}\n"
|
||||
"mask: [read, index]\nstages: [walk, extract, sanitize, catalog]\n"
|
||||
"sanitize: {mode: flag}\nsweep: manual\n"
|
||||
"sink: {catalog: local, vectors: local, plaintext: never}\n" % corpus)
|
||||
r = self.run_pyz("sweep", man, "--out", os.path.join(work, "out"), "--once",
|
||||
"--allow-unsigned", cwd=work, env={"PYTHONPATH": ""})
|
||||
self.assertEqual(r.returncode, 0, r.stderr)
|
||||
self.assertIn("ok", r.stderr)
|
||||
with open(os.path.join(work, "out", "catalog.jsonl")) as fh:
|
||||
card = json.loads(fh.readline())
|
||||
self.assertIn("[PII_EMAIL_1]", card["head"])
|
||||
finally:
|
||||
shutil.rmtree(work, ignore_errors=True)
|
||||
|
||||
|
||||
class DefaultOutTests(unittest.TestCase):
|
||||
def test_single_root_lives_in_dot_agent(self):
|
||||
from .__main__ import default_out
|
||||
from .manifest import validate_document
|
||||
m = validate_document({"freeagent": 1, "id": "d", "owner": {"host": "h"},
|
||||
"scope": {"roots": ["/srv/docs"]}, "mask": ["read"], "stages": ["walk"],
|
||||
"sweep": "manual", "sink": {"catalog": "local", "vectors": "local", "plaintext": "never"}})
|
||||
self.assertEqual(default_out(m), os.path.join(os.path.abspath("/srv/docs"), ".agent"))
|
||||
m2 = validate_document(dict(m.raw, scope={"roots": ["/a", "/b"]}))
|
||||
self.assertEqual(default_out(m2), os.path.join("build", "freeagent", "d"))
|
||||
|
||||
|
||||
class DoctorTests(unittest.TestCase):
|
||||
def test_checks_and_render(self):
|
||||
res = dr.checks(env={"EMBED_URL": "http://127.0.0.1:9", "FREEAGENT_CA_ROOT": "/nonexistent"})
|
||||
byname = {x["check"]: x for x in res}
|
||||
self.assertTrue(byname["python >= 3.8"]["ok"])
|
||||
self.assertFalse(byname["embedding endpoint http://127.0.0.1:9"]["ok"])
|
||||
self.assertFalse(byname["embedding endpoint http://127.0.0.1:9"]["required"])
|
||||
self.assertFalse(byname["FREEAGENT_CA_ROOT"]["ok"])
|
||||
self.assertIn("missing file", byname["FREEAGENT_CA_ROOT"]["detail"])
|
||||
text = dr.render(res)
|
||||
self.assertIn("python >= 3.8", text)
|
||||
self.assertTrue(text.endswith("pass") or "required missing" in text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,246 @@
|
||||
"""Unit tests for the catalog stage. Run: python3 -m unittest discover .
|
||||
|
||||
The summarizer is a loopback stub (same process, test only) speaking both
|
||||
chat wire formats and echoing a deterministic summary built from the
|
||||
prompt, so cache and fail-closed behaviour are checked without a model.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
from . import catalog as cat
|
||||
from .extract import extract
|
||||
from .manifest import validate_document
|
||||
from .sanitize import sanitize
|
||||
from .walk import walk
|
||||
|
||||
|
||||
def manifest(root, summarize=None):
|
||||
doc = {
|
||||
"freeagent": 1, "id": "t", "owner": {"host": "h"},
|
||||
"scope": {"roots": [root]},
|
||||
"mask": ["read", "search", "index"] + (["summarize"] if summarize else []),
|
||||
"stages": ["walk", "extract", "sanitize", "catalog"], "sweep": "manual",
|
||||
"sanitize": {"mode": "flag"},
|
||||
"sink": {"catalog": "local", "vectors": "local", "plaintext": "never"},
|
||||
}
|
||||
if summarize:
|
||||
doc["summarize"] = {"model": "stub-chat", "per": summarize}
|
||||
return validate_document(doc)
|
||||
|
||||
|
||||
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 _StubChat(BaseHTTPRequestHandler):
|
||||
serve = ("ollama", "openai")
|
||||
calls = []
|
||||
fail = False
|
||||
|
||||
def do_POST(self):
|
||||
n = int(self.headers.get("Content-Length", "0"))
|
||||
req = json.loads(self.rfile.read(n))
|
||||
prompt = req["messages"][-1]["content"]
|
||||
_StubChat.calls.append((self.path, req["model"], prompt))
|
||||
if _StubChat.fail or (self.path == "/api/chat" and "ollama" not in self.serve) \
|
||||
or (self.path == "/v1/chat/completions" and "openai" not in self.serve) \
|
||||
or self.path not in ("/api/chat", "/v1/chat/completions"):
|
||||
self.send_response(500 if _StubChat.fail else 404)
|
||||
self.send_header("Content-Length", "0")
|
||||
self.end_headers()
|
||||
return
|
||||
body_text = prompt.split("\n\n", 1)[-1]
|
||||
summary = '<think>hmm</think> "SUMMARY of: %s"' % body_text[:40].replace("\n", " ")
|
||||
if self.path == "/api/chat":
|
||||
body = {"message": {"role": "assistant", "content": summary}}
|
||||
else:
|
||||
body = {"choices": [{"message": {"role": "assistant", "content": summary}}]}
|
||||
data = json.dumps(body).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
|
||||
def log_message(self, *a):
|
||||
pass
|
||||
|
||||
|
||||
class UnitTests(unittest.TestCase):
|
||||
def test_title_head_keywords(self):
|
||||
text = "# Quarterly Budget Review\n\nSpending on routers and firewalls rose. Routers again.\n"
|
||||
self.assertEqual(cat.title_of(text), "Quarterly Budget Review")
|
||||
self.assertTrue(cat.head_of(text).startswith("# Quarterly Budget Review Spending"))
|
||||
counts = cat.term_counts(text)
|
||||
self.assertEqual(counts["routers"], 2)
|
||||
self.assertNotIn("and", counts)
|
||||
kws = cat.keywords(counts, {"routers": 1}, 1, 3)
|
||||
self.assertEqual(kws[0], "routers")
|
||||
self.assertEqual(kws, cat.keywords(counts, {"routers": 1}, 1, 3)) # deterministic
|
||||
|
||||
def test_mask_name(self):
|
||||
detect = cat._rules_detect()
|
||||
# the email rule legitimately spans dots and hyphens, so a name that
|
||||
# embeds an address is masked conservatively — whole span, not just
|
||||
# the address; the card keeps mime/size so nothing is lost but the name
|
||||
self.assertEqual(cat.mask_name("invoice-m@example.com.pdf", detect), "[PII_EMAIL]")
|
||||
self.assertEqual(cat.mask_name("invoice m@example.com.pdf", detect), "invoice [PII_EMAIL]")
|
||||
self.assertEqual(cat.mask_name("notes.txt", detect), "notes.txt")
|
||||
self.assertEqual(cat.mask_name("call 415-555-0173 notes", detect), "call [PII_PHONE] notes")
|
||||
|
||||
def test_clean_summary(self):
|
||||
self.assertEqual(cat.clean_summary('<think>x</think> "A b" '), "A b")
|
||||
self.assertEqual(len(cat.clean_summary("x" * 2000)), cat.SUMMARY_MAX_CHARS)
|
||||
|
||||
|
||||
class Case(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.srv = ThreadingHTTPServer(("127.0.0.1", 0), _StubChat)
|
||||
threading.Thread(target=cls.srv.serve_forever, daemon=True).start()
|
||||
cls.url = "http://127.0.0.1:%d" % cls.srv.server_address[1]
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls.srv.shutdown()
|
||||
cls.srv.server_close()
|
||||
|
||||
def setUp(self):
|
||||
self.root = tempfile.mkdtemp(prefix="fa-cat-")
|
||||
self.out = tempfile.mkdtemp(prefix="fa-catout-")
|
||||
_StubChat.serve, _StubChat.calls, _StubChat.fail = ("ollama", "openai"), [], False
|
||||
put(self.root, "finance/budget.md", "# Budget 2026\n\nRouters and firewalls. Contact m@example.com.\n")
|
||||
put(self.root, "finance/invoice-m@example.com.txt", "invoice for routers 4111 1111 1111 1111\n")
|
||||
put(self.root, "notes/orchard.txt", "apples pears plums harvest\n")
|
||||
put(self.root, "blob.bin", "\x00\x01")
|
||||
self.env = {"SUMMARIZE_URL": self.url}
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.root, ignore_errors=True)
|
||||
shutil.rmtree(self.out, ignore_errors=True)
|
||||
|
||||
def run_stage(self, m):
|
||||
walk(m, self.out)
|
||||
extract(m, self.out, use_external=False)
|
||||
sanitize(m, self.out, env={})
|
||||
s = cat.catalog(m, self.out, env=self.env)
|
||||
with open(os.path.join(self.out, "catalog.jsonl"), encoding="utf-8") as fh:
|
||||
files = {json.loads(l)["path"]: json.loads(l) for l in fh}
|
||||
with open(os.path.join(self.out, "catalog.dirs.jsonl"), encoding="utf-8") as fh:
|
||||
dirs = {json.loads(l)["path"]: json.loads(l) for l in fh}
|
||||
return s, files, dirs
|
||||
|
||||
|
||||
class StageTests(Case):
|
||||
def test_cards_without_summarizer(self):
|
||||
s, files, dirs = self.run_stage(manifest(self.root))
|
||||
self.assertIsNone(s["summarizer"])
|
||||
self.assertEqual(_StubChat.calls, [])
|
||||
b = files["finance/budget.md"]
|
||||
self.assertEqual(b["title"], "Budget 2026")
|
||||
self.assertIn("[PII_EMAIL_1]", b["head"])
|
||||
self.assertNotIn("m@example.com", json.dumps(b))
|
||||
self.assertEqual(b["entities"], {"EMAIL": 1})
|
||||
self.assertIn("routers", b["keywords"])
|
||||
self.assertEqual(b["text"], "ok")
|
||||
self.assertEqual(b["mime"], "text/markdown")
|
||||
self.assertNotIn("summary", b)
|
||||
# filename with an email is masked in the card
|
||||
self.assertIn("finance/[PII_EMAIL]", files) # conservative whole-span mask
|
||||
self.assertEqual(files["finance/[PII_EMAIL]"]["mime"], "text/plain")
|
||||
self.assertEqual(s["counts"]["names_masked"], 1)
|
||||
# placeholders never become keywords
|
||||
self.assertNotIn("pii", b["keywords"])
|
||||
self.assertNotIn("email", b["keywords"])
|
||||
# a file with no text still gets a card
|
||||
self.assertEqual(files["blob.bin"]["text"], "unsupported")
|
||||
self.assertNotIn("head", files["blob.bin"])
|
||||
# directories: root "" aggregates everything
|
||||
self.assertEqual(dirs[""]["files"], 4)
|
||||
self.assertEqual(dirs["finance"]["files"], 2)
|
||||
self.assertEqual(dirs["finance"]["with_text"], 2)
|
||||
self.assertEqual(dirs["finance"]["entities"], {"EMAIL": 1, "CREDIT_CARD": 1})
|
||||
self.assertIn("routers", dirs["finance"]["keywords"])
|
||||
self.assertIn("apples", dirs["notes"]["keywords"])
|
||||
self.assertEqual(dirs["finance"]["types"]["text/markdown"], 1)
|
||||
self.assertEqual(s["counts"], dict(files=4, with_text=3, dirs=3, summaries=0,
|
||||
summaries_reused=0, summary_errors=0, names_masked=1))
|
||||
|
||||
def test_parent_directory_names_are_masked_too(self):
|
||||
# a directory named with an address must not leak through
|
||||
# a file card's path, which is what seal ships
|
||||
put(self.root, "marisol.vega@example.com/notes.txt", "plain\n")
|
||||
s, files, dirs = self.run_stage(manifest(self.root))
|
||||
self.assertIn("[PII_EMAIL]/notes.txt", files)
|
||||
self.assertIn("[PII_EMAIL]", dirs)
|
||||
blob = json.dumps(files) + json.dumps(dirs)
|
||||
self.assertNotIn("marisol.vega", blob)
|
||||
self.assertEqual(s["counts"]["names_masked"], 2) # the invoice file + the directory
|
||||
|
||||
def test_catalog_carries_no_values(self):
|
||||
_, files, dirs = self.run_stage(manifest(self.root))
|
||||
blob = json.dumps(files) + json.dumps(dirs)
|
||||
for leak in ("m@example.com", "4111"):
|
||||
self.assertNotIn(leak, blob)
|
||||
|
||||
def test_summaries_file_and_directory_with_cache(self):
|
||||
m = manifest(self.root, summarize=["file", "directory"])
|
||||
s, files, dirs = self.run_stage(m)
|
||||
self.assertTrue(s["summarizer"].startswith("stub-chat@"))
|
||||
self.assertTrue(s["summarizer"].endswith("(ollama)"))
|
||||
self.assertTrue(files["finance/budget.md"]["summary"].startswith("SUMMARY of:"))
|
||||
self.assertNotIn("<think>", files["finance/budget.md"]["summary"])
|
||||
self.assertIsNone(files["blob.bin"]["summary"])
|
||||
self.assertTrue(dirs["finance"]["summary"].startswith("SUMMARY of:"))
|
||||
self.assertEqual(s["counts"]["summaries"], 3 + 3) # 3 texts + 3 dirs
|
||||
# the model only ever saw sanitized text
|
||||
for _, _, prompt in _StubChat.calls:
|
||||
self.assertNotIn("m@example.com", prompt)
|
||||
self.assertNotIn("4111", prompt)
|
||||
# second run: all summaries from cache, no calls
|
||||
_StubChat.calls = []
|
||||
s2, _, _ = self.run_stage(m)
|
||||
self.assertEqual((s2["counts"]["summaries"], s2["counts"]["summaries_reused"]), (0, 6))
|
||||
self.assertEqual(_StubChat.calls, [])
|
||||
|
||||
def test_openai_autodetect(self):
|
||||
_StubChat.serve = ("openai",)
|
||||
s, files, _ = self.run_stage(manifest(self.root, summarize=["file"]))
|
||||
self.assertTrue(s["summarizer"].endswith("(openai)"))
|
||||
self.assertTrue(files["notes/orchard.txt"]["summary"])
|
||||
|
||||
def test_summarizer_failure_stops_but_cards_are_written(self):
|
||||
_StubChat.fail = True
|
||||
s, files, dirs = self.run_stage(manifest(self.root, summarize=["file", "directory"]))
|
||||
self.assertIsNotNone(s["stopped"])
|
||||
self.assertEqual(s["counts"]["summary_errors"], 1) # stopped after the first
|
||||
self.assertEqual(len(files), 4)
|
||||
self.assertEqual(len(dirs), 3)
|
||||
errs = [c for c in list(files.values()) + list(dirs.values()) if c.get("summary_error")]
|
||||
self.assertEqual(len(errs), 1)
|
||||
self.assertTrue(all(c.get("summary") is None for c in files.values()))
|
||||
|
||||
def test_unreachable_summarizer(self):
|
||||
m = manifest(self.root, summarize=["file"])
|
||||
walk(m, self.out)
|
||||
extract(m, self.out, use_external=False)
|
||||
sanitize(m, self.out, env={})
|
||||
s = cat.catalog(m, self.out, env={"SUMMARIZE_URL": "http://127.0.0.1:9"})
|
||||
self.assertIn("summarizer", s["stopped"])
|
||||
|
||||
def test_requires_earlier_stages(self):
|
||||
with self.assertRaises(FileNotFoundError):
|
||||
cat.catalog(manifest(self.root), self.out, env=self.env)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,262 @@
|
||||
"""Unit tests for chunk / embed / search. Run: python3 -m unittest discover .
|
||||
|
||||
The embedding endpoint is a stub HTTP server on loopback (same process,
|
||||
test only) speaking both wire formats with a deterministic bag-of-words
|
||||
vector, so search ranking is meaningful without a model.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
from . import embed as em
|
||||
from .extract import extract
|
||||
from .manifest import validate_document
|
||||
from .sanitize import sanitize
|
||||
from .walk import walk
|
||||
|
||||
|
||||
def manifest(root, model="stub-model", dim=64, size=8, overlap=2):
|
||||
return validate_document({
|
||||
"freeagent": 1, "id": "t", "owner": {"host": "h"},
|
||||
"scope": {"roots": [root]},
|
||||
"mask": ["read", "search", "index"],
|
||||
"stages": ["walk", "extract", "sanitize", "chunk", "embed"], "sweep": "manual",
|
||||
"sanitize": {"mode": "flag"},
|
||||
"embed": {"model": model, "dim": dim, "chunk": {"size_tokens": size, "overlap_tokens": overlap}},
|
||||
"sink": {"catalog": "local", "vectors": "local", "plaintext": "never"},
|
||||
})
|
||||
|
||||
|
||||
def bow_vector(text, dim):
|
||||
v = [0.0] * dim
|
||||
for w in re.findall(r"\w+", text.lower()):
|
||||
h = int(hashlib.md5(w.encode()).hexdigest(), 16)
|
||||
v[h % dim] += 1.0
|
||||
return v
|
||||
|
||||
|
||||
class _Stub(BaseHTTPRequestHandler):
|
||||
dim = 64
|
||||
serve = ("ollama", "openai")
|
||||
calls = []
|
||||
|
||||
def do_POST(self):
|
||||
n = int(self.headers.get("Content-Length", "0"))
|
||||
req = json.loads(self.rfile.read(n))
|
||||
_Stub.calls.append((self.path, req["model"], len(req["input"])))
|
||||
inputs = req["input"] if isinstance(req["input"], list) else [req["input"]]
|
||||
if self.path == "/api/embed" and "ollama" in self.serve:
|
||||
body = {"embeddings": [bow_vector(t, self.dim) for t in inputs]}
|
||||
elif self.path == "/v1/embeddings" and "openai" in self.serve:
|
||||
body = {"data": [{"index": i, "embedding": bow_vector(t, self.dim)}
|
||||
for i, t in enumerate(inputs)]}
|
||||
else:
|
||||
self.send_response(404)
|
||||
self.send_header("Content-Length", "0")
|
||||
self.end_headers()
|
||||
return
|
||||
data = json.dumps(body).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
|
||||
def log_message(self, *a):
|
||||
pass
|
||||
|
||||
|
||||
class ChunkUnitTests(unittest.TestCase):
|
||||
def test_windows_cover_text_and_overlap(self):
|
||||
text = " ".join("w%d" % i for i in range(20))
|
||||
wins = list(em.chunk_text(text, 8, 2))
|
||||
self.assertEqual(wins[0][0], 0)
|
||||
self.assertEqual(wins[-1][1], len(text))
|
||||
# 20 tokens, step 6: windows at 0, 6, 12 — the third reaches the end
|
||||
self.assertEqual([em.count_tokens(w[2]) for w in wins], [8, 8, 8])
|
||||
# step = 6: window 2 starts at token 6 = "w6"
|
||||
self.assertTrue(wins[1][2].startswith("w6 "))
|
||||
# every char is inside at least one window
|
||||
covered = set()
|
||||
for s, e, _ in wins:
|
||||
covered.update(range(s, e))
|
||||
self.assertEqual(len(covered), len(text))
|
||||
|
||||
def test_short_and_empty(self):
|
||||
self.assertEqual(list(em.chunk_text("", 8, 2)), [])
|
||||
self.assertEqual(list(em.chunk_text(" \n ", 8, 2)), [])
|
||||
wins = list(em.chunk_text("just three tokens", 8, 2))
|
||||
self.assertEqual(len(wins), 1)
|
||||
self.assertEqual(wins[0][2], "just three tokens")
|
||||
|
||||
def test_punctuation_is_a_token_and_ids_are_content_addressed(self):
|
||||
self.assertEqual(em.count_tokens("Hi, there!"), 4)
|
||||
self.assertEqual(em.chunk_id("d", "x"), em.chunk_id("d", "x"))
|
||||
self.assertNotEqual(em.chunk_id("d", "x"), em.chunk_id("e", "x"))
|
||||
|
||||
def test_pack_roundtrip(self):
|
||||
v = [0.5, -1.25, 3.0]
|
||||
self.assertEqual(list(em.unpack(em.pack(v))), v)
|
||||
self.assertAlmostEqual(em.cosine([1, 0], [1, 0]), 1.0)
|
||||
self.assertAlmostEqual(em.cosine([1, 0], [0, 1]), 0.0)
|
||||
self.assertEqual(em.cosine([0, 0], [1, 1]), 0.0)
|
||||
|
||||
|
||||
class Case(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.srv = ThreadingHTTPServer(("127.0.0.1", 0), _Stub)
|
||||
threading.Thread(target=cls.srv.serve_forever, daemon=True).start()
|
||||
cls.url = "http://127.0.0.1:%d" % cls.srv.server_address[1]
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls.srv.shutdown()
|
||||
cls.srv.server_close()
|
||||
|
||||
def setUp(self):
|
||||
self.root = tempfile.mkdtemp(prefix="fa-em-")
|
||||
self.out = tempfile.mkdtemp(prefix="fa-emout-")
|
||||
_Stub.dim, _Stub.serve, _Stub.calls = 64, ("ollama", "openai"), []
|
||||
self.put("apples.txt", "apples pears plums. " * 12 + "orchard harvest.")
|
||||
self.put("routers.txt", "router switch firewall vlan. " * 12 + "network segment.")
|
||||
self.put("twin.txt", "apples pears plums. " * 12 + "orchard harvest.")
|
||||
self.env = {"EMBED_URL": self.url, "EMBED_BATCH": "4"}
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.root, ignore_errors=True)
|
||||
shutil.rmtree(self.out, ignore_errors=True)
|
||||
|
||||
def put(self, name, text):
|
||||
with open(os.path.join(self.root, name), "w", encoding="utf-8") as fh:
|
||||
fh.write(text)
|
||||
|
||||
def prep(self, m):
|
||||
walk(m, self.out)
|
||||
extract(m, self.out, use_external=False)
|
||||
sanitize(m, self.out, env={})
|
||||
|
||||
def chunks(self):
|
||||
with open(os.path.join(self.out, "chunks.jsonl"), encoding="utf-8") as fh:
|
||||
return [json.loads(l) for l in fh]
|
||||
|
||||
def db(self, sql, *params):
|
||||
conn = sqlite3.connect(em.vectors_path(self.out))
|
||||
try:
|
||||
return conn.execute(sql, params).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class StageTests(Case):
|
||||
def test_chunk_stage(self):
|
||||
m = manifest(self.root)
|
||||
self.prep(m)
|
||||
s = em.chunk(m, self.out)
|
||||
rows = self.chunks()
|
||||
self.assertEqual(s["counts"]["docs"], 2) # twin shares content with apples
|
||||
self.assertEqual(s["counts"]["chunks"], len(rows))
|
||||
self.assertGreater(len(rows), 4)
|
||||
docs = {r["doc"] for r in rows}
|
||||
self.assertEqual(len(docs), 2)
|
||||
self.assertEqual(rows, sorted(rows, key=lambda r: (r["doc"], r["idx"])))
|
||||
# deterministic
|
||||
em.chunk(m, self.out)
|
||||
self.assertEqual(rows, self.chunks())
|
||||
|
||||
def test_chunk_requires_sanitize(self):
|
||||
with self.assertRaises(FileNotFoundError):
|
||||
em.chunk(manifest(self.root), self.out)
|
||||
|
||||
def test_embed_probe_refuses_wrong_dim(self):
|
||||
m = manifest(self.root, dim=32) # stub serves 64
|
||||
self.prep(m)
|
||||
em.chunk(m, self.out)
|
||||
with self.assertRaises(em.EmbedError) as cm:
|
||||
em.embed(m, self.out, env=self.env)
|
||||
self.assertIn("serves 64-dim", str(cm.exception))
|
||||
self.assertFalse(os.path.exists(em.vectors_path(self.out))) # refused before opening the store
|
||||
|
||||
def test_embed_unreachable_fails_closed(self):
|
||||
m = manifest(self.root)
|
||||
self.prep(m)
|
||||
em.chunk(m, self.out)
|
||||
with self.assertRaises(em.EmbedError):
|
||||
em.embed(m, self.out, env={"EMBED_URL": "http://127.0.0.1:9"})
|
||||
|
||||
def test_embed_writes_reuses_and_keeps_old_models(self):
|
||||
m = manifest(self.root)
|
||||
self.prep(m)
|
||||
em.chunk(m, self.out)
|
||||
s1 = em.embed(m, self.out, env=self.env)
|
||||
n = s1["counts"]["unique"] # repeated paragraphs share a chunk id
|
||||
self.assertLess(n, s1["counts"]["chunks"])
|
||||
self.assertEqual(s1["counts"]["embedded"], n)
|
||||
self.assertEqual(s1["endpoint"], "stub-model@%s(ollama)" % self.url)
|
||||
# batches of 4 + one probe call
|
||||
embed_calls = [c for c in _Stub.calls if c[0] == "/api/embed"]
|
||||
self.assertEqual(sum(c[2] for c in embed_calls), n + 1)
|
||||
self.assertTrue(all(c[2] <= 4 for c in embed_calls))
|
||||
# second run: nothing re-embedded
|
||||
s2 = em.embed(m, self.out, env=self.env)
|
||||
self.assertEqual((s2["counts"]["embedded"], s2["counts"]["reused"]), (0, n))
|
||||
# a second model ADDS rows; the first model's rows survive
|
||||
m2 = manifest(self.root, model="other-model")
|
||||
s3 = em.embed(m2, self.out, env=self.env)
|
||||
self.assertEqual(s3["counts"]["embedded"], n)
|
||||
self.assertEqual(sorted(x["model"] for x in s3["models_in_store"]), ["other-model", "stub-model"])
|
||||
self.assertEqual(self.db("SELECT COUNT(*) FROM vectors WHERE model='stub-model'")[0][0], n)
|
||||
|
||||
def test_stale_vectors_are_kept_and_counted(self):
|
||||
m = manifest(self.root)
|
||||
self.prep(m)
|
||||
em.chunk(m, self.out)
|
||||
em.embed(m, self.out, env=self.env)
|
||||
before = self.db("SELECT COUNT(*) FROM vectors")[0][0]
|
||||
self.put("routers.txt", "completely different content now.")
|
||||
self.prep(m)
|
||||
em.chunk(m, self.out)
|
||||
s = em.embed(m, self.out, env=self.env)
|
||||
self.assertGreater(s["stale_vectors"], 0)
|
||||
self.assertGreaterEqual(self.db("SELECT COUNT(*) FROM vectors")[0][0], before)
|
||||
# unchanged doc's chunks were reused, not re-embedded
|
||||
self.assertGreater(s["counts"]["reused"], 0)
|
||||
|
||||
def test_openai_autodetect_and_forced_api(self):
|
||||
_Stub.serve = ("openai",)
|
||||
m = manifest(self.root)
|
||||
self.prep(m)
|
||||
em.chunk(m, self.out)
|
||||
s = em.embed(m, self.out, env=self.env)
|
||||
self.assertTrue(s["endpoint"].endswith("(openai)"))
|
||||
with self.assertRaises(em.EmbedError):
|
||||
em.embed(m, self.out, env=dict(self.env, EMBED_API="ollama"))
|
||||
with self.assertRaises(em.EmbedError):
|
||||
em.build_endpoint(m, dict(self.env, EMBED_API="magic"))
|
||||
|
||||
def test_search_ranks_the_right_document(self):
|
||||
m = manifest(self.root)
|
||||
self.prep(m)
|
||||
em.chunk(m, self.out)
|
||||
em.embed(m, self.out, env=self.env)
|
||||
hits = em.search(m, self.out, "firewall vlan router", k=3, env=self.env)
|
||||
self.assertEqual(len(hits), 3)
|
||||
self.assertIn("router", hits[0]["text"])
|
||||
self.assertNotIn("apples", hits[0]["text"])
|
||||
hits = em.search(m, self.out, "pears and plums", k=1, env=self.env)
|
||||
self.assertIn("apples", hits[0]["text"])
|
||||
self.assertTrue(0 < hits[0]["score"] <= 1.0)
|
||||
# a model with no vectors yields nothing, not an error
|
||||
self.assertEqual(em.search(manifest(self.root, model="nope"), self.out, "x", env=self.env), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,438 @@
|
||||
"""Unit tests for the extract stage. Run: python3 -m unittest discover .
|
||||
|
||||
Every fixture is generated here with the stdlib (zip + XML for office
|
||||
formats, hand-built PDFs), so the suite needs no binary blobs and runs on
|
||||
any OS. External tools are mocked; the real-tool path is exercised only
|
||||
where the machine happens to have them.
|
||||
"""
|
||||
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import unittest
|
||||
import zipfile
|
||||
import zlib
|
||||
from unittest import mock
|
||||
|
||||
from . import extract as ex
|
||||
from .manifest import validate_document
|
||||
from .walk import walk
|
||||
|
||||
|
||||
def manifest(root, **scope):
|
||||
return validate_document({
|
||||
"freeagent": 1, "id": "t", "owner": {"host": "h"},
|
||||
"scope": dict({"roots": [root]}, **scope),
|
||||
"mask": ["read"], "stages": ["walk", "extract"], "sweep": "manual",
|
||||
"sink": {"catalog": "local", "vectors": "local", "plaintext": "never"},
|
||||
})
|
||||
|
||||
|
||||
def put(root, rel, data):
|
||||
path = os.path.join(root, *rel.split("/"))
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
with open(path, "wb") as fh:
|
||||
fh.write(data if isinstance(data, bytes) else data.encode("utf-8"))
|
||||
|
||||
|
||||
def zipped(members):
|
||||
buf = io.BytesIO()
|
||||
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
|
||||
for name, data in members.items():
|
||||
zf.writestr(name, data)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
DOCX_XML = ('<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">'
|
||||
'<w:body><w:p><w:r><w:t>Hello</w:t></w:r><w:r><w:tab/><w:t>World</w:t></w:r></w:p>'
|
||||
'<w:p><w:r><w:t>Second paragraph</w:t></w:r></w:p></w:body></w:document>')
|
||||
XLSX_SST = ('<sst xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">'
|
||||
'<si><t>Name</t></si><si><t>Ada</t></si></sst>')
|
||||
XLSX_SHEET = ('<worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">'
|
||||
'<sheetData><row><c t="s"><v>0</v></c><c><v>42</v></c></row>'
|
||||
'<row><c t="s"><v>1</v></c><c t="inlineStr"><is><t>inline</t></is></c></row>'
|
||||
'</sheetData></worksheet>')
|
||||
PPTX_SLIDE = ('<p:sld xmlns:p="http://schemas.openxmlformats.org/presentationml/2006/main" '
|
||||
'xmlns:a="http://schemas.openxmlformats.org/drawingml/2006/main">'
|
||||
'<p:cSld><p:spTree><p:sp><p:txBody><a:p><a:r><a:t>Slide </a:t></a:r>'
|
||||
'<a:r><a:t>title</a:t></a:r></a:p></p:txBody></p:sp></p:spTree></p:cSld></p:sld>')
|
||||
ODT_CONTENT = ('<office:document-content xmlns:office="urn:oasis:names:tc:opendocument:xmlns:office:1.0" '
|
||||
'xmlns:text="urn:oasis:names:tc:opendocument:xmlns:text:1.0" '
|
||||
'xmlns:table="urn:oasis:names:tc:opendocument:xmlns:table:1.0"><office:body><office:text>'
|
||||
'<text:h>Heading</text:h><text:p>Body <text:span>span</text:span></text:p>'
|
||||
'<table:table><table:table-row><table:table-cell><text:p>c1</text:p></table:table-cell>'
|
||||
'<table:table-cell><text:p>c2</text:p></table:table-cell></table:table-row></table:table>'
|
||||
'</office:text></office:body></office:document-content>')
|
||||
|
||||
|
||||
def pdf_bytes(content, compress=False):
|
||||
stream = zlib.compress(content) if compress else content
|
||||
filt = b"/Filter /FlateDecode " if compress else b""
|
||||
return (b"%PDF-1.4\n1 0 obj << /Type /Catalog >> endobj\n"
|
||||
b"2 0 obj << " + filt + b"/Length " + str(len(stream)).encode() + b" >>\nstream\n"
|
||||
+ stream + b"\nendstream\nendobj\n%%EOF\n")
|
||||
|
||||
|
||||
class ExtractorUnitTests(unittest.TestCase):
|
||||
def test_decode_bytes(self):
|
||||
self.assertEqual(ex.decode_bytes("héllo".encode("utf-8")), ("héllo", "utf-8"))
|
||||
self.assertEqual(ex.decode_bytes(b"\xef\xbb\xbfhi"), ("hi", "utf-8-sig"))
|
||||
self.assertEqual(ex.decode_bytes("hi".encode("utf-16"))[0], "hi")
|
||||
self.assertEqual(ex.decode_bytes("hi there".encode("utf-16-le")), ("hi there", "utf-16-le"))
|
||||
self.assertEqual(ex.decode_bytes(b"caf\xe9"), ("café", "cp1252"))
|
||||
|
||||
def test_html_to_text(self):
|
||||
out = ex.html_to_text("<html><head><title>T</title><style>x{}</style></head>"
|
||||
"<body><p>One & two</p><script>bad()</script><div>Three</div></body>")
|
||||
self.assertIn("One & two", out)
|
||||
self.assertIn("Three", out)
|
||||
self.assertNotIn("bad()", out)
|
||||
self.assertNotIn("x{}", out)
|
||||
|
||||
def test_docx(self):
|
||||
text = ex.extract_docx(zipped({"word/document.xml": DOCX_XML}))
|
||||
self.assertEqual(text.strip().splitlines(), ["Hello\tWorld", "Second paragraph"])
|
||||
|
||||
def test_xlsx(self):
|
||||
text = ex.extract_xlsx(zipped({"xl/sharedStrings.xml": XLSX_SST,
|
||||
"xl/worksheets/sheet1.xml": XLSX_SHEET}))
|
||||
self.assertIn("Name\t42", text)
|
||||
self.assertIn("Ada\tinline", text)
|
||||
|
||||
def test_pptx(self):
|
||||
text = ex.extract_pptx(zipped({"ppt/slides/slide2.xml": PPTX_SLIDE,
|
||||
"ppt/slides/slide1.xml": PPTX_SLIDE.replace("title", "first")}))
|
||||
lines = [l for l in text.splitlines() if l]
|
||||
self.assertEqual(lines, ["# slide1", "Slide first", "# slide2", "Slide title"])
|
||||
|
||||
def test_odf(self):
|
||||
text = ex.extract_odf(zipped({"content.xml": ODT_CONTENT}))
|
||||
self.assertEqual(text.splitlines(), ["Heading", "Body span", "c1\tc2"])
|
||||
|
||||
def test_zip_member_cap(self):
|
||||
big = zipped({"content.xml": b"x" * (ex.ZIP_MEMBER_CAP + 1)})
|
||||
with self.assertRaises(ValueError):
|
||||
ex.extract_odf(big)
|
||||
|
||||
def test_office_refuses_a_dtd(self):
|
||||
# Office XML never carries a DTD; one is how entity expansion gets in.
|
||||
bomb = (b'<?xml version="1.0"?><!DOCTYPE d [<!ENTITY a "aaaaaaaaaa">'
|
||||
b'<!ENTITY b "&a;&a;&a;&a;&a;&a;&a;&a;&a;&a;">]>'
|
||||
b'<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">'
|
||||
b'<w:body><w:p><w:r><w:t>&b;</w:t></w:r></w:p></w:body></w:document>')
|
||||
with self.assertRaises(ValueError):
|
||||
ex.extract_docx(zipped({"word/document.xml": bomb}))
|
||||
|
||||
def test_office_text_cap_stops_with_the_text_so_far(self):
|
||||
rows = "".join('<row><c t="inlineStr"><is><t>row %d</t></is></c></row>' % i for i in range(200))
|
||||
sheet = ('<worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">'
|
||||
"<sheetData>%s</sheetData></worksheet>" % rows).encode()
|
||||
with mock.patch.object(ex, "OFFICE_TEXT_CAP", 300):
|
||||
text = ex.extract_xlsx(zipped({"xl/worksheets/sheet1.xml": sheet}))
|
||||
self.assertIn("row 0", text)
|
||||
self.assertNotIn("row 199", text)
|
||||
self.assertTrue(text.splitlines()[-1].startswith("[stopped reading here:"))
|
||||
|
||||
def test_office_decompressed_cap_counts_every_member(self):
|
||||
para = b"<w:p><w:r><w:t>x</w:t></w:r></w:p>" * 2000
|
||||
doc = (b'<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">'
|
||||
b"<w:body>" + para + b"</w:body></w:document>")
|
||||
files = {"word/document.xml": doc}
|
||||
files.update({"word/header%d.xml" % i: doc for i in range(1, 6)})
|
||||
with mock.patch.object(ex, "OFFICE_DOC_CAP", len(doc) * 3):
|
||||
text = ex.extract_docx(zipped(files))
|
||||
self.assertIn("[stopped reading here: the document expands past", text)
|
||||
|
||||
def test_office_parse_is_streamed(self):
|
||||
# Rows are dropped as they are read: the parsed tree never holds them all.
|
||||
import tracemalloc
|
||||
rows = "".join('<row><c t="inlineStr"><is><t>%s</t></is></c></row>' % ("v" * 50)
|
||||
for _ in range(60000))
|
||||
sheet = ('<worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">'
|
||||
"<sheetData>%s</sheetData></worksheet>" % rows).encode()
|
||||
data = zipped({"xl/worksheets/sheet1.xml": sheet})
|
||||
tracemalloc.start()
|
||||
text = ex.extract_xlsx(data)
|
||||
peak = tracemalloc.get_traced_memory()[1]
|
||||
tracemalloc.stop()
|
||||
self.assertEqual(text.count("v" * 50), 60000)
|
||||
# the XML is ~4.6 MB; a whole-tree parse peaks at many times that
|
||||
self.assertLess(peak, 3 * len(sheet), peak)
|
||||
|
||||
def test_office_drops_unwanted_siblings_linear(self):
|
||||
body = b"<w:bookmarkStart/><w:p><w:r><w:t>hi</w:t></w:r></w:p>" * 100000
|
||||
doc = (b'<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">'
|
||||
b"<w:body>" + body + b"</w:body></w:document>")
|
||||
import time as _t
|
||||
t0 = _t.monotonic()
|
||||
text = ex.extract_docx(zipped({"word/document.xml": doc}))
|
||||
self.assertLess(_t.monotonic() - t0, 10)
|
||||
self.assertEqual(text.count("hi"), 100000)
|
||||
self.assertNotIn("[stopped", text)
|
||||
|
||||
def test_office_one_paragraph_with_endless_parts_is_cut_and_reading_goes_on(self):
|
||||
runs = b"<w:r><w:t>a</w:t></w:r>" * 10
|
||||
doc = (b'<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">'
|
||||
b"<w:body><w:p>" + runs + b"</w:p><w:p><w:r><w:t>after</w:t></w:r></w:p>"
|
||||
b"</w:body></w:document>")
|
||||
with mock.patch.object(ex, "OFFICE_PARTS_PER_ITEM", 15):
|
||||
text = ex.extract_docx(zipped({"word/document.xml": doc}))
|
||||
lines = text.splitlines()
|
||||
self.assertTrue(lines[0].startswith("a") and len(lines[0]) < 10, lines[0])
|
||||
self.assertIn("after", lines)
|
||||
self.assertEqual(lines[-1], "[1 paragraph(s) or row(s) held more than 15 parts; each was cut there]")
|
||||
|
||||
def test_office_dropping_past_the_cap_is_linear(self):
|
||||
# A closing element is not its parent's last child when the pull
|
||||
# parser has read ahead; per-element remove() was quadratic (400k
|
||||
# runs: 60 s and the rest of the document lost).
|
||||
import time as _t
|
||||
runs = b"<w:r><w:t>a</w:t></w:r>" * 60000
|
||||
doc = (b'<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">'
|
||||
b"<w:body><w:p>" + runs + b"</w:p><w:p><w:r><w:t>TAIL</w:t></w:r></w:p>"
|
||||
b"</w:body></w:document>")
|
||||
with mock.patch.object(ex, "OFFICE_PARTS_PER_ITEM", 20000):
|
||||
t0 = _t.monotonic()
|
||||
text = ex.extract_docx(zipped({"word/document.xml": doc}))
|
||||
self.assertLess(_t.monotonic() - t0, 10)
|
||||
self.assertIn("TAIL", text.splitlines())
|
||||
|
||||
def test_office_cut_row_has_no_phantom_cells(self):
|
||||
cells = "".join('<c t="inlineStr"><is><t>v</t></is></c>' for _ in range(40))
|
||||
sheet = ('<worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">'
|
||||
"<sheetData><row>%s</row><row><c t=\"inlineStr\"><is><t>next</t></is></c></row>"
|
||||
"</sheetData></worksheet>" % cells).encode()
|
||||
with mock.patch.object(ex, "OFFICE_PARTS_PER_ITEM", 30):
|
||||
text = ex.extract_xlsx(zipped({"xl/worksheets/sheet1.xml": sheet}))
|
||||
first = text.splitlines()[1]
|
||||
self.assertTrue(all(c == "v" for c in first.split("\t")), first)
|
||||
self.assertIn("next", text.splitlines())
|
||||
|
||||
def test_office_utf16_is_read_and_its_dtd_still_refused(self):
|
||||
ns = 'xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"'
|
||||
good = ('<?xml version="1.0" encoding="UTF-16"?><w:document %s><w:body>'
|
||||
"<w:p><w:r><w:t>wide caf\u00e9</w:t></w:r></w:p></w:body></w:document>" % ns)
|
||||
for enc in ("utf-16", "utf-16-be"):
|
||||
data = good.encode(enc)
|
||||
if enc == "utf-16-be":
|
||||
data = b"\xfe\xff" + data
|
||||
self.assertEqual(ex.extract_docx(zipped({"word/document.xml": data})).strip(), "wide caf\u00e9")
|
||||
bomb = ('<?xml version="1.0" encoding="UTF-16"?><!DOCTYPE d [<!ENTITY a "A">]>'
|
||||
"<w:document %s><w:body><w:p><w:r><w:t>&a;</w:t></w:r></w:p></w:body></w:document>" % ns)
|
||||
with self.assertRaises(ValueError):
|
||||
ex.extract_docx(zipped({"word/document.xml": bomb.encode("utf-16")}))
|
||||
with self.assertRaises(ValueError): # UTF-32: not an Office encoding
|
||||
ex.extract_docx(zipped({"word/document.xml": good.replace("UTF-16", "UTF-32").encode("utf-32-le")}))
|
||||
|
||||
def test_docx_text_box_follows_the_paragraph_holding_it(self):
|
||||
ns = b'xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"'
|
||||
doc = (b"<w:document " + ns + b"><w:body>"
|
||||
b"<w:p><w:r><w:t>before</w:t></w:r></w:p>"
|
||||
b"<w:p><w:r><w:t>holder</w:t></w:r><w:r><w:txbxContent>"
|
||||
b"<w:p><w:r><w:t>boxed</w:t></w:r></w:p></w:txbxContent></w:r></w:p>"
|
||||
b"<w:p><w:r><w:t>after</w:t></w:r></w:p></w:body></w:document>")
|
||||
text = ex.extract_docx(zipped({"word/document.xml": doc}))
|
||||
self.assertEqual([l for l in text.splitlines() if l], ["before", "holder", "boxed", "after"])
|
||||
|
||||
def test_eml_plain_and_html(self):
|
||||
plain = b"From: a@b.c\nTo: d@e.f\nSubject: Hi\n\nbody text\n"
|
||||
out = ex.extract_eml(plain)
|
||||
self.assertIn("Subject: Hi", out)
|
||||
self.assertIn("body text", out)
|
||||
htm = (b"From: a@b.c\nSubject: H\nContent-Type: text/html\n\n"
|
||||
b"<p>rich <b>body</b></p><script>x</script>")
|
||||
out = ex.extract_eml(htm)
|
||||
self.assertIn("rich body", out)
|
||||
self.assertNotIn("<b>", out)
|
||||
|
||||
def test_rtf(self):
|
||||
rtf = (rb"{\rtf1\ansi{\fonttbl{\f0 Arial;}}{\colortbl;\red0;}"
|
||||
rb"\f0 Hello\par caf\'e9 \u8364? end\tab tabbed}")
|
||||
out = ex.extract_rtf(rtf)
|
||||
self.assertEqual(out, "Hello\ncafé € end\ttabbed")
|
||||
self.assertNotIn("Arial", out)
|
||||
|
||||
def test_pdf_builtin_plain_and_flate(self):
|
||||
content = b"BT /F1 12 Tf 72 700 Td (Hello) Tj T* [(Wor) -20 (ld)] TJ ET"
|
||||
self.assertIn("Hello", ex.extract_pdf_builtin(pdf_bytes(content)))
|
||||
out = ex.extract_pdf_builtin(pdf_bytes(content, compress=True))
|
||||
self.assertIn("Hello", out)
|
||||
self.assertIn("World", out.replace(" ", "").replace("\n", "") and out)
|
||||
|
||||
def test_pdf_builtin_refuses_garbage_and_encrypted(self):
|
||||
garbage = b"BT (\x01\x02\x03\x04\x05\x06\x07\x08\x0b\x0e\x0f\x10\x11\x12) Tj ET"
|
||||
self.assertIsNone(ex.extract_pdf_builtin(pdf_bytes(garbage)))
|
||||
self.assertIsNone(ex.extract_pdf_builtin(pdf_bytes(b"BT (x) Tj ET") + b"/Encrypt 5 0 R"))
|
||||
self.assertIsNone(ex.extract_pdf_builtin(pdf_bytes(b"no text ops here")))
|
||||
|
||||
def test_pdf_escapes(self):
|
||||
from . import pdftext as pt
|
||||
self.assertEqual(pt._Lexer(rb"(a\)b\\c\101\n)").token(), ("str", b"a)b\\cA\n"))
|
||||
|
||||
|
||||
class StageCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.root = tempfile.mkdtemp(prefix="fa-ex-")
|
||||
self.out = tempfile.mkdtemp(prefix="fa-exout-")
|
||||
put(self.root, "a.txt", "plain text\n")
|
||||
put(self.root, "copy-of-a.txt", "plain text\n")
|
||||
put(self.root, "latin.txt", "caf\xe9".encode("cp1252"))
|
||||
put(self.root, "page.html", "<p>hello <b>web</b></p>")
|
||||
put(self.root, "data.csv", "a,b\n1,2\n")
|
||||
put(self.root, "conf.json", '{"k": 1}')
|
||||
put(self.root, "mail.eml", "From: a@b.c\nSubject: S\n\nhi\n")
|
||||
put(self.root, "doc.docx", zipped({"word/document.xml": DOCX_XML}))
|
||||
put(self.root, "sheet.xlsx", zipped({"xl/sharedStrings.xml": XLSX_SST,
|
||||
"xl/worksheets/sheet1.xml": XLSX_SHEET}))
|
||||
put(self.root, "deck.pptx", zipped({"ppt/slides/slide1.xml": PPTX_SLIDE}))
|
||||
put(self.root, "text.odt", zipped({"content.xml": ODT_CONTENT}))
|
||||
put(self.root, "doc.pdf", pdf_bytes(b"BT (pdf words) Tj ET", compress=True))
|
||||
put(self.root, "scan.png", b"\x89PNG\r\n\x1a\n" + b"\x00" * 32)
|
||||
put(self.root, "blob.bin", b"\x00\x01\x02\x03")
|
||||
put(self.root, "empty.txt", "")
|
||||
put(self.root, "noext-empty", "")
|
||||
put(self.root, "broken.docx", b"PK\x03\x04not really a zip")
|
||||
put(self.root, "big.txt", "x" * 100_000)
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.root, ignore_errors=True)
|
||||
shutil.rmtree(self.out, ignore_errors=True)
|
||||
|
||||
def run_stage(self, use_external=False, **scope):
|
||||
m = manifest(self.root, **scope)
|
||||
walk(m, self.out)
|
||||
s = ex.extract(m, self.out, use_external=use_external)
|
||||
with open(os.path.join(self.out, "extract.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(ex.store_path(self.out, rec["sha256"]), encoding="utf-8") as fh:
|
||||
return fh.read()
|
||||
|
||||
|
||||
class StageTests(StageCase):
|
||||
def test_everything_builtin(self):
|
||||
s, idx = self.run_stage(max_file_bytes=50_000)
|
||||
self.assertEqual(idx["a.txt"]["status"], "ok")
|
||||
self.assertEqual(idx["a.txt"]["extractor"], "text:utf-8")
|
||||
self.assertEqual(self.text_of(idx["a.txt"]), "plain text\n")
|
||||
self.assertEqual(idx["latin.txt"]["extractor"], "text:cp1252")
|
||||
self.assertEqual(self.text_of(idx["latin.txt"]), "café\n")
|
||||
self.assertEqual(self.text_of(idx["page.html"]).strip(), "hello web")
|
||||
self.assertEqual(self.text_of(idx["data.csv"]), "a\tb\n1\t2\n")
|
||||
self.assertEqual(idx["conf.json"]["status"], "ok")
|
||||
self.assertIn("Subject: S", self.text_of(idx["mail.eml"]))
|
||||
self.assertIn("Hello\tWorld", self.text_of(idx["doc.docx"]))
|
||||
self.assertIn("Name\t42", self.text_of(idx["sheet.xlsx"]))
|
||||
self.assertIn("Slide title", self.text_of(idx["deck.pptx"]))
|
||||
self.assertIn("Heading", self.text_of(idx["text.odt"]))
|
||||
self.assertEqual(idx["doc.pdf"]["extractor"], "pdf-builtin")
|
||||
self.assertIn("pdf words", self.text_of(idx["doc.pdf"]))
|
||||
self.assertEqual(idx["scan.png"]["status"], "unsupported")
|
||||
self.assertIn("tesseract", idx["scan.png"]["error"])
|
||||
self.assertEqual(idx["blob.bin"]["status"], "unsupported")
|
||||
self.assertEqual(idx["empty.txt"]["status"], "empty")
|
||||
# same (empty) content as empty.txt -> deduped onto its record
|
||||
self.assertEqual(idx["noext-empty"]["status"], "empty")
|
||||
self.assertEqual(idx["noext-empty"]["sha256"], idx["empty.txt"]["sha256"])
|
||||
# the inode/x-empty branch on its own
|
||||
self.assertEqual(ex.Extractor(use_external=False).extract("/nonexistent", "inode/x-empty"),
|
||||
("ok", "empty", "", None))
|
||||
self.assertEqual(idx["broken.docx"]["status"], "error")
|
||||
self.assertEqual(idx["big.txt"]["status"], "skipped")
|
||||
self.assertEqual(idx["big.txt"]["error"], "too_large")
|
||||
# duplicate content: one store entry, both paths indexed
|
||||
self.assertEqual(idx["copy-of-a.txt"]["sha256"], idx["a.txt"]["sha256"])
|
||||
c = s["counts"]
|
||||
self.assertEqual(c["files"], len(idx))
|
||||
self.assertEqual(c["skipped"], 1)
|
||||
self.assertEqual(c["error"], 1)
|
||||
self.assertEqual(c["unsupported"], 2)
|
||||
self.assertEqual(c["empty"], 2) # empty.txt + its content-twin noext-empty
|
||||
self.assertEqual(s["tools"], {"pdftotext": None, "tesseract": None})
|
||||
|
||||
def test_reuse_across_sweeps(self):
|
||||
s1, _ = self.run_stage()
|
||||
s2, idx = self.run_stage()
|
||||
# everything that produced text (or was empty) last sweep, incl. the duplicate copy
|
||||
self.assertEqual(s2["counts"]["reused"], s1["counts"]["ok"] + s1["counts"]["empty"])
|
||||
self.assertEqual(s2["counts"]["ok"], 0)
|
||||
self.assertTrue(idx["a.txt"]["reused"])
|
||||
# unsupported verdicts are reused while the tool set is unchanged
|
||||
self.assertTrue(idx["scan.png"].get("reused"))
|
||||
# a changed file is re-extracted
|
||||
put(self.root, "a.txt", "new text\n")
|
||||
s3, idx3 = self.run_stage()
|
||||
self.assertEqual(self.text_of(idx3["a.txt"]), "new text\n")
|
||||
self.assertFalse(idx3["a.txt"].get("reused"))
|
||||
|
||||
def test_cached_verdicts_invalidate_on_extractor_upgrade(self):
|
||||
self.run_stage()
|
||||
with mock.patch.object(ex, "EXTRACT_VERSION", ex.EXTRACT_VERSION + 1):
|
||||
s, idx = self.run_stage()
|
||||
self.assertEqual(s["counts"]["reused"], 0)
|
||||
self.assertFalse(idx["a.txt"].get("reused"))
|
||||
self.assertFalse(idx["scan.png"].get("reused"))
|
||||
|
||||
def test_unsupported_retried_when_tools_change(self):
|
||||
self.run_stage()
|
||||
fake_tess = lambda path, tool: "OCR TEXT\n"
|
||||
with mock.patch.object(ex, "tool_path", lambda n: "/usr/bin/" + n), \
|
||||
mock.patch.object(ex, "extract_image_tool", fake_tess), \
|
||||
mock.patch.object(ex, "extract_pdf_tool", lambda p, t: "from poppler\n"):
|
||||
s, idx = self.run_stage(use_external=True)
|
||||
self.assertEqual(idx["scan.png"]["status"], "ok")
|
||||
self.assertEqual(idx["scan.png"]["extractor"], "tesseract")
|
||||
self.assertEqual(self.text_of(idx["scan.png"]), "OCR TEXT\n")
|
||||
self.assertEqual(s["tools"]["pdftotext"], "/usr/bin/pdftotext")
|
||||
# the pdf was already extracted by the builtin path and is reused, not re-run
|
||||
self.assertTrue(idx["doc.pdf"].get("reused"))
|
||||
|
||||
def test_tool_failure_is_an_error_entry(self):
|
||||
def boom(path, tool):
|
||||
raise RuntimeError("tesseract exited 1: no language data")
|
||||
with mock.patch.object(ex, "tool_path", lambda n: "/usr/bin/" + n if n == "tesseract" else None), \
|
||||
mock.patch.object(ex, "extract_image_tool", boom):
|
||||
s, idx = self.run_stage(use_external=True)
|
||||
self.assertEqual(idx["scan.png"]["status"], "error")
|
||||
self.assertIn("no language data", idx["scan.png"]["error"])
|
||||
self.assertEqual(s["counts"]["error"], 2) # + broken.docx
|
||||
|
||||
def test_oversized_csv_field_is_an_error_entry_not_a_crash(self):
|
||||
# csv.Error is not a ValueError; a >128 KiB field used to escape the
|
||||
# handler and take the whole sweep down
|
||||
put(self.root, "huge.csv", "a,b\n1," + "x" * 200_000 + "\n")
|
||||
s, idx = self.run_stage(max_file_bytes=1_000_000)
|
||||
self.assertEqual(idx["huge.csv"]["status"], "error")
|
||||
self.assertIn("Error", idx["huge.csv"]["error"])
|
||||
self.assertEqual(s["counts"]["error"], 2) # + broken.docx
|
||||
|
||||
def test_requires_walk_first(self):
|
||||
with self.assertRaises(FileNotFoundError):
|
||||
ex.extract(manifest(self.root), self.out)
|
||||
|
||||
def test_index_sorted_and_normalised(self):
|
||||
put(self.root, "crlf.txt", "one\r\ntwo\r\n\x00")
|
||||
_, idx = self.run_stage()
|
||||
self.assertEqual(self.text_of(idx["crlf.txt"]), "one\ntwo\n")
|
||||
paths = list(idx)
|
||||
self.assertEqual(paths, sorted(paths))
|
||||
|
||||
|
||||
class RunToolTests(unittest.TestCase):
|
||||
def test_run_tool_ok_and_failure(self):
|
||||
import sys
|
||||
out = ex.run_tool([sys.executable, "-c", "print('hi')"])
|
||||
self.assertEqual(out.strip(), b"hi")
|
||||
with self.assertRaises(RuntimeError) as cm:
|
||||
ex.run_tool([sys.executable, "-c", "import sys; sys.stderr.write('bad\\n'); sys.exit(3)"])
|
||||
self.assertIn("exited 3: bad", str(cm.exception))
|
||||
with self.assertRaises(RuntimeError):
|
||||
ex.run_tool(["/nonexistent/tool"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,273 @@
|
||||
"""Unit tests for the bastion agent manifest loader. Run: python3 -m unittest discover ."""
|
||||
|
||||
import copy
|
||||
import hashlib
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
import yaml
|
||||
|
||||
from .manifest import (
|
||||
ManifestError,
|
||||
load_manifest,
|
||||
signing_digest,
|
||||
signing_payload,
|
||||
validate_document,
|
||||
)
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
EXAMPLE = os.path.join(HERE, "..", "manifests", "home-documents.yaml")
|
||||
|
||||
MINIMAL = {
|
||||
"freeagent": 1,
|
||||
"id": "docs",
|
||||
"owner": {"host": "owner-host"},
|
||||
"scope": {"roots": ["~/Documents"]},
|
||||
"mask": ["read"],
|
||||
"stages": ["walk"],
|
||||
"sweep": "manual",
|
||||
"sink": {"catalog": "local", "vectors": "local", "plaintext": "never"},
|
||||
}
|
||||
|
||||
FULL = {
|
||||
"freeagent": 1,
|
||||
"id": "docs",
|
||||
"owner": {"host": "owner-host"},
|
||||
"scope": {"roots": ["/srv/a", "~/b"], "exclude": ["**/.git/**"]},
|
||||
"mask": ["read", "search", "index", "summarize"],
|
||||
"stages": ["walk", "extract", "sanitize", "chunk", "embed", "catalog", "seal"],
|
||||
"sanitize": {"mode": "sterilize", "verify": True},
|
||||
"embed": {"model": "m", "dim": 1024, "chunk": {"size_tokens": 512, "overlap_tokens": 64}},
|
||||
"summarize": {"model": "s", "per": ["file"]},
|
||||
"sweep": "daily",
|
||||
"sink": {"catalog": "host", "vectors": "local", "plaintext": "never", "seal_to": ["k1"]},
|
||||
"grants": [{"host": "other-host", "ops": ["search", "read_range"], "until": "2026-12-31"}],
|
||||
}
|
||||
|
||||
|
||||
def errors_of(doc, **kw):
|
||||
try:
|
||||
validate_document(doc, **kw)
|
||||
except ManifestError as exc:
|
||||
return exc.errors
|
||||
return []
|
||||
|
||||
|
||||
def with_(base, **changes):
|
||||
doc = copy.deepcopy(base)
|
||||
for k, v in changes.items():
|
||||
if v is None:
|
||||
doc.pop(k, None)
|
||||
else:
|
||||
doc[k] = v
|
||||
return doc
|
||||
|
||||
|
||||
class ValidTests(unittest.TestCase):
|
||||
def test_minimal(self):
|
||||
m = validate_document(MINIMAL)
|
||||
self.assertEqual(m.stages, ("walk",))
|
||||
self.assertEqual(m.scope.max_file_bytes, 256 * 1024 * 1024)
|
||||
self.assertIsNone(m.sanitize)
|
||||
self.assertIsNone(m.embed)
|
||||
self.assertEqual(m.grants, ())
|
||||
self.assertFalse(m.signed)
|
||||
|
||||
def test_full(self):
|
||||
m = validate_document(FULL)
|
||||
self.assertEqual(m.sanitize.scope, "docs") # defaults to the manifest id
|
||||
self.assertTrue(m.sanitize.verify)
|
||||
self.assertEqual(m.embed.chunk_overlap_tokens, 64)
|
||||
self.assertEqual(m.summarize.per, ("file",))
|
||||
self.assertEqual(m.sink.seal_to, ("k1",))
|
||||
self.assertEqual(m.grants[0].host, "other-host")
|
||||
self.assertEqual(m.grants[0].until.isoformat(), "2026-12-31")
|
||||
|
||||
def test_example_file_loads(self):
|
||||
m = load_manifest(EXAMPLE)
|
||||
self.assertEqual(m.id, "home-documents")
|
||||
self.assertEqual(m.stages[-1], "seal")
|
||||
self.assertEqual(len(m.grants), 1)
|
||||
|
||||
def test_yaml_date_grant(self):
|
||||
# PyYAML parses a bare 2026-12-31 as a date; both forms are accepted
|
||||
doc = yaml.safe_load("until: 2026-12-31")
|
||||
g = with_(FULL, grants=[{"host": "h", "ops": ["search"], "until": doc["until"]}])
|
||||
self.assertEqual(errors_of(g), [])
|
||||
|
||||
|
||||
class DoctrineTests(unittest.TestCase):
|
||||
"""The rules that exist because of what the product promises."""
|
||||
|
||||
def test_plaintext_is_an_invariant(self):
|
||||
doc = with_(MINIMAL, sink={"catalog": "local", "vectors": "local", "plaintext": "host"})
|
||||
self.assertTrue(any("invariant" in e for e in errors_of(doc)))
|
||||
|
||||
def test_vectors_local_only(self):
|
||||
doc = with_(MINIMAL, sink={"catalog": "local", "vectors": "host", "plaintext": "never"})
|
||||
self.assertTrue(any("vectors" in e for e in errors_of(doc)))
|
||||
|
||||
def test_read_is_mandatory(self):
|
||||
self.assertTrue(any("mandatory" in e for e in errors_of(with_(MINIMAL, mask=["search"]))))
|
||||
|
||||
def test_write_does_not_exist(self):
|
||||
self.assertTrue(any("mask" in e for e in errors_of(with_(MINIMAL, mask=["read", "write"]))))
|
||||
|
||||
def test_stage_needs_mask_op(self):
|
||||
doc = with_(FULL, mask=["read", "search", "summarize"])
|
||||
errs = errors_of(doc)
|
||||
self.assertTrue(any("'embed' requires 'index'" in e for e in errs))
|
||||
self.assertTrue(any("'catalog' requires 'index'" in e for e in errs))
|
||||
|
||||
def test_host_catalog_needs_seal(self):
|
||||
doc = with_(FULL, stages=["walk", "extract", "sanitize", "chunk", "embed", "catalog"],
|
||||
grants=None, sink={"catalog": "host", "vectors": "local", "plaintext": "never"})
|
||||
self.assertTrue(any("unsealed" in e for e in errors_of(doc)))
|
||||
|
||||
def test_grant_needs_seal(self):
|
||||
doc = with_(FULL, stages=["walk", "extract", "sanitize", "chunk", "embed", "catalog"],
|
||||
sink={"catalog": "local", "vectors": "local", "plaintext": "never"})
|
||||
self.assertTrue(any("grants: any grant requires the 'seal'" in e for e in errors_of(doc)))
|
||||
|
||||
def test_grant_search_needs_mask_search(self):
|
||||
doc = with_(FULL, mask=["read", "index", "summarize"])
|
||||
self.assertTrue(any("op 'search' requires 'search' in mask" in e for e in errors_of(doc)))
|
||||
|
||||
def test_grant_to_self_refused(self):
|
||||
doc = with_(FULL, grants=[{"host": "owner-host", "ops": ["search"], "until": "2026-12-31"}])
|
||||
self.assertTrue(any("owner" in e for e in errors_of(doc)))
|
||||
|
||||
def test_require_signed(self):
|
||||
self.assertTrue(any("signature: required" in e
|
||||
for e in errors_of(MINIMAL, require_signed=True)))
|
||||
signed = with_(MINIMAL, signature={"alg": "ecdsa-p256-sha256", "cert": "c.crt",
|
||||
"value": "AAAA"})
|
||||
self.assertEqual(errors_of(signed, require_signed=True), [])
|
||||
|
||||
|
||||
class StageTests(unittest.TestCase):
|
||||
def test_walk_first(self):
|
||||
self.assertTrue(any("first" in e for e in errors_of(with_(MINIMAL, stages=["extract"]))))
|
||||
|
||||
def test_canonical_order(self):
|
||||
doc = with_(FULL, stages=["walk", "chunk", "extract", "sanitize", "embed", "catalog", "seal"])
|
||||
self.assertTrue(any("canonical order" in e for e in errors_of(doc)))
|
||||
|
||||
def test_prerequisite_stage(self):
|
||||
doc = with_(MINIMAL, mask=["read", "index"], stages=["walk", "catalog"])
|
||||
self.assertTrue(any("'catalog' requires 'sanitize'" in e for e in errors_of(doc)))
|
||||
# chunk and catalog read the sanitized store: extract alone is not enough
|
||||
doc = with_(MINIMAL, mask=["read", "index"], stages=["walk", "extract", "catalog"])
|
||||
self.assertTrue(any("'catalog' requires 'sanitize'" in e for e in errors_of(doc)))
|
||||
doc = with_(FULL, stages=["walk", "extract", "chunk", "embed", "catalog", "seal"], sanitize=None)
|
||||
self.assertTrue(any("'chunk' requires 'sanitize'" in e for e in errors_of(doc)))
|
||||
|
||||
def test_unknown_stage(self):
|
||||
self.assertTrue(errors_of(with_(MINIMAL, stages=["walk", "upload"])))
|
||||
|
||||
def test_section_iff_stage(self):
|
||||
errs = errors_of(with_(MINIMAL, sanitize={"mode": "flag"}))
|
||||
self.assertTrue(any("not in stages" in e for e in errs))
|
||||
errs = errors_of(with_(MINIMAL, stages=["walk", "extract", "sanitize"]))
|
||||
self.assertTrue(any("sanitize: section required" in e for e in errs))
|
||||
errs = errors_of(with_(MINIMAL, mask=["read", "index"], sanitize={"mode": "flag"},
|
||||
stages=["walk", "extract", "sanitize", "chunk", "embed"]))
|
||||
self.assertTrue(any("embed: section required" in e for e in errs))
|
||||
|
||||
def test_summarize_needs_mask_and_catalog(self):
|
||||
doc = with_(MINIMAL, summarize={"model": "s", "per": ["file"]})
|
||||
errs = errors_of(doc)
|
||||
self.assertTrue(any("'summarize' is not in mask" in e for e in errs))
|
||||
self.assertTrue(any("requires the 'catalog' stage" in e for e in errs))
|
||||
|
||||
|
||||
class FieldTests(unittest.TestCase):
|
||||
def test_unknown_keys_everywhere(self):
|
||||
self.assertTrue(any("unknown key 'extra'" in e for e in errors_of(with_(MINIMAL, extra=1))))
|
||||
doc = with_(MINIMAL, scope={"roots": ["/a"], "excludes": []})
|
||||
self.assertTrue(any("unknown key 'excludes'" in e for e in errors_of(doc)))
|
||||
|
||||
def test_relative_root_refused(self):
|
||||
doc = with_(MINIMAL, scope={"roots": ["Documents"]})
|
||||
self.assertTrue(any("absolute" in e for e in errors_of(doc)))
|
||||
self.assertTrue(errors_of(with_(MINIMAL, scope={"roots": ["C:relative"]})))
|
||||
|
||||
def test_windows_and_unc_roots_accepted(self):
|
||||
for root in (r"C:\Users\me\Documents", "D:/data", r"\\nas\share\docs", "~", "~/x", "/x"):
|
||||
self.assertEqual(errors_of(with_(MINIMAL, scope={"roots": [root]})), [], root)
|
||||
|
||||
def test_schema_version(self):
|
||||
self.assertTrue(any("schema version" in e for e in errors_of(with_(MINIMAL, freeagent=2))))
|
||||
|
||||
def test_slugs(self):
|
||||
self.assertTrue(errors_of(with_(MINIMAL, id="Home Docs")))
|
||||
self.assertTrue(errors_of(with_(MINIMAL, owner={"host": "Owner"})))
|
||||
|
||||
def test_sweep_vocabulary(self):
|
||||
self.assertTrue(any("sweep" in e for e in errors_of(with_(MINIMAL, sweep="weekly"))))
|
||||
|
||||
def test_embed_overlap_smaller_than_size(self):
|
||||
doc = with_(FULL, embed={"model": "m", "dim": 8,
|
||||
"chunk": {"size_tokens": 64, "overlap_tokens": 64}})
|
||||
self.assertTrue(any("smaller than size_tokens" in e for e in errors_of(doc)))
|
||||
|
||||
def test_flag_mode_rejects_verify(self):
|
||||
doc = with_(FULL, sanitize={"mode": "flag", "verify": True})
|
||||
self.assertTrue(any("only valid with mode 'sterilize'" in e for e in errors_of(doc)))
|
||||
|
||||
def test_bad_grant_date(self):
|
||||
doc = with_(FULL, grants=[{"host": "h", "ops": ["search"], "until": "2026-02-30"}])
|
||||
self.assertTrue(any("real YYYY-MM-DD" in e for e in errors_of(doc)))
|
||||
|
||||
def test_bad_signature_base64(self):
|
||||
doc = with_(MINIMAL, signature={"alg": "ecdsa-p256-sha256", "cert": "c", "value": "*!*"})
|
||||
self.assertTrue(any("base64" in e for e in errors_of(doc)))
|
||||
|
||||
def test_all_errors_collected(self):
|
||||
doc = with_(MINIMAL, freeagent=9, id="BAD", sweep="never", mask=["search"])
|
||||
self.assertGreaterEqual(len(errors_of(doc)), 4)
|
||||
|
||||
|
||||
class SigningTests(unittest.TestCase):
|
||||
def test_payload_excludes_signature_and_is_canonical(self):
|
||||
signed = with_(FULL, signature={"alg": "ecdsa-p256-sha256", "cert": "c", "value": "AAAA"})
|
||||
self.assertEqual(signing_payload(signed), signing_payload(FULL))
|
||||
# key order in the source must not change the bytes
|
||||
reordered = dict(reversed(list(FULL.items())))
|
||||
self.assertEqual(signing_payload(reordered), signing_payload(FULL))
|
||||
self.assertNotIn(b" ", signing_payload(FULL))
|
||||
|
||||
def test_digest_matches_payload(self):
|
||||
self.assertEqual(signing_digest(FULL),
|
||||
hashlib.sha256(signing_payload(FULL)).hexdigest())
|
||||
|
||||
def test_yaml_date_and_quoted_date_sign_identically(self):
|
||||
bare = yaml.safe_load("until: 2026-12-31")["until"]
|
||||
a = with_(FULL, grants=[{"host": "h", "ops": ["search"], "until": bare}])
|
||||
b = with_(FULL, grants=[{"host": "h", "ops": ["search"], "until": "2026-12-31"}])
|
||||
self.assertEqual(signing_payload(a), signing_payload(b))
|
||||
|
||||
def test_any_field_change_changes_digest(self):
|
||||
self.assertNotEqual(signing_digest(FULL), signing_digest(with_(FULL, sweep="manual")))
|
||||
|
||||
|
||||
class FileTests(unittest.TestCase):
|
||||
def test_missing_file(self):
|
||||
with self.assertRaises(ManifestError):
|
||||
load_manifest("/nonexistent/manifest.yaml")
|
||||
|
||||
def test_invalid_yaml(self):
|
||||
fd, path = tempfile.mkstemp(suffix=".yaml")
|
||||
with os.fdopen(fd, "w") as fh:
|
||||
fh.write("freeagent: [1\n")
|
||||
try:
|
||||
with self.assertRaises(ManifestError) as cm:
|
||||
load_manifest(path)
|
||||
self.assertIn("not valid YAML", cm.exception.errors[0])
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,541 @@
|
||||
"""The agent's built-in PDF reader (pdftext) on the shapes real writers make.
|
||||
|
||||
The regex reader it replaced returned NOTHING for a browser-printed PDF: text
|
||||
there is glyph codes in hex strings through Type0 / Identity-H fonts, mapped
|
||||
back by /ToUnicode CMaps, one Td per glyph run. These tests build such files
|
||||
with a small writer (no third-party code, no user documents) and read them.
|
||||
"""
|
||||
|
||||
import random
|
||||
import time
|
||||
import unittest
|
||||
import zlib
|
||||
from unittest import mock
|
||||
|
||||
from . import extract as ex
|
||||
from . import pdftext as pt
|
||||
|
||||
CMAP = b"""/CIDInit /ProcSet findresource begin 12 dict begin begincmap
|
||||
1 begincodespacerange <0000> <FFFF> endcodespacerange
|
||||
2 beginbfchar
|
||||
<0003> <0020>
|
||||
<0010> <00660069>
|
||||
endbfchar
|
||||
1 beginbfrange
|
||||
<0024> <003D> <0041>
|
||||
endbfrange
|
||||
endcmap CMapName currentdict /CMap defineresource pop end end
|
||||
"""
|
||||
|
||||
|
||||
def _enc(text):
|
||||
"""Glyph codes for our CMap: A-Z -> 0x24.., space -> 3, 'fi' -> 0x10."""
|
||||
out = b""
|
||||
i = 0
|
||||
while i < len(text):
|
||||
if text[i:i + 2] == "fi":
|
||||
out += b"\x00\x10"
|
||||
i += 2
|
||||
continue
|
||||
c = text[i]
|
||||
out += b"\x00\x03" if c == " " else (0x24 + ord(c) - 0x41).to_bytes(2, "big")
|
||||
i += 1
|
||||
return b"<" + out.hex().upper().encode() + b">"
|
||||
|
||||
|
||||
class Pdf:
|
||||
def __init__(self):
|
||||
self.objs = []
|
||||
|
||||
def add(self, body, stream=None, compress=True):
|
||||
n = len(self.objs) + 1
|
||||
if stream is not None:
|
||||
data = zlib.compress(stream) if compress else stream
|
||||
flt = b"/Filter /FlateDecode " if compress else b""
|
||||
body = body.rstrip(b">") + b" " + flt + b"/Length " + str(len(data)).encode() + b" >>"
|
||||
self.objs.append(b"%d 0 obj\n" % n + body + b"\nstream\n" + data + b"\nendstream\nendobj\n")
|
||||
else:
|
||||
self.objs.append(b"%d 0 obj\n" % n + body + b"\nendobj\n")
|
||||
return n
|
||||
|
||||
def bytes(self):
|
||||
return b"%PDF-1.7\n" + b"".join(self.objs) + b"trailer << /Root 1 0 R >>\n%%EOF\n"
|
||||
|
||||
|
||||
def browser_like(pages, xobject=False):
|
||||
"""A catalog -> pages tree with RESOURCES ON THE PARENT (inherited), a
|
||||
Type0 Identity-H font with a ToUnicode CMap, and per-glyph-run Td moves."""
|
||||
pdf = Pdf()
|
||||
pdf.add(b"<< /Type /Catalog /Pages 2 0 R >>") # 1
|
||||
pdf.add(b"PLACEHOLDER") # 2 pages
|
||||
cmap = pdf.add(b"<< >>", CMAP) # 3
|
||||
font = pdf.add(b"<< /Type /Font /Subtype /Type0 /BaseFont /F /Encoding /Identity-H "
|
||||
b"/ToUnicode %d 0 R >>" % cmap) # 4
|
||||
kids = []
|
||||
form = None
|
||||
if xobject:
|
||||
form = pdf.add(b"<< /Type /XObject /Subtype /Form /BBox [0 0 10 10] >>",
|
||||
b"BT /F1 9 Tf 1 0 0 1 10 10 Tm " + _enc("FOOTER") + b" Tj ET")
|
||||
for lines in pages:
|
||||
content = b"BT /F1 12 Tf 1 0 0 1 50 700 Tm "
|
||||
y = 700
|
||||
for line in lines:
|
||||
words = line.split(" ")
|
||||
# one Td per word run on the SAME line (x moves only)
|
||||
for k, w in enumerate(words):
|
||||
run = (w + " ") if k < len(words) - 1 else w
|
||||
content += b"[" + _enc(run) + b"] TJ 20 0 Td "
|
||||
y -= 14
|
||||
content += b"1 0 0 1 50 %d Tm " % y
|
||||
content += b"ET"
|
||||
if form:
|
||||
content += b" /X1 Do"
|
||||
c = pdf.add(b"<< >>", content)
|
||||
kids.append(pdf.add(b"<< /Type /Page /Parent 2 0 R /Contents %d 0 R >>" % c))
|
||||
res = b"/Resources << /Font << /F1 %d 0 R >>" % font
|
||||
if form:
|
||||
res += b" /XObject << /X1 %d 0 R >>" % form
|
||||
res += b" >>"
|
||||
pdf.objs[1] = (b"2 0 obj\n<< /Type /Pages /Count %d /Kids [%s] %s >>\nendobj\n"
|
||||
% (len(kids), b" ".join(b"%d 0 R" % k for k in kids), res))
|
||||
return pdf.bytes()
|
||||
|
||||
|
||||
class PdfTextTests(unittest.TestCase):
|
||||
def test_type0_identity_h_with_tounicode(self):
|
||||
data = browser_like([["HELLO WORLD", "SECOND LINE"]])
|
||||
self.assertEqual(pt.extract_text(data), "HELLO WORLD\nSECOND LINE")
|
||||
|
||||
def test_the_old_reader_found_nothing_here(self):
|
||||
# The shape of the regression: no literal (…) strings at all.
|
||||
data = browser_like([["HELLO WORLD"]])
|
||||
for chunk in data.split(b"stream\n")[1:]:
|
||||
raw = chunk.split(b"\nendstream")[0]
|
||||
try:
|
||||
self.assertNotIn(b"HELLO", zlib.decompress(raw))
|
||||
except zlib.error:
|
||||
pass
|
||||
|
||||
def test_words_are_not_shredded_by_per_run_moves(self):
|
||||
out = pt.extract_text(browser_like([["THE QUICK BROWN FOX"]]))
|
||||
self.assertEqual(out, "THE QUICK BROWN FOX")
|
||||
|
||||
def test_ligature_codes_map_back_to_letters(self):
|
||||
out = pt.extract_text(browser_like([["CLARIfiES"]]))
|
||||
self.assertEqual(out, "CLARIFIES".replace("FI", "fi"))
|
||||
|
||||
def test_pages_in_order_and_inherited_resources(self):
|
||||
out = pt.extract_text(browser_like([["PAGE ONE"], ["PAGE TWO"], ["PAGE THREE"]]))
|
||||
self.assertEqual(out.split("\f"), ["PAGE ONE", "PAGE TWO", "PAGE THREE"])
|
||||
|
||||
def test_form_xobject_text_is_read(self):
|
||||
out = pt.extract_text(browser_like([["BODY"]], xobject=True))
|
||||
self.assertIn("BODY", out)
|
||||
self.assertIn("FOOTER", out)
|
||||
|
||||
def test_object_streams(self):
|
||||
# The page tree and font live inside an /ObjStm (PDF 1.5+ writers).
|
||||
cmap_raw = zlib.compress(CMAP)
|
||||
content = zlib.compress(b"BT /F1 12 Tf " + _enc("PACKED") + b" Tj ET")
|
||||
inner = [
|
||||
b"<< /Type /Catalog /Pages 2 0 R >>",
|
||||
b"<< /Type /Pages /Count 1 /Kids [3 0 R] >>",
|
||||
b"<< /Type /Page /Parent 2 0 R /Contents 6 0 R "
|
||||
b"/Resources << /Font << /F1 4 0 R >> >> >>",
|
||||
b"<< /Type /Font /Subtype /Type0 /Encoding /Identity-H /ToUnicode 5 0 R >>",
|
||||
]
|
||||
header, body, off = b"", b"", 0
|
||||
for i, o in enumerate(inner, start=1):
|
||||
header += b"%d %d " % (i, off)
|
||||
body += o + b" "
|
||||
off += len(o) + 1
|
||||
packed = zlib.compress(header + body)
|
||||
doc = (b"%PDF-1.7\n"
|
||||
+ (b"5 0 obj\n<< /Filter /FlateDecode /Length %d >>\nstream\n" % len(cmap_raw))
|
||||
+ cmap_raw + b"\nendstream\nendobj\n"
|
||||
+ (b"6 0 obj\n<< /Filter /FlateDecode /Length %d >>\nstream\n" % len(content))
|
||||
+ content + b"\nendstream\nendobj\n"
|
||||
+ (b"7 0 obj\n<< /Type /ObjStm /N 4 /First %d /Filter /FlateDecode /Length %d >>\nstream\n"
|
||||
% (len(header), len(packed)))
|
||||
+ packed + b"\nendstream\nendobj\n%%EOF\n")
|
||||
self.assertEqual(pt.extract_text(doc), "PACKED")
|
||||
|
||||
def test_extract_pdf_builtin_uses_it(self):
|
||||
self.assertEqual(ex.extract_pdf_builtin(browser_like([["HELLO"]])), "HELLO")
|
||||
|
||||
def test_hostile_input_never_raises(self):
|
||||
for junk in (b"", b"%PDF-1.4\n", b"%PDF-1.4\n1 0 obj << /Type /Catalog /Pages 1 0 R >> endobj",
|
||||
b"%PDF-1.4\n1 0 obj << /Length 99999999 >>\nstream\n\x78\x9c\xff\xff\nendstream\nendobj"):
|
||||
ex.extract_pdf_builtin(junk) # returns None or text; must not raise
|
||||
|
||||
def test_reference_cycles_do_not_hang(self):
|
||||
doc = (b"%PDF-1.4\n1 0 obj << /Type /Catalog /Pages 2 0 R >> endobj\n"
|
||||
b"2 0 obj << /Type /Pages /Kids [2 0 R] >> endobj\n%%EOF")
|
||||
self.assertIsNone(pt.extract_text(doc))
|
||||
|
||||
|
||||
HELVETICA = (b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica /FirstChar 32 /LastChar 126 "
|
||||
b"/Widths [" + b" ".join([b"500"] * 95) + b"] >>")
|
||||
|
||||
|
||||
def simple_doc(streams, font=HELVETICA, objects=(), page_extra=b""):
|
||||
"""One page, font /F1, content from `streams` (each its own stream in
|
||||
/Contents), `objects` appended as-is after the page."""
|
||||
pdf = Pdf()
|
||||
pdf.add(b"<< /Type /Catalog /Pages 2 0 R >>") # 1
|
||||
pdf.add(b"<< /Type /Pages /Count 1 /Kids [4 0 R] >>") # 2
|
||||
pdf.add(font) # 3
|
||||
first = len(pdf.objs) + 2
|
||||
refs = b" ".join(b"%d 0 R" % (first + i) for i in range(len(streams)))
|
||||
pdf.add(b"<< /Type /Page /Parent 2 0 R /Resources << /Font << /F1 3 0 R >> %s >> "
|
||||
b"/Contents [%s] >>" % (page_extra, refs)) # 4
|
||||
for s in streams:
|
||||
pdf.add(b"<< >>", s)
|
||||
for o in objects:
|
||||
pdf.objs.append(o)
|
||||
return pdf.bytes()
|
||||
|
||||
|
||||
class HostileInput(unittest.TestCase):
|
||||
"""Review B1/B2: a user's file is arbitrary bytes. Every shape here once
|
||||
raised past the extractor or ran without bound."""
|
||||
|
||||
def test_kids_that_is_not_an_array(self):
|
||||
doc = (b"%PDF-1.4\n1 0 obj << /Type /Catalog /Pages 2 0 R >> endobj\n"
|
||||
b"2 0 obj << /Type /Pages /Kids 5 >> endobj\n%%EOF")
|
||||
self.assertIsNone(pt.extract_text(doc))
|
||||
|
||||
def test_objstm_offsets_that_are_not_integers(self):
|
||||
packed = zlib.compress(b"1 1.e999 2 1e5 << /Type /Catalog >> ")
|
||||
doc = (b"%PDF-1.7\n" + b"7 0 obj\n<< /Type /ObjStm /N 2 /First 14 /Filter /FlateDecode "
|
||||
b"/Length %d >>\nstream\n" % len(packed) + packed + b"\nendstream\nendobj\n%%EOF")
|
||||
self.assertIsNone(pt.extract_text(doc))
|
||||
|
||||
def test_bfrange_destination_past_its_width(self):
|
||||
cmap = (b"1 begincodespacerange <0000> <FFFF> endcodespacerange\n"
|
||||
b"1 beginbfrange <0000> <FFFF> <FFF0> endbfrange\n")
|
||||
pdf = Pdf()
|
||||
pdf.add(b"<< /Type /Catalog /Pages 2 0 R >>")
|
||||
pdf.add(b"<< /Type /Pages /Count 1 /Kids [5 0 R] >>")
|
||||
tu = pdf.add(b"<< >>", cmap)
|
||||
pdf.add(b"<< /Type /Font /Subtype /Type0 /ToUnicode %d 0 R >>" % tu)
|
||||
pdf.add(b"<< /Type /Page /Parent 2 0 R /Resources << /Font << /F1 4 0 R >> >> /Contents 6 0 R >>")
|
||||
pdf.add(b"<< >>", b"BT /F1 12 Tf <00000001> Tj ET")
|
||||
# In-range codes map (U+FFF0, U+FFF1); the run past 0xFFFF stops, never raises.
|
||||
self.assertEqual(pt.extract_text(pdf.bytes()), "\ufff0\ufff1")
|
||||
|
||||
def test_deep_nesting_is_skipped_not_a_recursion_error(self):
|
||||
doc = (b"%PDF-1.4\n9 0 obj " + b"[" * 100000 + b" endobj\n"
|
||||
+ simple_doc([b"BT /F1 12 Tf (STILL READ) Tj ET"])[len(b"%PDF-1.7\n"):])
|
||||
self.assertEqual(pt.extract_text(doc), "STILL READ")
|
||||
|
||||
def test_form_fan_out_is_bounded(self):
|
||||
# X1 draws X2 ten times, X2 draws X3 ten times … six levels: 10^6
|
||||
# invocations unbounded. The budget stops it quickly with the text
|
||||
# read so far.
|
||||
objs, first = [], 6
|
||||
for lvl in range(6):
|
||||
num = first + lvl
|
||||
if lvl < 5:
|
||||
body = b"BT /F1 12 Tf (L%d) Tj ET " % lvl + b"/X Do " * 10
|
||||
res = b"/Resources << /Font << /F1 3 0 R >> /XObject << /X %d 0 R >> >>" % (num + 1)
|
||||
else:
|
||||
body, res = b"BT /F1 12 Tf (LEAF) Tj ET", b"/Resources << /Font << /F1 3 0 R >> >>"
|
||||
objs.append(b"%d 0 obj\n<< /Type /XObject /Subtype /Form %s /Length %d >>\nstream\n"
|
||||
% (num, res, len(body)) + body + b"\nendstream\nendobj\n")
|
||||
doc = simple_doc([b"/X Do"], objects=objs, page_extra=b"/XObject << /X 6 0 R >>")
|
||||
t0 = time.monotonic()
|
||||
out = pt.extract_text(doc)
|
||||
self.assertLess(time.monotonic() - t0, 10)
|
||||
self.assertTrue(out.startswith("L0"))
|
||||
|
||||
def test_cmap_entries_are_bounded(self):
|
||||
cmap = b"1 begincodespacerange <0000> <FFFF> endcodespacerange\n" + \
|
||||
b"1 beginbfrange <0000> <FFFF> <0041> endbfrange\n" * 2000
|
||||
pdf = Pdf()
|
||||
pdf.add(b"<< /Type /Catalog /Pages 2 0 R >>")
|
||||
pdf.add(b"<< /Type /Pages /Count 1 /Kids [5 0 R] >>")
|
||||
tu = pdf.add(b"<< >>", cmap)
|
||||
pdf.add(b"<< /Type /Font /Subtype /Type0 /ToUnicode %d 0 R >>" % tu)
|
||||
pdf.add(b"<< /Type /Page /Parent 2 0 R /Resources << /Font << /F1 4 0 R >> >> /Contents 6 0 R >>")
|
||||
pdf.add(b"<< >>", b"BT /F1 12 Tf <0000> Tj ET")
|
||||
b = pt.Budget()
|
||||
t0 = time.monotonic()
|
||||
pt.extract_text(pdf.bytes(), b)
|
||||
self.assertLess(time.monotonic() - t0, 10)
|
||||
self.assertEqual(b.reason, "cmap entries")
|
||||
|
||||
def test_a_stream_is_decoded_once(self):
|
||||
doc = simple_doc([b"BT /F1 12 Tf (ONCE) Tj ET /X Do /X Do /X Do /X Do"],
|
||||
objects=[b"6 0 obj\n<< /Type /XObject /Subtype /Form /Length 30 >>\nstream\n"
|
||||
b"BT /F1 12 Tf (FORM) Tj ET \nendstream\nendobj\n"],
|
||||
page_extra=b"/XObject << /X 6 0 R >>")
|
||||
with mock.patch.object(pt._Doc, "_decode", autospec=True, side_effect=pt._Doc._decode) as dec:
|
||||
self.assertIn("FORM", pt.extract_text(doc))
|
||||
self.assertEqual(sum(1 for c in dec.call_args_list if c.args[2].startswith(b"BT /F1 12 Tf (FORM)")), 1)
|
||||
|
||||
def test_missing_endobj_is_linear(self):
|
||||
doc = b"%PDF-1.4\n" + b"1 0 obj << /A 1 >>\n" * 60000 + b"endobj\n"
|
||||
t0 = time.monotonic()
|
||||
pt.extract_text(doc)
|
||||
self.assertLess(time.monotonic() - t0, 5)
|
||||
|
||||
def test_wall_time_budget(self):
|
||||
b = pt.Budget(seconds=0)
|
||||
out = pt.extract_text(browser_like([["A"]] * 30), b)
|
||||
self.assertEqual(b.reason, "time")
|
||||
self.assertIsNone(out)
|
||||
|
||||
def test_mutated_real_shapes_never_raise(self):
|
||||
rnd = random.Random(1)
|
||||
seeds = [browser_like([["HELLO WORLD"], ["PAGE TWO"]], xobject=True),
|
||||
simple_doc([b"BT /F1 12 Tf (plain) Tj ET"])]
|
||||
for _ in range(400):
|
||||
d = bytearray(rnd.choice(seeds))
|
||||
for _ in range(rnd.randint(1, 12)):
|
||||
i = rnd.randrange(len(d))
|
||||
d[i:i + rnd.randint(1, 4)] = bytes(rnd.randrange(256) for _ in range(rnd.randint(0, 4)))
|
||||
ex.extract_pdf_builtin(bytes(d)) # None or text; never raises
|
||||
|
||||
|
||||
class Layout(unittest.TestCase):
|
||||
"""Review non-blocking items: where words and lines break."""
|
||||
|
||||
def test_bt_resets_the_text_matrix(self):
|
||||
# Two blocks on the SAME line: Td in the second block is from the
|
||||
# origin again, not from where the first block left off.
|
||||
out = pt.extract_text(simple_doc([b"BT /F1 10 Tf 50 700 Td (LEFT) Tj ET "
|
||||
b"BT /F1 10 Tf 200 700 Td (RIGHT) Tj ET"]))
|
||||
self.assertEqual(out, "LEFT RIGHT")
|
||||
|
||||
def test_words_placed_by_position_get_a_space_from_widths(self):
|
||||
# No space glyph: the gap between HELLO (5 x 500/1000 x 10 = 25 wide)
|
||||
# and a run moved 40 to the right is the only word boundary.
|
||||
out = pt.extract_text(simple_doc([b"BT /F1 10 Tf 0 0 Td (HELLO) Tj 40 0 Td (WORLD) Tj ET"]))
|
||||
self.assertEqual(out, "HELLO WORLD")
|
||||
|
||||
def test_runs_that_continue_a_word_stay_joined(self):
|
||||
out = pt.extract_text(simple_doc([b"BT /F1 10 Tf 0 0 Td (HEL) Tj 15 0 Td (LO) Tj ET"]))
|
||||
self.assertEqual(out, "HELLO")
|
||||
|
||||
def test_without_widths_no_space_is_invented(self):
|
||||
font = b"<< /Type /Font /Subtype /Type1 /BaseFont /SomeFace >>"
|
||||
out = pt.extract_text(simple_doc([b"BT /F1 10 Tf 0 0 Td (HEL) Tj 15 0 Td (LO) Tj ET"], font=font))
|
||||
self.assertEqual(out, "HELLO")
|
||||
|
||||
def test_cid_widths_from_w_array(self):
|
||||
pdf = Pdf()
|
||||
pdf.add(b"<< /Type /Catalog /Pages 2 0 R >>")
|
||||
pdf.add(b"<< /Type /Pages /Count 1 /Kids [6 0 R] >>")
|
||||
tu = pdf.add(b"<< >>", CMAP)
|
||||
desc = pdf.add(b"<< /Type /Font /Subtype /CIDFontType2 /DW 1000 /W [36 [600 600 600 600 600 600]] >>")
|
||||
pdf.add(b"<< /Type /Font /Subtype /Type0 /Encoding /Identity-H /DescendantFonts [%d 0 R] "
|
||||
b"/ToUnicode %d 0 R >>" % (desc, tu))
|
||||
pdf.add(b"<< /Type /Page /Parent 2 0 R /Resources << /Font << /F1 5 0 R >> >> /Contents 7 0 R >>")
|
||||
# "AB" at 600/1000 x 10 = 12 wide; the next run 20 on is a new word.
|
||||
pdf.add(b"<< >>", b"BT /F1 10 Tf 0 0 Td " + _enc("AB") + b" Tj 20 0 Td " + _enc("CD") + b" Tj ET")
|
||||
self.assertEqual(pt.extract_text(pdf.bytes()), "AB CD")
|
||||
|
||||
def test_winansi_encoding(self):
|
||||
font = b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica /Encoding /WinAnsiEncoding >>"
|
||||
out = pt.extract_text(simple_doc([b"BT /F1 10 Tf (caf\\351 \\223q\\224 \\200) Tj ET"], font=font))
|
||||
self.assertEqual(out, "caf\u00e9 \u201cq\u201d \u20ac")
|
||||
|
||||
def test_macroman_encoding(self):
|
||||
font = (b"<< /Type /Font /Subtype /Type1 /BaseFont /Times "
|
||||
b"/Encoding << /BaseEncoding /MacRomanEncoding >> >>")
|
||||
out = pt.extract_text(simple_doc([b"BT /F1 10 Tf (caf\\216) Tj ET"], font=font))
|
||||
self.assertEqual(out, "caf\u00e9")
|
||||
|
||||
def test_font_carries_across_a_pages_content_streams(self):
|
||||
pdf = Pdf()
|
||||
pdf.add(b"<< /Type /Catalog /Pages 2 0 R >>")
|
||||
pdf.add(b"<< /Type /Pages /Count 1 /Kids [5 0 R] >>")
|
||||
tu = pdf.add(b"<< >>", CMAP)
|
||||
pdf.add(b"<< /Type /Font /Subtype /Type0 /Encoding /Identity-H /ToUnicode %d 0 R >>" % tu)
|
||||
pdf.add(b"<< /Type /Page /Parent 2 0 R /Resources << /Font << /F1 4 0 R >> >> "
|
||||
b"/Contents [6 0 R 7 0 R] >>")
|
||||
pdf.add(b"<< >>", b"BT /F1 12 Tf")
|
||||
pdf.add(b"<< >>", _enc("SPLIT") + b" Tj ET")
|
||||
self.assertEqual(pt.extract_text(pdf.bytes()), "SPLIT")
|
||||
|
||||
def test_incremental_update_later_definition_wins(self):
|
||||
doc = simple_doc([b"BT /F1 12 Tf (OLD) Tj ET"])
|
||||
body = b"BT /F1 12 Tf (NEW) Tj ET"
|
||||
doc += (b"5 0 obj\n<< /Length %d >>\nstream\n" % len(body)) + body + b"\nendstream\nendobj\n%%EOF\n"
|
||||
self.assertEqual(pt.extract_text(doc), "NEW")
|
||||
|
||||
def test_stream_holding_endobj_text_is_read_whole(self):
|
||||
pdf = Pdf()
|
||||
pdf.add(b"<< /Type /Catalog /Pages 2 0 R >>")
|
||||
pdf.add(b"<< /Type /Pages /Count 1 /Kids [4 0 R] >>")
|
||||
pdf.add(HELVETICA)
|
||||
pdf.add(b"<< /Type /Page /Parent 2 0 R /Resources << /Font << /F1 3 0 R >> >> /Contents 5 0 R >>")
|
||||
pdf.add(b"<< >>", b"BT /F1 12 Tf (a endobj b) Tj 0 -20 Td (AFTER) Tj ET", compress=False)
|
||||
self.assertEqual(pt.extract_text(pdf.bytes()), "a endobj b\nAFTER")
|
||||
|
||||
|
||||
class ReviewRound2(unittest.TestCase):
|
||||
"""Second review: lexing is what repeats, so lexing is what is charged;
|
||||
one bad value costs one stream, not the document."""
|
||||
|
||||
def _pages_sharing(self, body, npages):
|
||||
pdf = Pdf()
|
||||
pdf.add(b"<< /Type /Catalog /Pages 2 0 R >>")
|
||||
kids = b" ".join(b"%d 0 R" % (5 + i) for i in range(npages))
|
||||
pdf.add(b"<< /Type /Pages /Count %d /Kids [%s] >>" % (npages, kids))
|
||||
pdf.add(HELVETICA)
|
||||
pdf.add(b"<< >>", body) # 4
|
||||
for _ in range(npages):
|
||||
pdf.add(b"<< /Type /Page /Parent 2 0 R /Resources << /Font << /F1 3 0 R >> >> "
|
||||
b"/Contents 4 0 R >>")
|
||||
return pdf.bytes()
|
||||
|
||||
def test_one_stream_lexed_by_many_pages_is_bounded(self):
|
||||
doc = self._pages_sharing(b"BT /F1 12 Tf (hello) Tj ET " + b"1 " * 1_000_000, 40)
|
||||
b = pt.Budget(seconds=1)
|
||||
t0 = time.monotonic()
|
||||
out = pt.extract_text(doc, b)
|
||||
self.assertLess(time.monotonic() - t0, 5)
|
||||
self.assertEqual(b.reason, "time")
|
||||
self.assertTrue(out.startswith("hello"))
|
||||
|
||||
def test_token_cap(self):
|
||||
doc = self._pages_sharing(b"BT /F1 12 Tf (hi) Tj ET " + b"1 " * 50_000, 10)
|
||||
with mock.patch.object(pt, "MAX_TOKENS", 120_000):
|
||||
b = pt.Budget()
|
||||
pt.extract_text(doc, b)
|
||||
self.assertEqual(b.reason, "tokens")
|
||||
|
||||
def test_a_shared_tounicode_is_parsed_once(self):
|
||||
pdf = Pdf()
|
||||
pdf.add(b"<< /Type /Catalog /Pages 2 0 R >>")
|
||||
pdf.add(b"<< /Type /Pages /Count 1 /Kids [6 0 R] >>")
|
||||
tu = pdf.add(b"<< >>", CMAP) # 3
|
||||
pdf.add(b"<< /Type /Font /Subtype /Type0 /ToUnicode %d 0 R >>" % tu) # 4
|
||||
pdf.add(b"<< /Type /Font /Subtype /Type0 /ToUnicode %d 0 R >>" % tu) # 5
|
||||
pdf.add(b"<< /Type /Page /Parent 2 0 R /Resources << /Font << /F1 4 0 R /F2 5 0 R >> >> "
|
||||
b"/Contents 7 0 R >>")
|
||||
pdf.add(b"<< >>", b"BT /F1 12 Tf " + _enc("AB") + b" Tj /F2 12 Tf " + _enc("CD") + b" Tj ET")
|
||||
with mock.patch.object(pt._Font, "_parse", autospec=True, side_effect=pt._Font._parse) as parse:
|
||||
self.assertEqual(pt.extract_text(pdf.bytes()), "ABCD")
|
||||
self.assertEqual(parse.call_count, 1)
|
||||
|
||||
def test_words_far_apart_without_widths_are_separated(self):
|
||||
font = b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>"
|
||||
for ops in (b"BT /F1 10 Tf 50 700 Td (Name) Tj 150 0 Td (Value) Tj ET",
|
||||
b"BT /F1 10 Tf 1 0 0 1 50 700 Tm (Name) Tj 1 0 0 1 200 700 Tm (Value) Tj ET"):
|
||||
self.assertEqual(pt.extract_text(simple_doc([ops], font=font)), "Name Value")
|
||||
|
||||
def test_standard_font_without_widths_uses_its_real_widths(self):
|
||||
# Helvetica W is 944/1000: (WWWW) ends at 37.76 at size 10. The rest of
|
||||
# the word placed exactly there is the same word; a guessed 500 would
|
||||
# have read a 17.76-unit gap and split it.
|
||||
font = b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>"
|
||||
out = pt.extract_text(simple_doc([b"BT /F1 10 Tf 0 0 Td (WWWW) Tj 37.76 0 Td (ORD) Tj ET"],
|
||||
font=font))
|
||||
self.assertEqual(out, "WWWWORD")
|
||||
self.assertEqual(pt._standard_widths(b"ABCDEF+Arial,Bold")[1][ord("W") - 32], 944)
|
||||
self.assertEqual(pt._standard_widths(b"Courier-Oblique"), (32, [600] * 95))
|
||||
self.assertIsNone(pt._standard_widths(b"SomeFace"))
|
||||
|
||||
def test_unknown_font_without_widths_guesses_wide(self):
|
||||
font = b"<< /Type /Font /Subtype /Type1 /BaseFont /SomeFace >>"
|
||||
out = pt.extract_text(simple_doc([b"BT /F1 10 Tf 0 0 Td (WWWW) Tj 37.76 0 Td (ORD) Tj ET"],
|
||||
font=font))
|
||||
self.assertEqual(out, "WWWWORD")
|
||||
|
||||
def test_superscript_rise_is_not_a_new_line(self):
|
||||
out = pt.extract_text(simple_doc([b"BT /F1 10 Tf 0 0 Td (E=mc) Tj 20 3 Td (2) Tj "
|
||||
b"5 -3 Td ( is) Tj ET"]))
|
||||
self.assertEqual(out, "E=mc2 is")
|
||||
|
||||
def test_the_word_stream_inside_a_string_owns_no_bytes(self):
|
||||
# The stray object sits BEFORE every page object.
|
||||
doc = simple_doc([b"BT /F1 12 Tf (BODY) Tj ET"])
|
||||
head, rest = doc.split(b"1 0 obj", 1)
|
||||
doc = head + b"9 0 obj\n<< /Title (a live stream\n) >>\nendobj\n1 0 obj" + rest
|
||||
self.assertEqual(pt.extract_text(doc), "BODY")
|
||||
|
||||
def test_indirect_length(self):
|
||||
body = b"BT /F1 12 Tf (INDIRECT) Tj ET"
|
||||
raw = (b"%PDF-1.4\n1 0 obj << /Type /Catalog /Pages 2 0 R >> endobj\n"
|
||||
b"2 0 obj << /Type /Pages /Count 1 /Kids [3 0 R] >> endobj\n"
|
||||
b"3 0 obj << /Type /Page /Parent 2 0 R /Contents 4 0 R >> endobj\n"
|
||||
b"4 0 obj << /Length 5 0 R >>\nstream\n" + body + b"\nendstream\nendobj\n"
|
||||
b"5 0 obj %d endobj\n%%%%EOF" % len(body))
|
||||
self.assertEqual(pt.extract_text(raw), "INDIRECT")
|
||||
|
||||
def test_one_bad_value_costs_one_stream_not_the_document(self):
|
||||
bad_font = b"<< /Type /Font /Subtype /Type1 /Encoding [/a /b] >>"
|
||||
self.assertEqual(pt.extract_text(simple_doc([b"BT /F1 12 Tf (still) Tj ET"], font=bad_font)),
|
||||
"still")
|
||||
huge = b"BT /F1 12 Tf (hello world) Tj [(a) 1" + b"0" * 400 + b" (b)] TJ ET"
|
||||
self.assertTrue(pt.extract_text(simple_doc([huge])).startswith("hello world"))
|
||||
pdf = Pdf()
|
||||
pdf.add(b"<< /Type /Catalog /Pages 2 0 R >>")
|
||||
pdf.add(b"<< /Type /Pages /Count 2 /Kids [4 0 R 5 0 R] >>")
|
||||
pdf.add(HELVETICA)
|
||||
pdf.add(b"<< /Type /Page /Parent 2 0 R /Resources << /Font << /F1 3 0 R >> >> /Contents 6 0 R >>")
|
||||
pdf.add(b"<< /Type /Page /Parent 2 0 R /Resources << /Font << /F1 3 0 R >> >> /Contents 7 0 R >>")
|
||||
pdf.add(b"<< >>", b"BT /F1 12 Tf (PAGE ONE) Tj ET")
|
||||
pdf.add(b"<< >>", b"BT /F1 12 Tf (PAGE TWO) Tj ET")
|
||||
real = pt._Reader.content
|
||||
calls = []
|
||||
|
||||
def first_page_breaks(self, data, *a, **k):
|
||||
calls.append(1)
|
||||
if len(calls) == 1:
|
||||
raise TypeError("a malformed value")
|
||||
return real(self, data, *a, **k)
|
||||
with mock.patch.object(pt._Reader, "content", first_page_breaks):
|
||||
self.assertEqual(pt.extract_text(pdf.bytes()).split("\f"), ["", "PAGE TWO"])
|
||||
|
||||
def test_cid_pages_without_a_map_do_not_fall_back_to_junk(self):
|
||||
pdf = Pdf()
|
||||
pdf.add(b"<< /Type /Catalog /Pages 2 0 R >>")
|
||||
pdf.add(b"<< /Type /Pages /Count 1 /Kids [4 0 R] >>")
|
||||
pdf.add(b"<< /Type /Font /Subtype /Type0 /Encoding /Identity-H >>")
|
||||
pdf.add(b"<< /Type /Page /Parent 2 0 R /Resources << /Font << /F1 3 0 R >> >> /Contents 5 0 R >>")
|
||||
pdf.add(b"<< >>", b"BT /F1 12 Tf <00410042> Tj ET")
|
||||
self.assertIsNone(pt.extract_text(pdf.bytes()))
|
||||
|
||||
def test_index_checks_the_clock(self):
|
||||
data = b"%PDF-1.4\n" + b"".join(b"%d 0 obj [1 2 3] endobj\n" % i for i in range(200_000))
|
||||
b = pt.Budget(seconds=0)
|
||||
pt.extract_text(data, b)
|
||||
self.assertEqual(b.reason, "time")
|
||||
|
||||
|
||||
|
||||
class ReviewRound4(unittest.TestCase):
|
||||
"""Memory, not just time: a lexer step must not allocate per byte."""
|
||||
|
||||
def _peak(self, fn):
|
||||
import tracemalloc
|
||||
tracemalloc.start()
|
||||
fn()
|
||||
peak = tracemalloc.get_traced_memory()[1]
|
||||
tracemalloc.stop()
|
||||
return peak
|
||||
|
||||
def test_spaced_hex_string_is_stripped_without_per_byte_objects(self):
|
||||
data = b"<" + b"0 " * 2_000_000 + b"41>"
|
||||
peak = self._peak(lambda: pt._Lexer(data).token())
|
||||
self.assertEqual(pt._Lexer(b"<4 1 4 2>").token(), ("str", b"AB"))
|
||||
self.assertLess(peak, 3 * len(data), peak)
|
||||
|
||||
def test_comment_lines_are_skipped_without_per_line_state(self):
|
||||
data = b"%\n" * 1_000_000 + b"(x)"
|
||||
peak = self._peak(lambda: pt._Lexer(data).token())
|
||||
self.assertEqual(pt._Lexer(data).token(), ("str", b"x"))
|
||||
self.assertLess(peak, len(data), peak)
|
||||
|
||||
def test_a_long_escape_run_checks_the_clock(self):
|
||||
b = pt.Budget(seconds=0)
|
||||
with self.assertRaises(pt._Spent):
|
||||
pt._Lexer(b"(" + b"\\a" * 300_000 + b")", budget=b).token()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,288 @@
|
||||
"""Unit tests for the sweep scheduler and the manifest gate. Run:
|
||||
python3 -m unittest discover .
|
||||
|
||||
The gate tests generate a real EC key + self-signed cert with openssl (on
|
||||
PATH here and in CI); the loop is driven with a fake clock and sleep so
|
||||
a 'daily' cadence runs in milliseconds.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
import yaml
|
||||
|
||||
from . import run as rn
|
||||
from .manifest import load_manifest
|
||||
|
||||
HAVE_OPENSSL = bool(shutil.which("openssl"))
|
||||
|
||||
|
||||
def write_manifest(path, root, stages, cn="owner-host", sweep="manual", extra=None):
|
||||
doc = {
|
||||
"freeagent": 1, "id": "t", "owner": {"host": cn},
|
||||
"scope": {"roots": [root]}, "mask": ["read", "index"],
|
||||
"stages": stages, "sweep": sweep,
|
||||
"sink": {"catalog": "local", "vectors": "local", "plaintext": "never"},
|
||||
}
|
||||
if "sanitize" in stages:
|
||||
doc["sanitize"] = {"mode": "flag"}
|
||||
if extra:
|
||||
doc.update(extra)
|
||||
with open(path, "w", encoding="utf-8") as fh:
|
||||
yaml.safe_dump(doc, fh, sort_keys=False)
|
||||
return path
|
||||
|
||||
|
||||
def make_identity(dirpath, cn):
|
||||
key = os.path.join(dirpath, "host.key")
|
||||
cert = os.path.join(dirpath, "host.crt")
|
||||
subprocess.run(["openssl", "ecparam", "-name", "prime256v1", "-genkey", "-noout", "-out", key],
|
||||
check=True, capture_output=True)
|
||||
subprocess.run(["openssl", "req", "-new", "-x509", "-key", key, "-out", cert, "-days", "2",
|
||||
"-subj", "/CN=%s" % cn], check=True, capture_output=True)
|
||||
return key, cert
|
||||
|
||||
|
||||
class Case(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.root = tempfile.mkdtemp(prefix="fa-run-")
|
||||
self.work = tempfile.mkdtemp(prefix="fa-runwork-")
|
||||
self.out = os.path.join(self.work, "out")
|
||||
with open(os.path.join(self.root, "a.txt"), "w", encoding="utf-8") as fh:
|
||||
fh.write("hello m@example.com\n")
|
||||
self.manifest = os.path.join(self.work, "m.yaml")
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.root, ignore_errors=True)
|
||||
shutil.rmtree(self.work, ignore_errors=True)
|
||||
|
||||
def runs(self):
|
||||
with open(os.path.join(self.out, "run.jsonl"), encoding="utf-8") as fh:
|
||||
return [json.loads(l) for l in fh]
|
||||
|
||||
|
||||
@unittest.skipUnless(HAVE_OPENSSL, "openssl not on PATH")
|
||||
class GateTests(Case):
|
||||
def test_sign_then_verify(self):
|
||||
write_manifest(self.manifest, self.root, ["walk"])
|
||||
key, cert = make_identity(self.work, "owner-host")
|
||||
rn.sign_manifest_file(self.manifest, key, "host.crt")
|
||||
m = load_manifest(self.manifest, require_signed=True)
|
||||
gate = rn.verify_signature(m, {"FREEAGENT_CA_ROOT": cert}, manifest_dir=self.work)
|
||||
self.assertEqual(gate["cn"], "owner-host")
|
||||
self.assertTrue(gate["chained"])
|
||||
# no root, no gate: "signed by the key next to the manifest" proves nothing
|
||||
with self.assertRaises(rn.GateError) as cm:
|
||||
rn.verify_signature(m, {}, manifest_dir=self.work)
|
||||
self.assertIn("FREEAGENT_CA_ROOT", str(cm.exception))
|
||||
# re-signing replaces the block rather than stacking a second one
|
||||
rn.sign_manifest_file(self.manifest, key, "host.crt")
|
||||
with open(self.manifest, encoding="utf-8") as fh:
|
||||
self.assertEqual(fh.read().count("signature:"), 1)
|
||||
|
||||
def test_tampered_payload_fails(self):
|
||||
write_manifest(self.manifest, self.root, ["walk"])
|
||||
key, cert = make_identity(self.work, "owner-host")
|
||||
rn.sign_manifest_file(self.manifest, key, "host.crt")
|
||||
with open(self.manifest, encoding="utf-8") as fh:
|
||||
text = fh.read()
|
||||
with open(self.manifest, "w", encoding="utf-8") as fh:
|
||||
fh.write(text.replace("sweep: manual", "sweep: daily"))
|
||||
m = load_manifest(self.manifest, require_signed=True)
|
||||
with self.assertRaises(rn.GateError) as cm:
|
||||
rn.verify_signature(m, {"FREEAGENT_CA_ROOT": cert}, manifest_dir=self.work)
|
||||
self.assertIn("does not verify", str(cm.exception))
|
||||
|
||||
def test_wrong_cn_fails(self):
|
||||
write_manifest(self.manifest, self.root, ["walk"], cn="owner-host")
|
||||
key, cert = make_identity(self.work, "someone-else")
|
||||
rn.sign_manifest_file(self.manifest, key, "host.crt")
|
||||
m = load_manifest(self.manifest, require_signed=True)
|
||||
with self.assertRaises(rn.GateError) as cm:
|
||||
rn.verify_signature(m, {"FREEAGENT_CA_ROOT": cert}, manifest_dir=self.work)
|
||||
self.assertIn("not the manifest owner", str(cm.exception))
|
||||
|
||||
def test_multi_cn_cert_refused(self):
|
||||
# a cert carrying the owner's CN next to its real one must not pass
|
||||
write_manifest(self.manifest, self.root, ["walk"], cn="owner-host")
|
||||
key = os.path.join(self.work, "host.key")
|
||||
cert = os.path.join(self.work, "host.crt")
|
||||
subprocess.run(["openssl", "ecparam", "-name", "prime256v1", "-genkey", "-noout", "-out", key],
|
||||
check=True, capture_output=True)
|
||||
subprocess.run(["openssl", "req", "-new", "-x509", "-key", key, "-out", cert, "-days", "2",
|
||||
"-subj", "/CN=attacker/CN=owner-host"], check=True, capture_output=True)
|
||||
rn.sign_manifest_file(self.manifest, key, "host.crt")
|
||||
m = load_manifest(self.manifest, require_signed=True)
|
||||
with self.assertRaises(rn.GateError) as cm:
|
||||
rn.verify_signature(m, {"FREEAGENT_CA_ROOT": cert}, manifest_dir=self.work)
|
||||
self.assertIn("exactly one CN", str(cm.exception))
|
||||
|
||||
def test_ca_chain(self):
|
||||
write_manifest(self.manifest, self.root, ["walk"])
|
||||
key, cert = make_identity(self.work, "owner-host")
|
||||
rn.sign_manifest_file(self.manifest, key, "host.crt")
|
||||
m = load_manifest(self.manifest, require_signed=True)
|
||||
# self-signed: chains to itself, not to a stranger's root
|
||||
gate = rn.verify_signature(m, {"FREEAGENT_CA_ROOT": cert}, manifest_dir=self.work)
|
||||
self.assertTrue(gate["chained"])
|
||||
_, other_root = make_identity(tempfile.mkdtemp(prefix="fa-ca-"), "other-ca")
|
||||
with self.assertRaises(rn.GateError) as cm:
|
||||
rn.verify_signature(m, {"FREEAGENT_CA_ROOT": other_root}, manifest_dir=self.work)
|
||||
self.assertIn("chain", str(cm.exception))
|
||||
|
||||
def test_unsigned_refused_unless_allowed(self):
|
||||
write_manifest(self.manifest, self.root, ["walk"])
|
||||
rec, _ = rn.run_once(self.manifest, self.out, env={})
|
||||
self.assertFalse(rec["ok"])
|
||||
self.assertIn("unsigned", rec["error"])
|
||||
self.assertEqual(rec["stages"], [])
|
||||
# the refusal is on the append-only record
|
||||
self.assertEqual(len(self.runs()), 1)
|
||||
self.assertIn("gate:", self.runs()[0]["error"])
|
||||
rec, _ = rn.run_once(self.manifest, self.out, env={}, allow_unsigned=True)
|
||||
self.assertTrue(rec["ok"])
|
||||
self.assertTrue(rec["allow_unsigned"])
|
||||
self.assertIn("bypassed", rec["gate"])
|
||||
|
||||
|
||||
class RunTests(Case):
|
||||
def test_stages_in_order_and_logged(self):
|
||||
write_manifest(self.manifest, self.root, ["walk", "extract", "sanitize", "catalog"])
|
||||
rec, m = rn.run_once(self.manifest, self.out, env={}, allow_unsigned=True, now="2026-08-28T00:00:00Z")
|
||||
self.assertTrue(rec["ok"], rec)
|
||||
self.assertEqual([s["stage"] for s in rec["stages"]], ["walk", "extract", "sanitize", "catalog"])
|
||||
self.assertTrue(all(s["status"] == "ok" for s in rec["stages"]))
|
||||
self.assertEqual(rec["stages"][0]["files"], 1)
|
||||
self.assertEqual(rec["stages"][2]["ok"], 1)
|
||||
self.assertTrue(os.path.exists(os.path.join(self.out, "catalog.jsonl")))
|
||||
self.assertEqual(len(self.runs()), 1)
|
||||
self.assertEqual(self.runs()[0]["started"], "2026-08-28T00:00:00Z")
|
||||
|
||||
def test_default_out_is_the_folder_serve_uses(self):
|
||||
# without --out, sweep must write where serve reads: <root>/.agent
|
||||
write_manifest(self.manifest, self.root, ["walk"])
|
||||
rec, m = rn.run_once(self.manifest, None, env={}, allow_unsigned=True)
|
||||
self.assertTrue(rec["ok"], rec)
|
||||
want = os.path.join(os.path.abspath(self.root), ".agent")
|
||||
self.assertEqual(rec["out"], want)
|
||||
from .__main__ import default_out # what serve and every stage command use
|
||||
self.assertEqual(default_out(m), want)
|
||||
self.assertTrue(os.path.exists(os.path.join(want, "run.jsonl")))
|
||||
|
||||
def test_failure_stops_the_run(self):
|
||||
# sterilize with an unstartable recall backend stops at sanitize
|
||||
write_manifest(self.manifest, self.root, ["walk", "extract", "sanitize", "catalog"],
|
||||
extra={"sanitize": {"mode": "sterilize"}})
|
||||
rec, _ = rn.run_once(self.manifest, self.out, env={"PII_RECALL": "nope"}, allow_unsigned=True)
|
||||
self.assertFalse(rec["ok"])
|
||||
self.assertEqual([s["stage"] for s in rec["stages"]], ["walk", "extract", "sanitize"])
|
||||
self.assertEqual(rec["stages"][-1]["status"], "failed")
|
||||
self.assertIn("PII_RECALL", rec["error"])
|
||||
self.assertFalse(os.path.exists(os.path.join(self.out, "catalog.jsonl")))
|
||||
|
||||
def test_invalid_manifest_is_a_recorded_error(self):
|
||||
with open(self.manifest, "w") as fh:
|
||||
fh.write("freeagent: 2\n")
|
||||
rec, m = rn.run_once(self.manifest, self.out, env={}, allow_unsigned=True)
|
||||
self.assertIsNone(m)
|
||||
self.assertIn("manifest:", rec["error"])
|
||||
self.assertIn("manifest:", self.runs()[0]["error"]) # refused manifests are logged too
|
||||
|
||||
def test_lock_blocks_and_reclaims(self):
|
||||
write_manifest(self.manifest, self.root, ["walk"])
|
||||
os.makedirs(self.out, exist_ok=True)
|
||||
with rn.Lock(self.out):
|
||||
with self.assertRaises(rn.RunError):
|
||||
with rn.Lock(self.out):
|
||||
pass
|
||||
rec, _ = rn.run_once(self.manifest, self.out, env={}, allow_unsigned=True)
|
||||
self.assertIn("another sweep holds", rec["error"])
|
||||
# a stale lock (dead pid) is reclaimed
|
||||
with open(os.path.join(self.out, "sweep.lock"), "w") as fh:
|
||||
json.dump({"pid": 2 ** 22 + 12345, "since": 0}, fh)
|
||||
rec, _ = rn.run_once(self.manifest, self.out, env={}, allow_unsigned=True)
|
||||
self.assertTrue(rec["ok"])
|
||||
self.assertIn("lock_reclaimed", rec)
|
||||
self.assertFalse(os.path.exists(os.path.join(self.out, "sweep.lock")))
|
||||
# a LIVE holder keeps its lock however old the file is (never stolen)
|
||||
with open(os.path.join(self.out, "sweep.lock"), "w") as fh:
|
||||
json.dump({"pid": os.getpid(), "since": 0}, fh)
|
||||
rec, _ = rn.run_once(self.manifest, self.out, env={}, allow_unsigned=True)
|
||||
self.assertIn("another sweep holds", rec["error"])
|
||||
os.unlink(os.path.join(self.out, "sweep.lock"))
|
||||
# a truncated lock is judged by file age, not instantly stale
|
||||
with open(os.path.join(self.out, "sweep.lock"), "w") as fh:
|
||||
fh.write("{")
|
||||
rec, _ = rn.run_once(self.manifest, self.out, env={}, allow_unsigned=True)
|
||||
self.assertIn("another sweep holds", rec["error"])
|
||||
os.unlink(os.path.join(self.out, "sweep.lock"))
|
||||
|
||||
def test_next_due(self):
|
||||
self.assertEqual(rn.next_due("hourly", 1000, True, 0), 4600)
|
||||
self.assertEqual(rn.next_due("15m", 1000, True, 0), 1900)
|
||||
self.assertIsNone(rn.next_due("manual", 1000, True, 0))
|
||||
self.assertEqual(rn.next_due("daily", 1000, False, 1), 1060) # 1 min
|
||||
self.assertEqual(rn.next_due("daily", 1000, False, 3), 1240) # 4 min
|
||||
self.assertEqual(rn.next_due("15m", 1000, False, 20), 1900) # capped at cadence
|
||||
|
||||
|
||||
class SweepLoopTests(Case):
|
||||
def test_manual_runs_once(self):
|
||||
write_manifest(self.manifest, self.root, ["walk"], sweep="manual")
|
||||
sleeps = []
|
||||
last = rn.sweep(self.manifest, self.out, env={}, allow_unsigned=True,
|
||||
sleep=sleeps.append, clock=lambda: 1000.0, log=lambda s: None)
|
||||
self.assertTrue(last["ok"])
|
||||
self.assertEqual(sleeps, [])
|
||||
state = rn.load_state(self.out)
|
||||
self.assertIsNone(state["next_due"])
|
||||
self.assertEqual(state["consecutive_failures"], 0)
|
||||
|
||||
def test_cadence_loops_with_fake_clock(self):
|
||||
write_manifest(self.manifest, self.root, ["walk"], sweep="15m")
|
||||
clock = {"t": 1000.0}
|
||||
sleeps = []
|
||||
|
||||
def sleep(s):
|
||||
sleeps.append(s)
|
||||
clock["t"] += s
|
||||
|
||||
logs = []
|
||||
last = rn.sweep(self.manifest, self.out, env={}, allow_unsigned=True, sleep=sleep,
|
||||
clock=lambda: clock["t"], max_runs=3, log=logs.append)
|
||||
self.assertTrue(last["ok"])
|
||||
self.assertEqual(len(self.runs()), 3)
|
||||
self.assertEqual(sum(sleeps), 2 * 15 * 60) # slept out two full cadences
|
||||
self.assertTrue(all(s <= 30 for s in sleeps)) # in short naps, so a stop lands quickly
|
||||
state = rn.load_state(self.out)
|
||||
self.assertEqual(state["runs_this_process"], 3)
|
||||
self.assertEqual(state["next_due"], "1970-01-01T01:01:40Z") # 1000 + 2*900 + 900
|
||||
self.assertEqual(len(logs), 3)
|
||||
self.assertIn("ok", logs[0])
|
||||
|
||||
def test_failure_backoff_and_counter(self):
|
||||
write_manifest(self.manifest, self.root, ["walk", "extract", "sanitize"], sweep="hourly",
|
||||
extra={"sanitize": {"mode": "sterilize"}})
|
||||
clock = {"t": 1000.0}
|
||||
sleeps = []
|
||||
|
||||
def sleep(s):
|
||||
sleeps.append(s)
|
||||
clock["t"] += s
|
||||
|
||||
rn.sweep(self.manifest, self.out, env={"PII_RECALL": "nope"}, allow_unsigned=True,
|
||||
sleep=sleep, clock=lambda: clock["t"], max_runs=3, log=lambda s: None)
|
||||
state = rn.load_state(self.out)
|
||||
self.assertEqual(state["consecutive_failures"], 3)
|
||||
self.assertFalse(state["last_ok"])
|
||||
# retries at 60 s then 120 s, not the hourly cadence
|
||||
self.assertEqual(sum(sleeps), 60 + 120)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,383 @@
|
||||
"""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)
|
||||
self.assertNotIn("4111", 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)
|
||||
self.assertNotIn("4111", 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_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()
|
||||
@@ -0,0 +1,197 @@
|
||||
"""Unit tests for the seal stage. Run: python3 -m unittest discover .
|
||||
|
||||
Round trips use a real age identity generated with age-keygen when the
|
||||
binaries are on PATH (they are on the server and in CI); the envelope
|
||||
tests are skipped where age is absent, the fail-closed tests are not.
|
||||
"""
|
||||
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import tarfile
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
from . import seal as sl
|
||||
from .catalog import catalog
|
||||
from .extract import extract
|
||||
from .manifest import validate_document
|
||||
from .sanitize import sanitize
|
||||
from .walk import walk
|
||||
|
||||
HAVE_AGE = bool(shutil.which("age") and shutil.which("age-keygen"))
|
||||
|
||||
|
||||
def manifest(root, seal_to=(), stages=None):
|
||||
doc = {
|
||||
"freeagent": 1, "id": "t", "owner": {"host": "h"},
|
||||
"scope": {"roots": [root]},
|
||||
"mask": ["read", "search", "index"],
|
||||
"stages": stages or ["walk", "extract", "sanitize", "catalog", "seal"], "sweep": "manual",
|
||||
"sanitize": {"mode": "flag"},
|
||||
"sink": {"catalog": "host", "vectors": "local", "plaintext": "never", "seal_to": list(seal_to)},
|
||||
}
|
||||
return validate_document(doc)
|
||||
|
||||
|
||||
def keygen(dirpath):
|
||||
ident = os.path.join(dirpath, "key.txt")
|
||||
out = subprocess.run(["age-keygen", "-o", ident], capture_output=True, text=True, check=True)
|
||||
pub = [l for l in (out.stderr + out.stdout).splitlines() if "age1" in l][-1].split()[-1]
|
||||
return ident, pub
|
||||
|
||||
|
||||
class Case(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.root = tempfile.mkdtemp(prefix="fa-seal-")
|
||||
self.out = tempfile.mkdtemp(prefix="fa-sealout-")
|
||||
with open(os.path.join(self.root, "note.txt"), "w", encoding="utf-8") as fh:
|
||||
fh.write("Marisol's email is m@example.com and card 4111 1111 1111 1111.\n")
|
||||
m = manifest(self.root)
|
||||
walk(m, self.out)
|
||||
extract(m, self.out, use_external=False)
|
||||
sanitize(m, self.out, env={})
|
||||
catalog(m, self.out, env={})
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.root, ignore_errors=True)
|
||||
shutil.rmtree(self.out, ignore_errors=True)
|
||||
|
||||
|
||||
class FailClosedTests(Case):
|
||||
def test_no_owner_recipient(self):
|
||||
# age is faked present: seal checks for it first, and this test is
|
||||
# about the recipient refusal, whether or not age is installed
|
||||
with mock.patch.object(sl, "tool_path", lambda n: "/usr/bin/" + n):
|
||||
with self.assertRaises(sl.SealError) as cm:
|
||||
sl.seal(manifest(self.root), self.out, env={})
|
||||
self.assertIn("FREEAGENT_OWNER_RECIPIENT", str(cm.exception))
|
||||
self.assertFalse(os.path.exists(os.path.join(self.out, "sealed")))
|
||||
|
||||
def test_bad_recipient_strings(self):
|
||||
env = {"FREEAGENT_OWNER_RECIPIENT": "not-a-key"}
|
||||
with mock.patch.object(sl, "tool_path", lambda n: "/usr/bin/" + n):
|
||||
with self.assertRaises(sl.SealError) as cm:
|
||||
sl.seal(manifest(self.root), self.out, env=env)
|
||||
self.assertIn("not an age recipient", str(cm.exception))
|
||||
env = {"FREEAGENT_OWNER_RECIPIENT": "age1" + "q" * 58}
|
||||
with self.assertRaises(sl.SealError) as cm:
|
||||
sl.recipients_for(manifest(self.root, seal_to=["partner-key"]), env)
|
||||
self.assertIn("partner-key", str(cm.exception))
|
||||
|
||||
def test_no_age_on_path(self):
|
||||
env = {"FREEAGENT_OWNER_RECIPIENT": "age1" + "q" * 58}
|
||||
with mock.patch.object(sl, "tool_path", lambda n: None):
|
||||
with self.assertRaises(sl.SealError) as cm:
|
||||
sl.seal(manifest(self.root), self.out, env=env)
|
||||
self.assertIn("age", str(cm.exception))
|
||||
|
||||
def test_seal_not_in_stages(self):
|
||||
# (a `catalog: host` manifest cannot even omit seal — the schema refuses it)
|
||||
m = validate_document({
|
||||
"freeagent": 1, "id": "t", "owner": {"host": "h"}, "scope": {"roots": [self.root]},
|
||||
"mask": ["read", "index"], "stages": ["walk", "extract", "sanitize", "catalog"],
|
||||
"sweep": "manual", "sanitize": {"mode": "flag"},
|
||||
"sink": {"catalog": "local", "vectors": "local", "plaintext": "never"}})
|
||||
with self.assertRaises(sl.SealError):
|
||||
sl.seal(m, self.out, env={"FREEAGENT_OWNER_RECIPIENT": "age1" + "q" * 58})
|
||||
|
||||
def test_bundle_allowlist(self):
|
||||
plain, inv = sl.build_bundle(manifest(self.root), self.out, "20260828T000000Z")
|
||||
with tarfile.open(fileobj=io.BytesIO(plain)) as tar:
|
||||
names = tar.getnames()
|
||||
self.assertIn("catalog.jsonl", names)
|
||||
self.assertIn("catalog.dirs.jsonl", names)
|
||||
self.assertIn("manifest.json", names)
|
||||
self.assertIn("seal.manifest.json", names)
|
||||
self.assertIn("walk.summary.json", names)
|
||||
self.assertIn("sanitize.summary.json", names)
|
||||
for forbidden in ("walk.jsonl", "extract.jsonl", "sanitize.jsonl", "vectors.db", "vault.db"):
|
||||
self.assertFalse(any(n.endswith(forbidden) for n in names), forbidden)
|
||||
self.assertFalse(any(n.startswith(("extract/", "sanitize/")) for n in names))
|
||||
# nothing in the bundle is a value
|
||||
self.assertNotIn(b"m@example.com", plain)
|
||||
self.assertNotIn(b"4111", plain)
|
||||
self.assertIn(b"[PII_EMAIL_1]", plain) # sanitized head is there, as designed
|
||||
self.assertEqual(set(inv["members"]), set(names))
|
||||
|
||||
def test_bundle_reproducible(self):
|
||||
a, _ = sl.build_bundle(manifest(self.root), self.out, "20260828T000000Z")
|
||||
b, _ = sl.build_bundle(manifest(self.root), self.out, "20260828T000000Z")
|
||||
self.assertEqual(a, b)
|
||||
|
||||
|
||||
@unittest.skipUnless(HAVE_AGE, "age / age-keygen not on PATH")
|
||||
class RoundTripTests(Case):
|
||||
def test_seal_and_unseal(self):
|
||||
ident, pub = keygen(self.out)
|
||||
other_ident, other_pub = keygen(tempfile.mkdtemp(prefix="fa-k2-"))
|
||||
env = {"FREEAGENT_OWNER_RECIPIENT": pub}
|
||||
s = sl.seal(manifest(self.root, seal_to=[other_pub]), self.out, env=env, now="20260828T000000Z")
|
||||
sealed = os.path.join(self.out, s["file"])
|
||||
self.assertTrue(os.path.exists(sealed))
|
||||
self.assertEqual(s["recipients"], [pub, other_pub])
|
||||
self.assertGreater(s["bytes_sealed"], 0)
|
||||
self.assertIsNone(s["signature"])
|
||||
with open(sealed, "rb") as fh:
|
||||
blob = fh.read()
|
||||
self.assertTrue(blob.startswith(b"age-encryption.org/v1"))
|
||||
self.assertNotIn(b"PII_EMAIL", blob) # sealed: the catalog text is not visible
|
||||
with open(os.path.join(self.out, "sealed", "latest.json")) as fh:
|
||||
self.assertEqual(json.load(fh)["file"], s["file"])
|
||||
# both recipients can open it; a stranger cannot
|
||||
for who in (ident, other_ident):
|
||||
dest = tempfile.mkdtemp(prefix="fa-un-")
|
||||
inv = sl.unseal(sealed, who, dest)
|
||||
self.assertEqual(inv["manifest_id"], "t")
|
||||
with open(os.path.join(dest, "catalog.jsonl"), encoding="utf-8") as fh:
|
||||
self.assertIn("[PII_EMAIL_1]", fh.read())
|
||||
shutil.rmtree(dest, ignore_errors=True)
|
||||
stranger, _ = keygen(tempfile.mkdtemp(prefix="fa-k3-"))
|
||||
with self.assertRaises(sl.SealError):
|
||||
sl.unseal(sealed, stranger, tempfile.mkdtemp(prefix="fa-un-"))
|
||||
|
||||
def test_owner_recipient_deduped_and_signature(self):
|
||||
ident, pub = keygen(self.out)
|
||||
key = os.path.join(self.out, "sign.pem")
|
||||
subprocess.run(["openssl", "ecparam", "-name", "prime256v1", "-genkey", "-noout", "-out", key],
|
||||
check=True, capture_output=True)
|
||||
env = {"FREEAGENT_OWNER_RECIPIENT": pub, "SEAL_SIGN_KEY": key}
|
||||
s = sl.seal(manifest(self.root, seal_to=[pub]), self.out, env=env)
|
||||
self.assertEqual(s["recipients"], [pub])
|
||||
sig = os.path.join(self.out, "sealed", s["signature"])
|
||||
self.assertTrue(os.path.exists(sig))
|
||||
# the signature verifies against the public half
|
||||
pubkey = os.path.join(self.out, "sign.pub")
|
||||
subprocess.run(["openssl", "ec", "-in", key, "-pubout", "-out", pubkey],
|
||||
check=True, capture_output=True)
|
||||
v = subprocess.run(["openssl", "dgst", "-sha256", "-verify", pubkey, "-signature", sig,
|
||||
os.path.join(self.out, s["file"])], capture_output=True, text=True)
|
||||
self.assertEqual(v.returncode, 0, v.stderr)
|
||||
|
||||
def test_tampered_member_refused(self):
|
||||
ident, pub = keygen(self.out)
|
||||
env = {"FREEAGENT_OWNER_RECIPIENT": pub}
|
||||
s = sl.seal(manifest(self.root), self.out, env=env)
|
||||
# rebuild a bundle with a modified catalog but the old inventory, seal it, and unseal
|
||||
plain, _ = sl.build_bundle(manifest(self.root), self.out, "x")
|
||||
buf = io.BytesIO()
|
||||
with tarfile.open(fileobj=io.BytesIO(plain)) as src, tarfile.open(fileobj=buf, mode="w") as dst:
|
||||
for m in src.getmembers():
|
||||
data = src.extractfile(m).read()
|
||||
if m.name == "catalog.jsonl":
|
||||
data = data + b"\n{\"forged\": true}\n"
|
||||
m.size = len(data)
|
||||
dst.addfile(m, io.BytesIO(data))
|
||||
bad = os.path.join(self.out, "bad.tar.age")
|
||||
subprocess.run(["age", "-r", pub, "-o", bad], input=buf.getvalue(), check=True)
|
||||
with self.assertRaises(sl.SealError) as cm:
|
||||
sl.unseal(bad, ident, tempfile.mkdtemp(prefix="fa-un-"))
|
||||
self.assertIn("does not match the inventory", str(cm.exception))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,351 @@
|
||||
"""Tests for the query endpoint. Run: python3 -m unittest discover .
|
||||
|
||||
A miniature PKI is minted with openssl per test class: a root CA, a
|
||||
server cert, an owner client cert, a grantee client cert, an expired-
|
||||
grant client cert, and a client cert from a DIFFERENT CA. The server
|
||||
binds loopback here ONLY because client and server are the same test
|
||||
process; production binds the mesh interface (docs/serve.md).
|
||||
"""
|
||||
|
||||
import http.client
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import ssl
|
||||
import subprocess
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
|
||||
from . import embed as em
|
||||
from . import run as rn
|
||||
from . import serve as sv
|
||||
from .catalog import catalog
|
||||
from .extract import extract
|
||||
from .manifest import validate_document
|
||||
from .sanitize import sanitize
|
||||
from .test_embed import _Stub as EmbedStub # deterministic bag-of-words embedder
|
||||
from .walk import walk
|
||||
|
||||
HAVE_OPENSSL = bool(shutil.which("openssl"))
|
||||
import datetime as _dt
|
||||
TODAY = _dt.date(2026, 8, 28)
|
||||
|
||||
|
||||
def sh(*argv):
|
||||
subprocess.run(list(argv), check=True, capture_output=True)
|
||||
|
||||
|
||||
def mint_pki(d):
|
||||
"""root CA + server + clients; returns dict of paths."""
|
||||
p = lambda n: os.path.join(d, n)
|
||||
def root(name, cn):
|
||||
sh("openssl", "ecparam", "-name", "prime256v1", "-genkey", "-noout", "-out", p(name + ".key"))
|
||||
sh("openssl", "req", "-new", "-x509", "-key", p(name + ".key"), "-out", p(name + ".crt"), "-days", "2",
|
||||
"-subj", "/CN=" + cn, "-addext", "basicConstraints=critical,CA:TRUE",
|
||||
"-addext", "keyUsage=critical,keyCertSign,cRLSign")
|
||||
|
||||
root("ca", "test-root")
|
||||
root("other-ca", "other-root")
|
||||
|
||||
def leaf(name, cn, ca="ca", san=None):
|
||||
sh("openssl", "ecparam", "-name", "prime256v1", "-genkey", "-noout", "-out", p(name + ".key"))
|
||||
sh("openssl", "req", "-new", "-key", p(name + ".key"), "-out", p(name + ".csr"), "-subj", "/CN=" + cn)
|
||||
ext = p(name + ".ext")
|
||||
with open(ext, "w") as fh:
|
||||
fh.write("subjectAltName=%s\nbasicConstraints=CA:FALSE\n"
|
||||
"extendedKeyUsage=serverAuth,clientAuth\n" % (san or "DNS:" + cn))
|
||||
sh("openssl", "x509", "-req", "-in", p(name + ".csr"), "-CA", p(ca + ".crt"), "-CAkey", p(ca + ".key"),
|
||||
"-CAcreateserial", "-out", p(name + ".crt"), "-days", "2", "-extfile", ext)
|
||||
return p(name + ".crt"), p(name + ".key")
|
||||
|
||||
out = {"ca": p("ca.crt"), "other_ca": p("other-ca.crt")}
|
||||
out["server"] = leaf("server", "localhost", san="DNS:localhost,IP:127.0.0.1")
|
||||
out["owner"] = leaf("owner", "owner-host")
|
||||
out["partner"] = leaf("partner", "partner-host")
|
||||
out["expired"] = leaf("expired", "old-host")
|
||||
out["stranger"] = leaf("stranger", "owner-host", ca="other-ca") # right CN, wrong CA
|
||||
out["nobody"] = leaf("nobody", "nobody-host")
|
||||
return out
|
||||
|
||||
|
||||
@unittest.skipUnless(HAVE_OPENSSL, "openssl not on PATH")
|
||||
class ServeTests(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.work = tempfile.mkdtemp(prefix="fa-serve-")
|
||||
cls.pki = mint_pki(cls.work)
|
||||
cls.root = os.path.join(cls.work, "corpus")
|
||||
os.makedirs(cls.root)
|
||||
for name, text in (("apples.txt", "apples pears plums orchard harvest m@example.com\n"),
|
||||
("routers.txt", "router switch firewall vlan network segment\n")):
|
||||
with open(os.path.join(cls.root, name), "w", encoding="utf-8") as fh:
|
||||
fh.write(text)
|
||||
cls.out = os.path.join(cls.work, "out")
|
||||
cls.manifest = os.path.join(cls.work, "m.yaml")
|
||||
cls.write_manifest(grants=[{"host": "partner-host", "ops": ["search"], "until": "2026-12-31"},
|
||||
{"host": "old-host", "ops": ["search", "read_range"], "until": "2026-01-01"}])
|
||||
m = cls.load()
|
||||
# embedding stub
|
||||
from http.server import ThreadingHTTPServer
|
||||
EmbedStub.dim = 64
|
||||
cls.estub = ThreadingHTTPServer(("127.0.0.1", 0), EmbedStub)
|
||||
threading.Thread(target=cls.estub.serve_forever, daemon=True).start()
|
||||
cls.env = {"EMBED_URL": "http://127.0.0.1:%d" % cls.estub.server_address[1],
|
||||
"FREEAGENT_TLS_CERT": cls.pki["server"][0], "FREEAGENT_TLS_KEY": cls.pki["server"][1],
|
||||
"FREEAGENT_CA_ROOT": cls.pki["ca"]}
|
||||
walk(m, cls.out)
|
||||
extract(m, cls.out, use_external=False)
|
||||
sanitize(m, cls.out, env={})
|
||||
em.chunk(m, cls.out)
|
||||
em.embed(m, cls.out, env=cls.env)
|
||||
catalog(m, cls.out, env={})
|
||||
cls.srv = sv.build_server(cls.manifest, cls.out, env=cls.env, bind="127.0.0.1:0", today=TODAY)
|
||||
cls.port = cls.srv.server_address[1]
|
||||
threading.Thread(target=cls.srv.serve_forever, daemon=True).start()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls.srv.shutdown()
|
||||
cls.srv.server_close()
|
||||
cls.estub.shutdown()
|
||||
shutil.rmtree(cls.work, ignore_errors=True)
|
||||
|
||||
@classmethod
|
||||
def write_manifest(cls, grants):
|
||||
import yaml
|
||||
doc = {"freeagent": 1, "id": "t", "owner": {"host": "owner-host"},
|
||||
"scope": {"roots": [cls.root]}, "mask": ["read", "search", "index"],
|
||||
"stages": ["walk", "extract", "sanitize", "chunk", "embed", "catalog", "seal"],
|
||||
"sanitize": {"mode": "flag"},
|
||||
"embed": {"model": "stub-model", "dim": 64, "chunk": {"size_tokens": 8, "overlap_tokens": 2}},
|
||||
"sweep": "manual",
|
||||
"sink": {"catalog": "host", "vectors": "local", "plaintext": "never"},
|
||||
"grants": grants}
|
||||
with open(cls.manifest, "w", encoding="utf-8") as fh:
|
||||
yaml.safe_dump(doc, fh, sort_keys=False)
|
||||
rn.sign_manifest_file(cls.manifest, cls.pki["owner"][1], cls.pki["owner"][0])
|
||||
|
||||
@classmethod
|
||||
def load(cls):
|
||||
from .manifest import load_manifest
|
||||
return load_manifest(cls.manifest)
|
||||
|
||||
# ── client ───────────────────────────────────────────────────────────
|
||||
def client(self, who, ca=None):
|
||||
ctx = ssl.create_default_context(cafile=ca or self.pki["ca"])
|
||||
ctx.check_hostname = False # test server cert is for localhost/127.0.0.1; keep simple
|
||||
if who:
|
||||
ctx.load_cert_chain(*self.pki[who])
|
||||
return http.client.HTTPSConnection("127.0.0.1", self.port, context=ctx, timeout=10)
|
||||
|
||||
def call(self, who, method, path, body=None):
|
||||
c = self.client(who)
|
||||
data = json.dumps(body).encode() if body is not None else None
|
||||
c.request(method, path, body=data, headers={"Content-Type": "application/json"})
|
||||
r = c.getresponse()
|
||||
out = json.loads(r.read().decode("utf-8"))
|
||||
c.close()
|
||||
return r.status, out
|
||||
|
||||
def audit(self):
|
||||
with open(os.path.join(self.out, "serve.audit.jsonl"), encoding="utf-8") as fh:
|
||||
return [json.loads(l) for l in fh]
|
||||
|
||||
# ── tests ────────────────────────────────────────────────────────────
|
||||
def test_owner_everything(self):
|
||||
st, body = self.call("owner", "GET", "/health")
|
||||
self.assertEqual((st, body["ok"], body["serving"]), (200, True, True))
|
||||
st, body = self.call("owner", "POST", "/search", {"query": "firewall vlan", "k": 2})
|
||||
self.assertEqual(st, 200)
|
||||
self.assertEqual(len(body["hits"]), 2)
|
||||
self.assertIn("router", body["hits"][0]["text"])
|
||||
doc = body["hits"][0]["doc"]
|
||||
st, body = self.call("owner", "POST", "/read_range", {"doc": doc, "start": 0, "end": 20})
|
||||
self.assertEqual(st, 200)
|
||||
self.assertEqual(body["text"], "router switch firewa")
|
||||
self.assertEqual(body["total"], 44)
|
||||
st, body = self.call("owner", "GET", "/catalog")
|
||||
self.assertEqual(st, 200)
|
||||
self.assertEqual({c["path"] for c in body["cards"]}, {"apples.txt", "routers.txt"})
|
||||
# sanitized store only: the apples card head carries a token, not the address
|
||||
apples = [c for c in body["cards"] if c["path"] == "apples.txt"][0]
|
||||
self.assertIn("[PII_EMAIL_1]", apples["head"])
|
||||
st, body = self.call("owner", "GET", "/status")
|
||||
self.assertEqual(st, 200)
|
||||
# the map: directory cards, root "" aggregates everything
|
||||
st, body = self.call("owner", "GET", "/catalog/dirs")
|
||||
self.assertEqual((st, body["kind"]), (200, "dirs"))
|
||||
root = [c for c in body["cards"] if c["path"] == ""][0]
|
||||
self.assertEqual(root["files"], 2)
|
||||
self.assertEqual(root["entities"], {"EMAIL": 1})
|
||||
self.assertEqual(self.call("partner", "GET", "/catalog/dirs")[0], 403)
|
||||
# paging: total/offset/limit/next_offset; prefix filters by path
|
||||
st, body = self.call("owner", "GET", "/catalog?limit=1")
|
||||
self.assertEqual((st, body["total"], len(body["cards"]), body["next_offset"]), (200, 2, 1, 1))
|
||||
st, body = self.call("owner", "GET", "/catalog?limit=1&offset=1")
|
||||
self.assertEqual((len(body["cards"]), body["next_offset"]), (1, None))
|
||||
st, body = self.call("owner", "GET", "/catalog?prefix=rou")
|
||||
self.assertEqual([c["path"] for c in body["cards"]], ["routers.txt"])
|
||||
self.assertEqual(self.call("owner", "GET", "/catalog?limit=x")[0], 400)
|
||||
|
||||
def test_grantee_gets_only_its_ops(self):
|
||||
st, body = self.call("partner", "POST", "/search", {"query": "apples"})
|
||||
self.assertEqual(st, 200)
|
||||
self.assertIn("[PII_EMAIL_1]", body["hits"][0]["text"]) # sanitized, always
|
||||
doc = body["hits"][0]["doc"]
|
||||
st, body = self.call("partner", "POST", "/read_range", {"doc": doc, "start": 0, "end": 5})
|
||||
self.assertEqual(st, 403)
|
||||
self.assertIn("partner-host", body["error"])
|
||||
self.assertEqual(self.call("partner", "GET", "/catalog")[0], 403)
|
||||
self.assertEqual(self.call("partner", "GET", "/status")[0], 403)
|
||||
self.assertEqual(self.call("partner", "GET", "/health")[0], 200)
|
||||
|
||||
def test_expired_grant_and_unknown_cn(self):
|
||||
self.assertEqual(self.call("expired", "POST", "/search", {"query": "x"})[0], 403)
|
||||
self.assertEqual(self.call("nobody", "POST", "/search", {"query": "x"})[0], 403)
|
||||
self.assertEqual(self.call("nobody", "GET", "/health")[0], 200) # chained, so alive
|
||||
|
||||
def test_no_cert_or_foreign_ca_never_reaches_a_handler(self):
|
||||
with self.assertRaises((ssl.SSLError, ConnectionError, OSError)):
|
||||
self.call(None, "GET", "/health")
|
||||
with self.assertRaises((ssl.SSLError, ConnectionError, OSError)):
|
||||
self.call("stranger", "GET", "/health") # CN 'owner-host' but signed by another CA
|
||||
# and no audit line was written for either
|
||||
self.assertFalse(any(a["cn"] is None for a in self.audit()))
|
||||
|
||||
def test_read_range_is_bounded_and_sanitized_only(self):
|
||||
st, body = self.call("owner", "POST", "/search", {"query": "apples", "k": 1})
|
||||
doc = body["hits"][0]["doc"]
|
||||
st, body = self.call("owner", "POST", "/read_range", {"doc": doc, "start": 0, "end": 10 ** 9})
|
||||
self.assertEqual(st, 200)
|
||||
self.assertLessEqual(body["end"] - body["start"], sv.READ_RANGE_MAX)
|
||||
self.assertNotIn("m@example.com", body["text"])
|
||||
# one id per document: the extract and sanitized stores share the
|
||||
# content hash, and the route reads ONLY the sanitized directory —
|
||||
# so the extract-side plaintext for the same id is never what comes back
|
||||
with open(os.path.join(self.out, "extract.jsonl"), encoding="utf-8") as fh:
|
||||
ext_sha = [json.loads(l) for l in fh if "apples" in l][0]["sha256"]
|
||||
st, body = self.call("owner", "POST", "/read_range", {"doc": ext_sha, "start": 0, "end": 200})
|
||||
self.assertEqual(st, 200)
|
||||
self.assertIn("[PII_EMAIL_1]", body["text"])
|
||||
self.assertNotIn("m@example.com", body["text"])
|
||||
with open(os.path.join(self.out, "extract", ext_sha[:2], ext_sha + ".txt"), encoding="utf-8") as fh:
|
||||
self.assertIn("m@example.com", fh.read()) # the plaintext exists on disk; it has no route
|
||||
st, _ = self.call("owner", "POST", "/read_range", {"doc": "0" * 64, "start": 0, "end": 10})
|
||||
self.assertEqual(st, 404)
|
||||
self.assertEqual(self.call("owner", "POST", "/read_range", {"doc": "zz", "start": 0, "end": 1})[0], 400)
|
||||
self.assertEqual(self.call("owner", "POST", "/search", {"query": ""})[0], 400)
|
||||
self.assertEqual(self.call("owner", "POST", "/nope", {})[0], 404)
|
||||
|
||||
def test_silent_tcp_connection_does_not_stall_other_clients(self):
|
||||
# with the handshake on the accept loop one idle raw TCP
|
||||
# connection blocked every client; now the handshake runs in the
|
||||
# connection's own thread with its own timeout
|
||||
import socket
|
||||
staller = socket.create_connection(("127.0.0.1", self.port))
|
||||
try:
|
||||
st, body = self.call("owner", "GET", "/health")
|
||||
self.assertEqual((st, body["ok"]), (200, True))
|
||||
finally:
|
||||
staller.close()
|
||||
|
||||
def test_audit_hashes_query_never_logs_it(self):
|
||||
self.call("owner", "POST", "/search", {"query": "secret question about marisol"})
|
||||
lines = self.audit()
|
||||
last = [a for a in lines if a["op"] == "search"][-1]
|
||||
self.assertEqual(last["cn"], "owner-host")
|
||||
self.assertEqual(len(last["q"]), 16)
|
||||
blob = json.dumps(lines)
|
||||
self.assertNotIn("marisol", blob)
|
||||
|
||||
def test_revoke_grant_by_resigning(self):
|
||||
self.assertEqual(self.call("partner", "POST", "/search", {"query": "apples"})[0], 200)
|
||||
self.write_manifest(grants=[]) # drop the grant, re-sign
|
||||
self.assertEqual(self.call("partner", "POST", "/search", {"query": "apples"})[0], 403)
|
||||
self.assertEqual(self.call("owner", "POST", "/search", {"query": "apples"})[0], 200)
|
||||
# tamper without re-signing -> the endpoint refuses everything but health
|
||||
with open(self.manifest, encoding="utf-8") as fh:
|
||||
text = fh.read()
|
||||
with open(self.manifest, "w", encoding="utf-8") as fh:
|
||||
fh.write(text.replace("sweep: manual", "sweep: daily"))
|
||||
st, body = self.call("owner", "POST", "/search", {"query": "apples"})
|
||||
self.assertEqual(st, 503)
|
||||
self.assertIn("gate", body["error"])
|
||||
self.assertEqual(self.call("owner", "GET", "/health")[1]["serving"], False)
|
||||
# restore
|
||||
self.write_manifest(grants=[{"host": "partner-host", "ops": ["search"], "until": "2026-12-31"},
|
||||
{"host": "old-host", "ops": ["search", "read_range"], "until": "2026-01-01"}])
|
||||
self.assertEqual(self.call("owner", "POST", "/search", {"query": "apples"})[0], 200)
|
||||
|
||||
|
||||
class UnitTests(unittest.TestCase):
|
||||
def test_policy_fails_closed_when_the_manifest_cannot_be_read(self):
|
||||
# Review: an unreadable manifest (PermissionError, bad encoding) is
|
||||
# not a ManifestError, and it escaped after _mtime advanced — leaving
|
||||
# the PREVIOUS manifest in force.
|
||||
from unittest import mock
|
||||
work = tempfile.mkdtemp(prefix="fa-pol-")
|
||||
self.addCleanup(shutil.rmtree, work, ignore_errors=True)
|
||||
man = os.path.join(work, "m.yaml")
|
||||
with open(man, "w") as fh:
|
||||
fh.write("x: 1\n")
|
||||
good = object()
|
||||
with mock.patch.object(sv, "load_manifest", return_value=good), \
|
||||
mock.patch.object(sv, "verify_signature", return_value={}):
|
||||
pol = sv.Policy(man, {})
|
||||
self.assertIs(pol.manifest, good)
|
||||
os.utime(man, ns=(1, 1))
|
||||
with mock.patch.object(sv, "load_manifest", side_effect=PermissionError("denied")):
|
||||
pol.reload()
|
||||
self.assertIsNone(pol.manifest)
|
||||
self.assertIn("manifest unreadable", pol.gate_error)
|
||||
|
||||
# skipped rather than faked: the port file is written by build_server
|
||||
# only after it has loaded a real certificate chain and verified a signed
|
||||
# manifest, and minting those needs openssl
|
||||
@unittest.skipUnless(HAVE_OPENSSL, "openssl not on PATH: needed to mint the test PKI and sign the manifest")
|
||||
def test_port_file_written_after_bind(self):
|
||||
# a supervisor binds :0 and reads the port back — one atomic file
|
||||
work = tempfile.mkdtemp(prefix="fa-pf-")
|
||||
try:
|
||||
pki = mint_pki(work)
|
||||
man = os.path.join(work, "m.yaml")
|
||||
import yaml
|
||||
with open(man, "w") as fh:
|
||||
yaml.safe_dump({"freeagent": 1, "id": "p", "owner": {"host": "owner-host"},
|
||||
"scope": {"roots": [work]}, "mask": ["read"], "stages": ["walk"],
|
||||
"sweep": "manual",
|
||||
"sink": {"catalog": "local", "vectors": "local", "plaintext": "never"}}, fh)
|
||||
rn.sign_manifest_file(man, pki["owner"][1], pki["owner"][0])
|
||||
env = {"FREEAGENT_TLS_CERT": pki["server"][0], "FREEAGENT_TLS_KEY": pki["server"][1],
|
||||
"FREEAGENT_CA_ROOT": pki["ca"]}
|
||||
pf = os.path.join(work, "port")
|
||||
srv = sv.build_server(man, os.path.join(work, "out"), env=env, bind="127.0.0.1:0", port_file=pf)
|
||||
try:
|
||||
with open(pf) as fh:
|
||||
self.assertEqual(int(fh.read().strip()), srv.server_address[1])
|
||||
finally:
|
||||
srv.server_close()
|
||||
finally:
|
||||
shutil.rmtree(work, ignore_errors=True)
|
||||
|
||||
def test_parse_bind(self):
|
||||
self.assertEqual(sv.parse_bind("[::]:8470"), ("::", 8470))
|
||||
self.assertEqual(sv.parse_bind("10.0.0.5:9"), ("10.0.0.5", 9))
|
||||
self.assertEqual(sv.parse_bind("[fd00:1::31]:8470"), ("fd00:1::31", 8470))
|
||||
with self.assertRaises(sv.ServeError):
|
||||
sv.parse_bind("nope")
|
||||
|
||||
def test_peer_cn(self):
|
||||
self.assertEqual(sv.peer_cn({"subject": ((("commonName", "x"),),)}), "x")
|
||||
self.assertIsNone(sv.peer_cn({}))
|
||||
self.assertIsNone(sv.peer_cn(None))
|
||||
|
||||
def test_ssl_context_requires_all_three(self):
|
||||
with self.assertRaises(sv.ServeError):
|
||||
sv.build_ssl_context("", "", "")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,213 @@
|
||||
"""Tests for the release verifier. Run: python3 -m unittest discover ."""
|
||||
|
||||
import base64
|
||||
import contextlib
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
from . import build as bd
|
||||
from . import verify_release as vr
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
REPO = os.path.dirname(HERE)
|
||||
|
||||
|
||||
def _sign(seed, msg):
|
||||
"""RFC 8032 signing, for the tests only: the verifier never signs."""
|
||||
h = hashlib.sha512(seed).digest()
|
||||
a = int.from_bytes(h[:32], "little")
|
||||
a &= (1 << 254) - 8
|
||||
a |= 1 << 254
|
||||
pub = vr._compress(vr._mul(a, vr._G))
|
||||
r = int.from_bytes(hashlib.sha512(h[32:] + msg).digest(), "little") % vr._Q
|
||||
big_r = vr._compress(vr._mul(r, vr._G))
|
||||
k = int.from_bytes(hashlib.sha512(big_r + pub + msg).digest(), "little") % vr._Q
|
||||
return big_r + ((r + k * a) % vr._Q).to_bytes(32, "little")
|
||||
|
||||
|
||||
class Ed25519Tests(unittest.TestCase):
|
||||
# RFC 8032, section 7.1, tests 1 and 2.
|
||||
RFC = (
|
||||
("9d61b19deffd5a60ba844af492ec2cc44449c5697b326919703bac031cae7f60",
|
||||
"d75a980182b10ab7d54bfed3c964073a0ee172f3daa62325af021a68f707511a",
|
||||
"",
|
||||
"e5564300c360ac729086e2cc806e828a84877f1eb8e5d974d873e065224901555fb8821590a33bacc61e39701cf9b46bd25bf5f0595bbe24655141438e7a100b"),
|
||||
("4ccd089b28ff96da9db6c346ec114e0f5b8a319f35aba624da8cf6ed4fb8a6fb",
|
||||
"3d4017c3e843895a92b70aa74d1b7ebc9c982ccf2ec4968cc0cd55f12af4660c",
|
||||
"72",
|
||||
"92a009a9f0d4cab8720e820b5f642540a2b27b5416503f8fb3762223ebdb69da085ac1e43e15996e458f3613d0f11d8c387b2eaeb4302aeeb00d291612bb0c00"),
|
||||
)
|
||||
|
||||
def test_rfc8032_vectors(self):
|
||||
for seed, pub, msg, sig in self.RFC:
|
||||
seed, pub, msg, sig = (bytes.fromhex(x) for x in (seed, pub, msg, sig))
|
||||
self.assertEqual(vr.ed25519_public_key(seed), pub)
|
||||
self.assertTrue(vr.ed25519_verify(pub, msg, sig))
|
||||
self.assertEqual(_sign(seed, msg), sig)
|
||||
|
||||
def test_any_changed_bit_is_refused(self):
|
||||
seed, pub, msg, sig = (bytes.fromhex(x) for x in self.RFC[1])
|
||||
self.assertFalse(vr.ed25519_verify(pub, msg + b"\x00", sig))
|
||||
for i in (0, 31, 32, 63):
|
||||
bad = bytearray(sig)
|
||||
bad[i] ^= 1
|
||||
self.assertFalse(vr.ed25519_verify(pub, msg, bytes(bad)))
|
||||
other = vr.ed25519_public_key(b"\x02" * 32)
|
||||
self.assertFalse(vr.ed25519_verify(other, msg, sig))
|
||||
|
||||
def test_malformed_input_is_false_not_an_error(self):
|
||||
_, pub, msg, sig = (bytes.fromhex(x) for x in self.RFC[0])
|
||||
self.assertFalse(vr.ed25519_verify(pub[:31], msg, sig))
|
||||
self.assertFalse(vr.ed25519_verify(pub, msg, sig[:63]))
|
||||
# S at or above the group order is refused (signature malleability).
|
||||
s = int.from_bytes(sig[32:], "little") + vr._Q
|
||||
if s < 1 << 256:
|
||||
self.assertFalse(vr.ed25519_verify(pub, msg, sig[:32] + s.to_bytes(32, "little")))
|
||||
# A y coordinate that is on no point of the curve.
|
||||
self.assertFalse(vr.ed25519_verify(b"\x02" + b"\x00" * 31, msg, sig))
|
||||
|
||||
|
||||
class StatementTests(unittest.TestCase):
|
||||
def test_statement_bytes(self):
|
||||
self.assertEqual(vr.statement("bastion-agent", "0.1.13", "ab" * 32),
|
||||
b"privatae-release-signature-v1\nbastion-agent\n0.1.13\n" + b"ab" * 32 + b"\n")
|
||||
|
||||
def test_the_signing_tool_and_this_check_agree(self):
|
||||
# The vector the bastion's release-signing tool and the bastion share:
|
||||
# seed of 32 bytes of 1, this manifest, kind "bastion", this signature.
|
||||
# If the statement format drifted here, this would fail.
|
||||
pub = base64.b64encode(vr.ed25519_public_key(b"\x01" * 32)).decode()
|
||||
body = b'{"version":"1.2.3","binaries":{}}'
|
||||
sig = "tr5h3oXjBDmwvwQiZSriW4T6Y6oBFQDymfofLEGXMdTJSXfoQAKN+ajhCs89FO9vdCEZH3Y8mYzYXPZ7RzZVDw=="
|
||||
m, i = vr.check_signature(body, sig, keys=(pub,), kind="bastion")
|
||||
self.assertEqual((m["version"], i), ("1.2.3", 0))
|
||||
# Bound to its kind: a bastion manifest's signature is not an agent's.
|
||||
with self.assertRaises(vr.VerifyError):
|
||||
vr.check_signature(body, sig, keys=(pub,), kind="bastion-agent")
|
||||
|
||||
def test_signature_file_name(self):
|
||||
body = b'{"version":"0.1.13"}'
|
||||
self.assertEqual(vr.sig_name(body),
|
||||
"version.json.%s.sig" % hashlib.sha256(body).hexdigest()[:16])
|
||||
|
||||
def test_pinned_keys_are_two_valid_points(self):
|
||||
self.assertEqual(len(vr.RELEASE_KEYS), 2)
|
||||
for k in vr.RELEASE_KEYS:
|
||||
raw = base64.b64decode(k, validate=True)
|
||||
self.assertEqual(len(raw), 32)
|
||||
self.assertIsNotNone(vr._decompress(raw))
|
||||
|
||||
def test_refusals(self):
|
||||
seed = b"\x03" * 32
|
||||
pub = base64.b64encode(vr.ed25519_public_key(seed)).decode()
|
||||
|
||||
def signed(body):
|
||||
msg = vr.statement("bastion-agent", json.loads(body)["version"], hashlib.sha256(body).hexdigest())
|
||||
return base64.b64encode(_sign(seed, msg)).decode()
|
||||
|
||||
body = b'{"version":"0.2.0"}'
|
||||
self.assertEqual(vr.check_signature(body, signed(body), keys=(pub,))[1], 0)
|
||||
# The key that verified is reported by position: the spare is second.
|
||||
self.assertEqual(vr.check_signature(body, signed(body), keys=(vr.RELEASE_KEYS[0], pub))[1], 1)
|
||||
cases = (
|
||||
(b"not json", "AAAA"),
|
||||
(b'{"version":""}', signed(body)),
|
||||
(b'{"version":"0.2.0\\n"}', signed(body)),
|
||||
(body, "not base64!"),
|
||||
(body, base64.b64encode(b"short").decode()),
|
||||
(body + b" ", signed(body)), # the manifest changed after signing
|
||||
(body, signed(body)), # the production keys did not sign it
|
||||
)
|
||||
for i, (b, s) in enumerate(cases):
|
||||
keys = None if i == len(cases) - 1 else (pub,)
|
||||
with self.subTest(i=i), self.assertRaises(vr.VerifyError):
|
||||
vr.check_signature(b, s, keys=keys)
|
||||
|
||||
|
||||
@unittest.skipUnless(shutil.which("git"), "needs git")
|
||||
class EndToEndTests(unittest.TestCase):
|
||||
"""The whole check against a throwaway repository holding this tree."""
|
||||
|
||||
SEED = b"\x04" * 32
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.tmp = tempfile.mkdtemp(prefix="agent-verify-test-")
|
||||
cls.repo = os.path.join(cls.tmp, "repo")
|
||||
ignore = shutil.ignore_patterns("__pycache__", "*.pyc")
|
||||
shutil.copytree(HERE, os.path.join(cls.repo, "bastion_agent"), ignore=ignore)
|
||||
shutil.copytree(os.path.join(REPO, "pii"), os.path.join(cls.repo, "pii"), ignore=ignore)
|
||||
shutil.copy(os.path.join(REPO, "LICENSE"), cls.repo)
|
||||
git = ("git", "-C", cls.repo, "-c", "user.name=test", "-c", "user.email=test@example.invalid",
|
||||
"-c", "commit.gpgsign=false", "-c", "tag.gpgsign=false")
|
||||
run = lambda *a: subprocess.run(git + a, check=True, stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE, universal_newlines=True).stdout.strip()
|
||||
subprocess.run(("git", "init", "-q", cls.repo), check=True)
|
||||
run("add", "-A")
|
||||
run("commit", "-q", "-m", "release")
|
||||
cls.commit = run("rev-parse", "HEAD")
|
||||
cls.run_git = staticmethod(run)
|
||||
cls.version = bd.version()
|
||||
# The same files as this tree, so the same bytes as building this tree.
|
||||
cls.digest = hashlib.sha256(bd.build_bytes()).hexdigest()
|
||||
cls.pub = base64.b64encode(vr.ed25519_public_key(cls.SEED)).decode()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
shutil.rmtree(cls.tmp, ignore_errors=True)
|
||||
|
||||
def manifest(self, built_from, digest=None):
|
||||
name = "bastion-agent-%s.pyz" % self.version
|
||||
return json.dumps({"version": self.version, "built_from": built_from,
|
||||
"files": {name: {"sha256": digest or self.digest},
|
||||
name + ".sha256": {"sha256": "0" * 64}}}).encode()
|
||||
|
||||
def verify(self, body, seed=None):
|
||||
msg = vr.statement("bastion-agent", self.version, hashlib.sha256(body).hexdigest())
|
||||
sig = base64.b64encode(_sign(seed or self.SEED, msg))
|
||||
served = {"/version.json": body, "/" + vr.sig_name(body): sig}
|
||||
out, err = io.StringIO(), io.StringIO()
|
||||
with mock.patch.object(vr, "fetch", lambda url: served[url[len("https://test"):]]), \
|
||||
contextlib.redirect_stdout(out), contextlib.redirect_stderr(err):
|
||||
code = vr.main(["--base", "https://test"], repo=self.repo, keys=(self.pub,))
|
||||
return code, out.getvalue() + err.getvalue()
|
||||
|
||||
def test_tag_first_then_built_from(self):
|
||||
code, out = self.verify(self.manifest("0000000"))
|
||||
self.assertEqual(code, 1, out)
|
||||
self.assertIn("neither the tag v%s" % self.version, out)
|
||||
|
||||
code, out = self.verify(self.manifest(self.commit[:7]))
|
||||
self.assertEqual(code, 0, out)
|
||||
self.assertIn("source: built_from", out)
|
||||
self.assertIn("matches %s" % self.digest, out)
|
||||
|
||||
self.run_git("tag", "v" + self.version)
|
||||
try:
|
||||
code, out = self.verify(self.manifest("0000000"))
|
||||
self.assertEqual(code, 0, out)
|
||||
self.assertIn("source: tag v%s" % self.version, out)
|
||||
finally:
|
||||
self.run_git("tag", "-d", "v" + self.version)
|
||||
|
||||
def test_a_different_build_is_reported(self):
|
||||
code, out = self.verify(self.manifest(self.commit, digest="1" * 64))
|
||||
self.assertEqual(code, 1, out)
|
||||
self.assertIn("DIFFERS", out)
|
||||
|
||||
def test_a_foreign_signature_stops_before_any_checkout(self):
|
||||
code, out = self.verify(self.manifest(self.commit), seed=b"\x05" * 32)
|
||||
self.assertEqual(code, 1, out)
|
||||
self.assertIn("does not verify", out)
|
||||
self.assertNotIn("source:", out)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,271 @@
|
||||
"""Unit tests for the walk stage. Run: python3 -m unittest discover .
|
||||
|
||||
Builds a real temp tree so the walker is exercised on whatever OS runs the
|
||||
suite; Windows-only behaviours (junctions, casefold) are covered by the
|
||||
pieces that can be tested everywhere (the exclude matcher, _is_link fallback)
|
||||
and by the platform guard in each test.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
from .manifest import validate_document
|
||||
from .walk import Excluder, Walker, _glob_to_regex, load_hash_cache, walk
|
||||
|
||||
IS_WINDOWS = os.name == "nt"
|
||||
|
||||
|
||||
def manifest(root, **scope):
|
||||
doc = {
|
||||
"freeagent": 1, "id": "t", "owner": {"host": "h"},
|
||||
"scope": dict({"roots": [root]}, **scope),
|
||||
"mask": ["read"], "stages": ["walk"], "sweep": "manual",
|
||||
"sink": {"catalog": "local", "vectors": "local", "plaintext": "never"},
|
||||
}
|
||||
return validate_document(doc)
|
||||
|
||||
|
||||
def touch(path, data=b""):
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
with open(path, "wb") as fh:
|
||||
fh.write(data)
|
||||
|
||||
|
||||
def read_jsonl(path):
|
||||
with open(path, encoding="utf-8") as fh:
|
||||
return [json.loads(line) for line in fh]
|
||||
|
||||
|
||||
class TreeCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.root = tempfile.mkdtemp(prefix="fa-walk-")
|
||||
self.out = tempfile.mkdtemp(prefix="fa-out-")
|
||||
touch(os.path.join(self.root, "a.txt"), b"hello\n")
|
||||
touch(os.path.join(self.root, "docs", "b.pdf"), b"%PDF-1.7 fake")
|
||||
touch(os.path.join(self.root, "docs", "c.bin"), b"\x00\x01\x02")
|
||||
touch(os.path.join(self.root, "docs", "empty"), b"")
|
||||
touch(os.path.join(self.root, "node_modules", "x", "y.js"), b"1")
|
||||
touch(os.path.join(self.root, "secret.key"), b"k")
|
||||
touch(os.path.join(self.root, "report.docx"), b"PK\x03\x04zip")
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.root, ignore_errors=True)
|
||||
shutil.rmtree(self.out, ignore_errors=True)
|
||||
|
||||
def by_path(self, entries):
|
||||
return {e["path"]: e for e in entries}
|
||||
|
||||
|
||||
class WalkTests(TreeCase):
|
||||
def test_map_shape_and_sorting(self):
|
||||
s = walk(manifest(self.root), self.out, now="2026-08-28T00:00:00Z")
|
||||
entries = read_jsonl(os.path.join(self.out, "walk.jsonl"))
|
||||
paths = [e["path"] for e in entries]
|
||||
self.assertEqual(paths, sorted(paths))
|
||||
m = self.by_path(entries)
|
||||
self.assertEqual(m["a.txt"]["kind"], "file")
|
||||
self.assertEqual(m["a.txt"]["size"], 6)
|
||||
self.assertEqual(m["a.txt"]["mime"], "text/plain")
|
||||
self.assertEqual(m["a.txt"]["sha256"], hashlib.sha256(b"hello\n").hexdigest())
|
||||
self.assertTrue(m["a.txt"]["mtime"].endswith("Z"))
|
||||
self.assertEqual(m["docs"]["kind"], "dir")
|
||||
self.assertEqual(m["docs/b.pdf"]["mime"], "application/pdf")
|
||||
self.assertEqual(m["docs/c.bin"]["mime"], "application/octet-stream")
|
||||
self.assertEqual(m["docs/empty"]["mime"], "inode/x-empty")
|
||||
self.assertEqual(m["report.docx"]["mime"],
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document")
|
||||
self.assertEqual(s["entries"], len(entries))
|
||||
self.assertEqual(s["counts"]["files"], 7)
|
||||
self.assertEqual(s["counts"]["errors"], 0)
|
||||
self.assertEqual(s["walked_at"], "2026-08-28T00:00:00Z")
|
||||
self.assertTrue(s["manifest_digest"].startswith("sha256:"))
|
||||
self.assertEqual(s["roots"][0]["resolved"], os.path.abspath(self.root))
|
||||
# paths are '/'-separated on every OS
|
||||
self.assertTrue(all("\\" not in e["path"] for e in entries))
|
||||
|
||||
def test_agent_dir_is_always_excluded(self):
|
||||
touch(os.path.join(self.root, ".agent", "catalog.jsonl"), b"{}")
|
||||
touch(os.path.join(self.root, "sub", ".agent", "x"), b"x")
|
||||
touch(os.path.join(self.root, "other", ".Agent", "y"), b"y") # case-insensitive volumes
|
||||
walk(manifest(self.root), self.out)
|
||||
paths = {e["path"] for e in read_jsonl(os.path.join(self.out, "walk.jsonl"))}
|
||||
self.assertFalse(any(".agent" in p.lower() for p in paths), paths)
|
||||
|
||||
def test_excludes_prune(self):
|
||||
m = manifest(self.root, exclude=["node_modules", "*.key", "docs/c.*"])
|
||||
s = walk(m, self.out)
|
||||
paths = {e["path"] for e in read_jsonl(os.path.join(self.out, "walk.jsonl"))}
|
||||
self.assertNotIn("node_modules", paths)
|
||||
self.assertNotIn("node_modules/x/y.js", paths)
|
||||
self.assertNotIn("secret.key", paths)
|
||||
self.assertNotIn("docs/c.bin", paths)
|
||||
self.assertIn("docs/b.pdf", paths)
|
||||
self.assertEqual(s["counts"]["excluded"], 3)
|
||||
|
||||
def test_too_large_is_metadata_only(self):
|
||||
m = manifest(self.root, max_file_bytes=4)
|
||||
walk(m, self.out)
|
||||
e = self.by_path(read_jsonl(os.path.join(self.out, "walk.jsonl")))
|
||||
self.assertEqual(e["a.txt"]["skip"], "too_large")
|
||||
self.assertIsNone(e["a.txt"]["sha256"])
|
||||
self.assertEqual(e["a.txt"]["mime"], "text/plain") # still sniffed
|
||||
self.assertEqual(e["docs/c.bin"]["sha256"], hashlib.sha256(b"\x00\x01\x02").hexdigest())
|
||||
|
||||
def test_no_hash(self):
|
||||
s = walk(manifest(self.root), self.out, do_hash=False)
|
||||
self.assertEqual(s["counts"]["hashed"], 0)
|
||||
self.assertFalse(s["hashing"])
|
||||
e = self.by_path(read_jsonl(os.path.join(self.out, "walk.jsonl")))
|
||||
self.assertIsNone(e["a.txt"]["sha256"])
|
||||
|
||||
def test_hash_cache_reuse_and_invalidation(self):
|
||||
walk(manifest(self.root), self.out)
|
||||
s2 = walk(manifest(self.root), self.out)
|
||||
self.assertEqual(s2["counts"]["hashed"], 0)
|
||||
self.assertEqual(s2["counts"]["hash_reused"], 7)
|
||||
# a content change with a new mtime is re-hashed
|
||||
p = os.path.join(self.root, "a.txt")
|
||||
touch(p, b"changed\n")
|
||||
os.utime(p, ns=(1, 1_700_000_000_000_000_000))
|
||||
s3 = walk(manifest(self.root), self.out)
|
||||
self.assertEqual(s3["counts"]["hashed"], 1)
|
||||
e = self.by_path(read_jsonl(os.path.join(self.out, "walk.jsonl")))
|
||||
self.assertEqual(e["a.txt"]["sha256"], hashlib.sha256(b"changed\n").hexdigest())
|
||||
|
||||
def test_deterministic_output(self):
|
||||
walk(manifest(self.root), self.out)
|
||||
first = open(os.path.join(self.out, "walk.jsonl"), "rb").read()
|
||||
walk(manifest(self.root), self.out)
|
||||
self.assertEqual(first, open(os.path.join(self.out, "walk.jsonl"), "rb").read())
|
||||
|
||||
def test_missing_root_is_an_error_not_a_crash(self):
|
||||
missing = os.path.join(self.root, "nope")
|
||||
m = manifest(missing)
|
||||
s = walk(m, self.out)
|
||||
self.assertEqual(s["counts"]["errors"], 1)
|
||||
self.assertIn("error", s["roots"][0])
|
||||
self.assertEqual(s["entries"], 0)
|
||||
|
||||
def test_unreadable_dir_recorded(self):
|
||||
real = os.scandir
|
||||
|
||||
def flaky(path):
|
||||
if os.path.basename(path) == "docs":
|
||||
raise PermissionError(13, "Permission denied")
|
||||
return real(path)
|
||||
|
||||
with mock.patch("bastion_agent.walk.os.scandir", flaky):
|
||||
s = walk(manifest(self.root), self.out)
|
||||
e = self.by_path(read_jsonl(os.path.join(self.out, "walk.jsonl")))
|
||||
self.assertEqual(s["counts"]["errors"], 1)
|
||||
self.assertIn("Permission denied", e["docs"]["error"])
|
||||
self.assertNotIn("docs/b.pdf", e)
|
||||
|
||||
def test_unicode_paths(self):
|
||||
touch(os.path.join(self.root, "résumé — 履歴書.txt"), b"x")
|
||||
walk(manifest(self.root), self.out)
|
||||
e = self.by_path(read_jsonl(os.path.join(self.out, "walk.jsonl")))
|
||||
self.assertIn("résumé — 履歴書.txt", e)
|
||||
|
||||
def test_home_root_expands(self):
|
||||
with mock.patch.dict(os.environ, {"HOME": self.root, "USERPROFILE": self.root}):
|
||||
s = walk(manifest("~"), self.out)
|
||||
self.assertEqual(s["roots"][0]["resolved"], os.path.abspath(self.root))
|
||||
self.assertGreater(s["entries"], 0)
|
||||
|
||||
@unittest.skipIf(IS_WINDOWS, "symlink creation needs a privilege on Windows")
|
||||
def test_symlink_recorded_not_followed(self):
|
||||
os.symlink(os.path.join(self.root, "docs"), os.path.join(self.root, "link-to-docs"))
|
||||
os.symlink("a.txt", os.path.join(self.root, "link-to-a"))
|
||||
s = walk(manifest(self.root), self.out)
|
||||
e = self.by_path(read_jsonl(os.path.join(self.out, "walk.jsonl")))
|
||||
self.assertEqual(e["link-to-docs"]["kind"], "symlink")
|
||||
self.assertEqual(e["link-to-a"]["target"], "a.txt")
|
||||
self.assertNotIn("link-to-docs/b.pdf", e)
|
||||
self.assertEqual(s["counts"]["symlinks"], 2)
|
||||
|
||||
@unittest.skipIf(IS_WINDOWS, "symlink creation needs a privilege on Windows")
|
||||
def test_follow_symlinks_when_granted(self):
|
||||
os.symlink(os.path.join(self.root, "docs"), os.path.join(self.root, "link-to-docs"))
|
||||
walk(manifest(self.root, follow_symlinks=True), self.out)
|
||||
e = self.by_path(read_jsonl(os.path.join(self.out, "walk.jsonl")))
|
||||
self.assertEqual(e["link-to-docs"]["kind"], "dir")
|
||||
self.assertIn("link-to-docs/b.pdf", e)
|
||||
|
||||
|
||||
class ExcludeTests(unittest.TestCase):
|
||||
def test_glob_translation(self):
|
||||
import re
|
||||
rx = lambda p: re.compile("^" + _glob_to_regex(p) + "$")
|
||||
self.assertTrue(rx("**/node_modules/**").match("a/b/node_modules/c/d.js"))
|
||||
self.assertTrue(rx("**/node_modules/**").match("node_modules/c"))
|
||||
self.assertTrue(rx("*.key").match("id.key"))
|
||||
self.assertFalse(rx("*.key").match("a/id.key"))
|
||||
self.assertTrue(rx("docs/*.pdf").match("docs/x.pdf"))
|
||||
self.assertFalse(rx("docs/*.pdf").match("docs/sub/x.pdf"))
|
||||
self.assertTrue(rx("docs/**").match("docs/sub/x.pdf"))
|
||||
self.assertTrue(rx("[!a]?.txt").match("bx.txt"))
|
||||
self.assertFalse(rx("[!a]?.txt").match("ax.txt"))
|
||||
|
||||
def test_double_star_dir_pattern_prunes_the_dir(self):
|
||||
ex = Excluder(["**/.git/**"])
|
||||
self.assertTrue(ex.matches(".git"))
|
||||
self.assertTrue(ex.matches("a/.git"))
|
||||
self.assertTrue(ex.matches("a/.git/HEAD"))
|
||||
self.assertFalse(ex.matches("a/.gitignore"))
|
||||
|
||||
def test_segment_patterns_match_at_any_depth(self):
|
||||
ex = Excluder(["node_modules", "*.key", ".git"])
|
||||
self.assertTrue(ex.matches("node_modules"))
|
||||
self.assertTrue(ex.matches("a/b/node_modules"))
|
||||
self.assertTrue(ex.matches("deep/id.key"))
|
||||
self.assertTrue(ex.matches("proj/.git"))
|
||||
self.assertFalse(ex.matches("proj/gitignore"))
|
||||
|
||||
def test_path_patterns_are_anchored(self):
|
||||
ex = Excluder(["docs/private/**", "**/*.tmp"])
|
||||
self.assertTrue(ex.matches("docs/private/x"))
|
||||
self.assertTrue(ex.matches("docs/private")) # the dir itself is pruned
|
||||
self.assertFalse(ex.matches("docs"))
|
||||
self.assertFalse(ex.matches("other/docs/private/x"))
|
||||
self.assertTrue(ex.matches("a.tmp"))
|
||||
self.assertTrue(ex.matches("a/b/c.tmp"))
|
||||
|
||||
def test_backslash_patterns_normalised(self):
|
||||
ex = Excluder(["docs\\private\\**"])
|
||||
self.assertTrue(ex.matches("docs/private/x"))
|
||||
|
||||
def test_case_sensitivity_matches_os(self):
|
||||
ex = Excluder(["*.PDF"])
|
||||
self.assertEqual(ex.matches("x.pdf"), IS_WINDOWS)
|
||||
|
||||
|
||||
class CacheTests(unittest.TestCase):
|
||||
def test_cache_ignores_garbage_and_non_files(self):
|
||||
fd, path = tempfile.mkstemp(suffix=".jsonl")
|
||||
with os.fdopen(fd, "w") as fh:
|
||||
fh.write('{"root":0,"path":"a","kind":"file","size":1,"mtime_ns":2,"sha256":"h"}\n')
|
||||
fh.write('{"root":0,"path":"d","kind":"dir"}\n')
|
||||
fh.write("not json\n")
|
||||
try:
|
||||
c = load_hash_cache(path)
|
||||
finally:
|
||||
os.unlink(path)
|
||||
self.assertEqual(c, {(0, "a"): (1, 2, "h")})
|
||||
|
||||
def test_missing_cache_is_empty(self):
|
||||
self.assertEqual(load_hash_cache("/nonexistent/walk.jsonl"), {})
|
||||
|
||||
|
||||
# The Python 3.8 syntax floor is enforced for this package by
|
||||
# test_syntax_floor.py (shared with pii).
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,139 @@
|
||||
"""The manifest gate needs no openssl (x509gate).
|
||||
|
||||
A Windows machine ships no openssl, and the gate used to shell out to it: the
|
||||
agent refused every manifest there and never ran. These tests mint real
|
||||
certificates with openssl (a test-machine tool only) and then run the gate with openssl REMOVED from PATH, and cover
|
||||
the managed install's signing lane — the bastion signs `payload` with its
|
||||
own key and records it with `sign --value`.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import datetime as _dt
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
from . import run as _run
|
||||
from . import x509gate as x
|
||||
from .manifest import load_manifest
|
||||
from .test_run import write_manifest
|
||||
|
||||
HAVE_OPENSSL = bool(shutil.which("openssl"))
|
||||
|
||||
|
||||
|
||||
def sh(*argv, data=None):
|
||||
return subprocess.run(argv, input=data, check=True, stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE).stdout
|
||||
|
||||
|
||||
def _der_sign(key, data):
|
||||
return sh("openssl", "dgst", "-sha256", "-sign", key, data=data)
|
||||
|
||||
|
||||
@unittest.skipUnless(HAVE_OPENSSL, "openssl is needed to MINT the test certificates")
|
||||
class GateWithoutOpenssl(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.d = tempfile.mkdtemp()
|
||||
self.addCleanup(shutil.rmtree, self.d, ignore_errors=True)
|
||||
p = lambda n: os.path.join(self.d, n) # noqa: E731
|
||||
self.p = p
|
||||
sh("openssl", "ecparam", "-name", "prime256v1", "-genkey", "-noout", "-out", p("host.key"))
|
||||
sh("openssl", "req", "-new", "-x509", "-key", p("host.key"), "-out", p("host.crt"),
|
||||
"-days", "2", "-subj", "/CN=desk-example")
|
||||
os.makedirs(p("root"))
|
||||
write_manifest(p("m.yaml"), p("root"), ["walk"], cn="desk-example")
|
||||
self.env = {"FREEAGENT_CA_ROOT": p("host.crt")}
|
||||
|
||||
def _no_openssl(self):
|
||||
return mock.patch.object(shutil, "which", lambda name, *a, **k: None)
|
||||
|
||||
def _bastion_sign(self):
|
||||
"""What the managed install does: sign `payload` with its own key and
|
||||
record it with `sign --value` — no openssl on the machine."""
|
||||
der = _der_sign(self.p("host.key"), _run.manifest_payload(self.p("m.yaml")))
|
||||
with self._no_openssl():
|
||||
_run.write_signature(self.p("m.yaml"), der, self.p("host.crt"))
|
||||
|
||||
def test_gate_verifies_with_openssl_removed_from_path(self):
|
||||
self._bastion_sign()
|
||||
with self._no_openssl():
|
||||
m = load_manifest(self.p("m.yaml"), require_signed=True)
|
||||
got = _run.verify_signature(m, self.env, manifest_dir=self.d)
|
||||
self.assertEqual(got["cn"], "desk-example")
|
||||
self.assertTrue(got["chained"])
|
||||
|
||||
def test_tampered_manifest_refused_without_openssl(self):
|
||||
self._bastion_sign()
|
||||
with open(self.p("m.yaml")) as fh:
|
||||
text = fh.read()
|
||||
with open(self.p("m.yaml"), "w") as fh:
|
||||
fh.write(text.replace("sweep: manual", "sweep: hourly", 1))
|
||||
with self._no_openssl():
|
||||
m = load_manifest(self.p("m.yaml"), require_signed=True)
|
||||
with self.assertRaises(_run.GateError) as cm:
|
||||
_run.verify_signature(m, self.env, manifest_dir=self.d)
|
||||
self.assertIn("does not verify", str(cm.exception))
|
||||
|
||||
def test_cli_sign_value_records_a_bastion_signature(self):
|
||||
der = _der_sign(self.p("host.key"), _run.manifest_payload(self.p("m.yaml")))
|
||||
env = dict(os.environ, PATH="/nonexistent") # no openssl for the CLI either
|
||||
r = subprocess.run(
|
||||
[shutil.which("python3") or "python3", "-m", "bastion_agent", "sign", self.p("m.yaml"),
|
||||
"--value", base64.b64encode(der).decode(), "--cert", self.p("host.crt")],
|
||||
cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||||
env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
self.assertEqual(r.returncode, 0, r.stderr)
|
||||
with self._no_openssl():
|
||||
m = load_manifest(self.p("m.yaml"), require_signed=True)
|
||||
_run.verify_signature(m, self.env, manifest_dir=self.d)
|
||||
|
||||
def test_expired_cert_refused(self):
|
||||
self._bastion_sign()
|
||||
leaf = x.load_cert_file(self.p("host.crt"))
|
||||
later = _dt.datetime.now(_dt.timezone.utc) + _dt.timedelta(days=30)
|
||||
with self.assertRaises(x.CertError):
|
||||
x.verify_chain(leaf, [leaf], now=later)
|
||||
|
||||
def test_cert_signed_by_another_ca_refused(self):
|
||||
p = self.p
|
||||
for n in ("ca", "other"):
|
||||
sh("openssl", "ecparam", "-name", "prime256v1", "-genkey", "-noout", "-out", p(n + ".key"))
|
||||
sh("openssl", "req", "-new", "-x509", "-key", p(n + ".key"), "-out", p(n + ".crt"),
|
||||
"-days", "2", "-subj", "/CN=Same Name CA")
|
||||
sh("openssl", "ecparam", "-name", "prime256v1", "-genkey", "-noout", "-out", p("leaf.key"))
|
||||
sh("openssl", "req", "-new", "-key", p("leaf.key"), "-out", p("leaf.csr"), "-subj", "/CN=desk-example")
|
||||
sh("openssl", "x509", "-req", "-in", p("leaf.csr"), "-CA", p("other.crt"), "-CAkey", p("other.key"),
|
||||
"-CAcreateserial", "-out", p("leaf.crt"), "-days", "2")
|
||||
leaf = x.load_cert_file(p("leaf.crt"))
|
||||
ca = x.load_cert_file(p("ca.crt"))
|
||||
# Same subject, different key: the name matches, the signature must not.
|
||||
with self.assertRaises(x.CertError):
|
||||
x.verify_chain(leaf, [ca])
|
||||
x.verify_chain(leaf, [x.load_cert_file(p("other.crt"))])
|
||||
|
||||
def test_signature_bit_flips_never_verify(self):
|
||||
leaf = x.load_cert_file(self.p("host.crt"))
|
||||
for i in range(8):
|
||||
msg = os.urandom(40 + i)
|
||||
der = _der_sign(self.p("host.key"), msg)
|
||||
self.assertTrue(x.ecdsa_p256_sha256_verify(leaf.public_key, msg, der))
|
||||
bad = bytearray(der)
|
||||
bad[-1 - i] ^= 0x01
|
||||
self.assertFalse(x.ecdsa_p256_sha256_verify(leaf.public_key, msg, bytes(bad)))
|
||||
self.assertFalse(x.ecdsa_p256_sha256_verify(leaf.public_key, msg + b"x", der))
|
||||
|
||||
|
||||
class DerSafety(unittest.TestCase):
|
||||
def test_garbage_is_a_cert_error_not_a_crash(self):
|
||||
for junk in (b"", b"\x30", b"\x30\x84\xff\xff\xff\xff", os.urandom(64)):
|
||||
with self.assertRaises(x.CertError):
|
||||
x.Cert(junk)
|
||||
self.assertFalse(x.ecdsa_p256_sha256_verify(x._G, b"m", b"\x30\x00"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,318 @@
|
||||
"""Rebuild the published bastion agent from source and check it is what was published.
|
||||
|
||||
python3 -m bastion_agent.verify_release [--base URL]
|
||||
|
||||
Run it from a clone of this repository. It needs Python 3.8 or later, git,
|
||||
network access to the download address, and PyYAML at the version build.py
|
||||
pins (the build vendors it). It needs no key or account of ours.
|
||||
|
||||
It fetches the release manifest (version.json) and the signature beside it,
|
||||
checks the signature against the release keys the bastion pins, checks out the
|
||||
source of that version in a separate worktree, rebuilds the .pyz with that
|
||||
version's own build.py, and compares the sha256 with the manifest's. A match
|
||||
means the file every bastion downloads is exactly what this source produces,
|
||||
and the signature says the release key holder published it.
|
||||
|
||||
The source is the tag v<version>. The public repository is a snapshot of the
|
||||
development one, so the commit the manifest names (built_from) exists only
|
||||
there; the tag marks the same tree in the public history. When the tag is
|
||||
absent and built_from is present (a clone of the development repository),
|
||||
built_from is used instead.
|
||||
|
||||
The signature check is ed25519 in pure Python, because the standard library
|
||||
has no ed25519 and this has to run where nothing else is installed. It only
|
||||
verifies public data, so speed and side channels do not matter here; the code
|
||||
follows the reference implementation in RFC 8032, section 6.
|
||||
|
||||
This module is a tool for checking releases. build.py leaves it out of the
|
||||
.pyz: a verifier shipped inside the file it verifies would prove nothing.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import urllib.request
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
REPO = os.path.dirname(HERE)
|
||||
DEFAULT_BASE = "https://privatae.ai/bastion-agent"
|
||||
KIND = "bastion-agent"
|
||||
|
||||
# The public halves of the release signing keys: base64 of the raw 32-byte
|
||||
# ed25519 key, the primary first and then the spare kept offline. They are the
|
||||
# keys the bastion compiles in (productionReleaseKeys in client-go/release_sign.go
|
||||
# in the bastion's repository), because a bastion refuses any agent manifest
|
||||
# that does not verify against those. When the bastion rotates a key, this list
|
||||
# changes in the same way, or this check fails where a bastion would pass.
|
||||
RELEASE_KEYS = (
|
||||
"tWOyZm38DwTwbawSnvExQG/CQsVJBKYWc46arCBZj9E=",
|
||||
"CKL+J0VSdqrfY4AmxChvylJ1ha4e8GvcrYF2o17Ah+g=",
|
||||
)
|
||||
|
||||
FETCH_TIMEOUT = 30
|
||||
FETCH_MAX = 1 << 20 # a manifest or a signature is a few kilobytes
|
||||
|
||||
|
||||
class VerifyError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
# ---- ed25519 verification (RFC 8032, section 6) -----------------------------
|
||||
|
||||
_P = 2 ** 255 - 19
|
||||
_Q = 2 ** 252 + 27742317777372353535851937790883648493
|
||||
_D = -121665 * pow(121666, _P - 2, _P) % _P
|
||||
_SQRT_M1 = pow(2, (_P - 1) // 4, _P)
|
||||
|
||||
|
||||
def _inv(x):
|
||||
return pow(x, _P - 2, _P)
|
||||
|
||||
|
||||
def _add(a, b):
|
||||
# Extended homogeneous coordinates (X, Y, Z, T).
|
||||
ea = (a[1] - a[0]) * (b[1] - b[0]) % _P
|
||||
eb = (a[1] + a[0]) * (b[1] + b[0]) % _P
|
||||
ec = 2 * a[3] * b[3] * _D % _P
|
||||
ed = 2 * a[2] * b[2] % _P
|
||||
e, f, g, h = eb - ea, ed - ec, ed + ec, eb + ea
|
||||
return (e * f % _P, g * h % _P, f * g % _P, e * h % _P)
|
||||
|
||||
|
||||
def _mul(s, pt):
|
||||
acc = (0, 1, 1, 0)
|
||||
while s > 0:
|
||||
if s & 1:
|
||||
acc = _add(acc, pt)
|
||||
pt = _add(pt, pt)
|
||||
s >>= 1
|
||||
return acc
|
||||
|
||||
|
||||
def _equal(a, b):
|
||||
return (a[0] * b[2] - b[0] * a[2]) % _P == 0 and (a[1] * b[2] - b[1] * a[2]) % _P == 0
|
||||
|
||||
|
||||
def _recover_x(y, sign):
|
||||
if y >= _P:
|
||||
return None
|
||||
x2 = (y * y - 1) * _inv(_D * y * y + 1)
|
||||
if x2 == 0:
|
||||
return None if sign else 0
|
||||
x = pow(x2, (_P + 3) // 8, _P)
|
||||
if (x * x - x2) % _P != 0:
|
||||
x = x * _SQRT_M1 % _P
|
||||
if (x * x - x2) % _P != 0:
|
||||
return None
|
||||
if (x & 1) != sign:
|
||||
x = _P - x
|
||||
return x
|
||||
|
||||
|
||||
def _decompress(b):
|
||||
if len(b) != 32:
|
||||
return None
|
||||
y = int.from_bytes(b, "little")
|
||||
sign = y >> 255
|
||||
y &= (1 << 255) - 1
|
||||
x = _recover_x(y, sign)
|
||||
if x is None:
|
||||
return None
|
||||
return (x, y, 1, x * y % _P)
|
||||
|
||||
|
||||
def _compress(pt):
|
||||
zi = _inv(pt[2])
|
||||
x, y = pt[0] * zi % _P, pt[1] * zi % _P
|
||||
return (y | ((x & 1) << 255)).to_bytes(32, "little")
|
||||
|
||||
|
||||
_GY = 4 * _inv(5) % _P
|
||||
_G = (_recover_x(_GY, 0), _GY, 1, _recover_x(_GY, 0) * _GY % _P)
|
||||
|
||||
|
||||
def ed25519_verify(public, msg, sig):
|
||||
"""True when sig is a valid ed25519 signature of msg by the 32-byte key public."""
|
||||
if len(public) != 32 or len(sig) != 64:
|
||||
return False
|
||||
a = _decompress(public)
|
||||
r = _decompress(sig[:32])
|
||||
if a is None or r is None:
|
||||
return False
|
||||
s = int.from_bytes(sig[32:], "little")
|
||||
if s >= _Q:
|
||||
return False
|
||||
h = int.from_bytes(hashlib.sha512(sig[:32] + public + msg).digest(), "little") % _Q
|
||||
return _equal(_mul(s, _G), _add(r, _mul(h, a)))
|
||||
|
||||
|
||||
def ed25519_public_key(seed):
|
||||
"""The public key for a 32-byte seed. Used by the tests to check the
|
||||
arithmetic against published vectors; verification never needs it."""
|
||||
h = hashlib.sha512(seed).digest()
|
||||
a = int.from_bytes(h[:32], "little")
|
||||
a &= (1 << 254) - 8
|
||||
a |= 1 << 254
|
||||
return _compress(_mul(a, _G))
|
||||
|
||||
|
||||
# ---- the release signature --------------------------------------------------
|
||||
|
||||
def statement(kind, version, manifest_sha256_hex):
|
||||
"""The exact text a release signature covers. The bastion builds the same
|
||||
text (releaseStatement in client-go/release_sign.go), and so does the
|
||||
signing tool the key holder runs."""
|
||||
return ("privatae-release-signature-v1\n%s\n%s\n%s\n"
|
||||
% (kind, version, manifest_sha256_hex)).encode("utf-8")
|
||||
|
||||
|
||||
def sig_name(body):
|
||||
"""The signature file's name beside version.json: named by the first 16 hex
|
||||
of the manifest's sha256, so a signature can only pair with its manifest."""
|
||||
return "version.json.%s.sig" % hashlib.sha256(body).hexdigest()[:16]
|
||||
|
||||
|
||||
def check_signature(body, sig_text, keys=None, kind=KIND):
|
||||
"""Return the manifest (parsed) when sig_text verifies over body against
|
||||
one of keys, and the index of the key that verified. Raise VerifyError
|
||||
otherwise; there is no unsigned fallback."""
|
||||
try:
|
||||
m = json.loads(body.decode("utf-8"))
|
||||
except (UnicodeDecodeError, ValueError):
|
||||
raise VerifyError("the manifest is not JSON")
|
||||
ver = m.get("version") if isinstance(m, dict) else None
|
||||
if not isinstance(ver, str) or not ver or "\n" in ver or "\r" in ver:
|
||||
raise VerifyError("the manifest names no usable version")
|
||||
try:
|
||||
raw = base64.b64decode(sig_text.strip(), validate=True)
|
||||
except (ValueError, TypeError):
|
||||
raise VerifyError("the signature is not base64")
|
||||
if len(raw) != 64:
|
||||
raise VerifyError("the signature is malformed")
|
||||
msg = statement(kind, ver, hashlib.sha256(body).hexdigest())
|
||||
for i, k in enumerate(RELEASE_KEYS if keys is None else keys):
|
||||
pub = base64.b64decode(k)
|
||||
if ed25519_verify(pub, msg, raw):
|
||||
return m, i
|
||||
raise VerifyError("the signature does not verify against any pinned release key")
|
||||
|
||||
|
||||
# ---- fetching, checking out, rebuilding -------------------------------------
|
||||
|
||||
def fetch(url):
|
||||
# A named client: the download address sits behind Cloudflare, whose
|
||||
# Browser Integrity Check refuses Python's default user agent (error
|
||||
# 1010, a 403), the signature generic scrapers send.
|
||||
req = urllib.request.Request(url, headers={"User-Agent": "bastion-agent-verify-release"})
|
||||
with urllib.request.urlopen(req, timeout=FETCH_TIMEOUT) as resp:
|
||||
data = resp.read(FETCH_MAX + 1)
|
||||
if len(data) > FETCH_MAX:
|
||||
raise VerifyError("%s is larger than a release manifest can be" % url)
|
||||
return data
|
||||
|
||||
|
||||
def _git(repo, *args, check=True):
|
||||
return subprocess.run(("git", "-C", repo) + args, check=check,
|
||||
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
||||
universal_newlines=True)
|
||||
|
||||
|
||||
def _commit(repo, ref):
|
||||
r = _git(repo, "rev-parse", "-q", "--verify", ref + "^{commit}", check=False)
|
||||
return r.stdout.strip() if r.returncode == 0 else ""
|
||||
|
||||
|
||||
def pick_source(repo, version, built_from):
|
||||
"""(label, commit) of the source to rebuild: the tag v<version>, else the
|
||||
manifest's built_from, else VerifyError."""
|
||||
tag = "v" + version
|
||||
# Fetch the tag if the clone lacks it; a failure (offline, no such tag)
|
||||
# just leaves the local refs to decide.
|
||||
_git(repo, "fetch", "-q", "origin", "refs/tags/%s:refs/tags/%s" % (tag, tag), check=False)
|
||||
c = _commit(repo, "refs/tags/" + tag)
|
||||
if c:
|
||||
return "tag " + tag, c
|
||||
if built_from:
|
||||
c = _commit(repo, built_from)
|
||||
if c:
|
||||
return "built_from " + built_from, c
|
||||
raise VerifyError("this clone has neither the tag %s nor the commit %s. Check that "
|
||||
"the public repository has the tag (git fetch --tags origin); if it "
|
||||
"does not, the release predates tagging and can only be checked from a "
|
||||
"clone that has its built_from commit" % (tag, built_from or "(none named)"))
|
||||
|
||||
|
||||
def rebuild(repo, commit, version, work):
|
||||
"""sha256 of the .pyz that commit's own build.py produces."""
|
||||
src = os.path.join(work, "src")
|
||||
_git(repo, "worktree", "add", "-q", "--detach", src, commit)
|
||||
try:
|
||||
with open(os.path.join(src, "bastion_agent", "VERSION"), encoding="utf-8") as fh:
|
||||
got = fh.read().strip()
|
||||
if got != version:
|
||||
raise VerifyError("the source says version %s, the manifest %s" % (got, version))
|
||||
dist = os.path.join(work, "dist")
|
||||
r = subprocess.run((sys.executable, "-m", "bastion_agent.build", "--out", dist),
|
||||
cwd=src, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
||||
universal_newlines=True)
|
||||
if r.returncode != 0:
|
||||
raise VerifyError("the build failed: %s" % (r.stderr.strip() or r.stdout.strip()))
|
||||
with open(os.path.join(dist, "bastion-agent-%s.pyz" % version), "rb") as fh:
|
||||
return hashlib.sha256(fh.read()).hexdigest()
|
||||
finally:
|
||||
_git(repo, "worktree", "remove", "--force", src, check=False)
|
||||
|
||||
|
||||
def main(argv=None, repo=REPO, keys=None):
|
||||
ap = argparse.ArgumentParser(prog="bastion_agent.verify_release",
|
||||
description="Rebuild the published agent and compare it with the release manifest.")
|
||||
ap.add_argument("--base", default=os.environ.get("MANIFEST_BASE", DEFAULT_BASE),
|
||||
help="where version.json is served (default %s)" % DEFAULT_BASE)
|
||||
args = ap.parse_args(argv)
|
||||
base = args.base.rstrip("/")
|
||||
work = tempfile.mkdtemp(prefix="agent-verify-")
|
||||
try:
|
||||
body = fetch(base + "/version.json")
|
||||
sig = fetch(base + "/" + sig_name(body)).decode("ascii", "replace")
|
||||
m, key = check_signature(body, sig, keys)
|
||||
ver, built_from = m["version"], m.get("built_from") or ""
|
||||
print("published: v%s built from %s" % (ver, built_from or "(not named)"))
|
||||
print("signature: ok (release key %d of %d, %s)"
|
||||
% (key + 1, len(RELEASE_KEYS if keys is None else keys), "primary" if key == 0 else "spare"))
|
||||
files = m.get("files") if isinstance(m.get("files"), dict) else {}
|
||||
# Every .pyz the manifest lists: the artifact under its own name, and
|
||||
# the same bytes under the legacy name while that is still published.
|
||||
pyz = {n: (e or {}).get("sha256") for n, e in sorted(files.items()) if n.endswith(".pyz")}
|
||||
if "bastion-agent-%s.pyz" % ver not in pyz:
|
||||
raise VerifyError("the manifest lists no bastion-agent-%s.pyz" % ver)
|
||||
label, commit = pick_source(repo, ver, built_from)
|
||||
print("source: %s (%s)" % (label, commit[:12]))
|
||||
got = rebuild(repo, commit, ver, work)
|
||||
bad = 0
|
||||
for name, want in pyz.items():
|
||||
same = got == want
|
||||
bad += not same
|
||||
print("rebuild: %-28s %s" % (name, "matches %s" % got if same
|
||||
else "DIFFERS (rebuilt %s, published %s)" % (got, want)))
|
||||
return 1 if bad else 0
|
||||
except VerifyError as exc:
|
||||
print("verify: %s" % exc, file=sys.stderr)
|
||||
return 1
|
||||
except subprocess.CalledProcessError as exc:
|
||||
print("verify: git %s failed: %s" % (" ".join(exc.cmd[3:]), (exc.stderr or "").strip()), file=sys.stderr)
|
||||
return 1
|
||||
except OSError as exc: # network errors from urllib are OSErrors
|
||||
print("verify: %s" % exc, file=sys.stderr)
|
||||
return 1
|
||||
finally:
|
||||
shutil.rmtree(work, ignore_errors=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,476 @@
|
||||
"""bastion agent `walk` stage — the filesystem + metadata map.
|
||||
|
||||
The cheapest thing a host can plan from: every entry under the manifest's
|
||||
roots with path, kind, size, mtime, a sniffed MIME type and (for files within
|
||||
`max_file_bytes`) a SHA-256. Nothing is read beyond an 8 KiB sniff and the
|
||||
hash stream; nothing is written outside the output directory.
|
||||
|
||||
Portability is a requirement, not a nicety: this runs on the owner's Mac,
|
||||
Windows box or Linux server. So: stdlib only, no syntax newer than Python
|
||||
3.8, paths recorded `/`-separated relative to their root regardless of OS,
|
||||
`~` expanded per-OS, symlinks AND Windows junctions treated alike (recorded,
|
||||
never followed unless the manifest says so), and every entry sorted so two
|
||||
runs over the same tree produce byte-identical maps.
|
||||
|
||||
Output (docs/walk.md): `<out>/walk.jsonl` (one JSON object per
|
||||
entry, sorted by root index then path) and `<out>/walk.summary.json`. Both
|
||||
are written atomically. A previous walk.jsonl in the same directory is used
|
||||
as a hash cache: a file whose size and mtime_ns are unchanged keeps its
|
||||
hash without being re-read.
|
||||
"""
|
||||
|
||||
import datetime as _dt
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import re
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
from .manifest import AGENT_DIR, signing_digest
|
||||
|
||||
HASH_BUF = 1024 * 1024
|
||||
SNIFF_BYTES = 16
|
||||
TEXT_PROBE_BYTES = 8192
|
||||
|
||||
# (prefix, mime) — a small, deterministic magic table for the formats a
|
||||
# home corpus is made of. Anything else falls back to the extension, then
|
||||
# to a text/binary probe. Order matters only where prefixes overlap.
|
||||
MAGIC = (
|
||||
(b"%PDF-", "application/pdf"),
|
||||
(b"\x89PNG\r\n\x1a\n", "image/png"),
|
||||
(b"\xff\xd8\xff", "image/jpeg"),
|
||||
(b"GIF87a", "image/gif"),
|
||||
(b"GIF89a", "image/gif"),
|
||||
(b"RIFF", "application/x-riff"), # refined below: WEBP / WAV / AVI
|
||||
(b"PK\x03\x04", "application/zip"), # refined below: docx/xlsx/pptx by ext
|
||||
(b"\x1f\x8b", "application/gzip"),
|
||||
(b"BZh", "application/x-bzip2"),
|
||||
(b"\xfd7zXZ\x00", "application/x-xz"),
|
||||
(b"7z\xbc\xaf\x27\x1c", "application/x-7z-compressed"),
|
||||
(b"Rar!\x1a\x07", "application/vnd.rar"),
|
||||
(b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1", "application/x-ole-storage"),
|
||||
(b"SQLite format 3\x00", "application/vnd.sqlite3"),
|
||||
(b"\x7fELF", "application/x-executable"),
|
||||
(b"MZ", "application/x-dosexec"),
|
||||
(b"\xca\xfe\xba\xbe", "application/x-mach-binary"),
|
||||
(b"\xcf\xfa\xed\xfe", "application/x-mach-binary"),
|
||||
(b"ID3", "audio/mpeg"),
|
||||
(b"fLaC", "audio/flac"),
|
||||
(b"OggS", "audio/ogg"),
|
||||
(b"\x00\x00\x00\x1cftyp", "video/mp4"),
|
||||
(b"\x00\x00\x00\x18ftyp", "video/mp4"),
|
||||
(b"\x00\x00\x00\x20ftyp", "video/mp4"),
|
||||
(b"\x1aE\xdf\xa3", "video/x-matroska"),
|
||||
(b"{\\rtf", "application/rtf"),
|
||||
)
|
||||
ZIP_BY_EXT = {
|
||||
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
||||
".odt": "application/vnd.oasis.opendocument.text",
|
||||
".ods": "application/vnd.oasis.opendocument.spreadsheet",
|
||||
".odp": "application/vnd.oasis.opendocument.presentation",
|
||||
".epub": "application/epub+zip",
|
||||
".jar": "application/java-archive",
|
||||
".apk": "application/vnd.android.package-archive",
|
||||
}
|
||||
OLE_BY_EXT = {
|
||||
".doc": "application/msword",
|
||||
".xls": "application/vnd.ms-excel",
|
||||
".ppt": "application/vnd.ms-powerpoint",
|
||||
".msg": "application/vnd.ms-outlook",
|
||||
}
|
||||
# extension table kept here rather than the `mimetypes` module: that module
|
||||
# reads the OS registry / /etc/mime.types, so the same file would get a
|
||||
# different type on a different machine. Determinism wins.
|
||||
EXT_MIME = {
|
||||
".txt": "text/plain", ".md": "text/markdown", ".markdown": "text/markdown",
|
||||
".rst": "text/x-rst", ".csv": "text/csv", ".tsv": "text/tab-separated-values",
|
||||
".json": "application/json", ".jsonl": "application/x-ndjson",
|
||||
".yaml": "application/yaml", ".yml": "application/yaml", ".toml": "application/toml",
|
||||
".xml": "application/xml", ".html": "text/html", ".htm": "text/html",
|
||||
".css": "text/css", ".js": "text/javascript", ".mjs": "text/javascript",
|
||||
".ts": "text/x-typescript", ".tsx": "text/x-typescript", ".jsx": "text/javascript",
|
||||
".py": "text/x-python", ".go": "text/x-go", ".rs": "text/x-rust",
|
||||
".java": "text/x-java", ".c": "text/x-c", ".h": "text/x-c", ".cpp": "text/x-c++",
|
||||
".hpp": "text/x-c++", ".cs": "text/x-csharp", ".rb": "text/x-ruby",
|
||||
".php": "text/x-php", ".sh": "text/x-shellscript", ".bash": "text/x-shellscript",
|
||||
".ps1": "text/x-powershell", ".bat": "text/x-msdos-batch", ".sql": "application/sql",
|
||||
".ini": "text/plain", ".cfg": "text/plain", ".conf": "text/plain", ".log": "text/plain",
|
||||
".eml": "message/rfc822", ".ics": "text/calendar", ".vcf": "text/vcard",
|
||||
".svg": "image/svg+xml", ".bmp": "image/bmp", ".tif": "image/tiff", ".tiff": "image/tiff",
|
||||
".heic": "image/heic", ".webp": "image/webp",
|
||||
".mp3": "audio/mpeg", ".wav": "audio/wav", ".m4a": "audio/mp4", ".aac": "audio/aac",
|
||||
".mp4": "video/mp4", ".mov": "video/quicktime", ".avi": "video/x-msvideo",
|
||||
".mkv": "video/x-matroska", ".webm": "video/webm",
|
||||
".pdf": "application/pdf", ".zip": "application/zip", ".gz": "application/gzip",
|
||||
".tar": "application/x-tar", ".key": "application/x-pem-file", ".pem": "application/x-pem-file",
|
||||
}
|
||||
EXT_MIME.update(ZIP_BY_EXT)
|
||||
EXT_MIME.update(OLE_BY_EXT)
|
||||
|
||||
KIND_FILE, KIND_DIR, KIND_SYMLINK, KIND_OTHER = "file", "dir", "symlink", "other"
|
||||
SKIP_TOO_LARGE = "too_large"
|
||||
|
||||
_CASEFOLD = os.name == "nt"
|
||||
|
||||
|
||||
class WalkError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
# ── exclude globs ───────────────────────────────────────────────────────────
|
||||
|
||||
def _glob_to_regex(pat):
|
||||
"""`**` spans segments, `*`/`?` stay inside one, `[...]` passes through."""
|
||||
i, n, out = 0, len(pat), []
|
||||
while i < n:
|
||||
c = pat[i]
|
||||
if c == "*":
|
||||
if pat[i:i + 3] == "**/":
|
||||
out.append("(?:.*/)?")
|
||||
i += 3
|
||||
continue
|
||||
if pat[i:i + 2] == "**":
|
||||
out.append(".*")
|
||||
i += 2
|
||||
continue
|
||||
out.append("[^/]*")
|
||||
elif c == "?":
|
||||
out.append("[^/]")
|
||||
elif c == "[":
|
||||
j = pat.find("]", i + 1)
|
||||
if j == -1:
|
||||
out.append(re.escape(c))
|
||||
else:
|
||||
body = pat[i + 1:j]
|
||||
if body.startswith("!"):
|
||||
body = "^" + body[1:]
|
||||
out.append("[" + body.replace("\\", "\\\\") + "]")
|
||||
i = j
|
||||
else:
|
||||
out.append(re.escape(c))
|
||||
i += 1
|
||||
return "".join(out)
|
||||
|
||||
|
||||
class Excluder:
|
||||
"""Compiled `scope.exclude`. A pattern with a `/` matches the whole
|
||||
root-relative path; one without also matches any single segment (so
|
||||
`node_modules` or `*.key` prune at any depth, gitignore-style)."""
|
||||
|
||||
def __init__(self, patterns):
|
||||
flags = re.IGNORECASE if _CASEFOLD else 0
|
||||
self._path = []
|
||||
self._segment = []
|
||||
for p in patterns:
|
||||
p = p.replace("\\", "/").rstrip("/")
|
||||
rx = re.compile("^" + _glob_to_regex(p) + "$", flags)
|
||||
self._path.append(rx)
|
||||
if "/" not in p:
|
||||
self._segment.append(rx)
|
||||
elif p.endswith("/**"):
|
||||
# `docs/**` and `**/.git/**` prune the directory itself, not
|
||||
# just what is under it — otherwise the dir entry leaks into
|
||||
# the map while its contents are excluded
|
||||
self._path.append(re.compile("^" + _glob_to_regex(p[:-3]) + "$", flags))
|
||||
|
||||
def matches(self, relpath):
|
||||
if not relpath:
|
||||
return False
|
||||
for rx in self._path:
|
||||
if rx.match(relpath):
|
||||
return True
|
||||
if self._segment:
|
||||
for seg in relpath.split("/"):
|
||||
for rx in self._segment:
|
||||
if rx.match(seg):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# ── helpers ─────────────────────────────────────────────────────────────────
|
||||
|
||||
def _utc_iso(ns):
|
||||
secs, rem = divmod(ns, 1_000_000_000)
|
||||
dt = _dt.datetime.fromtimestamp(secs, _dt.timezone.utc)
|
||||
return dt.strftime("%Y-%m-%dT%H:%M:%S") + ".%06dZ" % (rem // 1000)
|
||||
|
||||
|
||||
def _relposix(root, path):
|
||||
rel = os.path.relpath(path, root)
|
||||
return "" if rel == os.curdir else rel.replace(os.sep, "/")
|
||||
|
||||
|
||||
def _is_link(entry):
|
||||
"""Symlink on every OS; on Windows also a junction (3.12+ exposes
|
||||
is_junction; older Pythons see junctions as directories, in which case
|
||||
the reparse tag on the lstat result is the tell)."""
|
||||
try:
|
||||
if entry.is_symlink():
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
is_junction = getattr(entry, "is_junction", None)
|
||||
if is_junction is not None:
|
||||
try:
|
||||
return bool(is_junction())
|
||||
except OSError:
|
||||
return False
|
||||
if os.name == "nt":
|
||||
try:
|
||||
st = entry.stat(follow_symlinks=False)
|
||||
tag = getattr(st, "st_reparse_tag", 0)
|
||||
return tag == 0xA0000003 # IO_REPARSE_TAG_MOUNT_POINT
|
||||
except OSError:
|
||||
return False
|
||||
return False
|
||||
|
||||
|
||||
def _sniff(path, ext):
|
||||
"""MIME by magic first, extension second, text/binary probe last."""
|
||||
try:
|
||||
with open(path, "rb") as fh:
|
||||
head = fh.read(TEXT_PROBE_BYTES)
|
||||
except OSError as exc:
|
||||
return None, "%s: %s" % (type(exc).__name__, exc.strerror or exc)
|
||||
for prefix, mime in MAGIC:
|
||||
if head.startswith(prefix):
|
||||
if mime == "application/zip":
|
||||
return ZIP_BY_EXT.get(ext, mime), None
|
||||
if mime == "application/x-ole-storage":
|
||||
return OLE_BY_EXT.get(ext, mime), None
|
||||
if mime == "application/x-riff":
|
||||
sub = head[8:12]
|
||||
return {b"WEBP": "image/webp", b"WAVE": "audio/wav",
|
||||
b"AVI ": "video/x-msvideo"}.get(sub, mime), None
|
||||
return mime, None
|
||||
if ext in EXT_MIME:
|
||||
return EXT_MIME[ext], None
|
||||
if not head:
|
||||
return "inode/x-empty", None
|
||||
if b"\x00" in head:
|
||||
return "application/octet-stream", None
|
||||
return "text/plain", None
|
||||
|
||||
|
||||
def _sha256(path):
|
||||
h = hashlib.sha256()
|
||||
with open(path, "rb") as fh:
|
||||
while True:
|
||||
buf = fh.read(HASH_BUF)
|
||||
if not buf:
|
||||
break
|
||||
h.update(buf)
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def load_hash_cache(jsonl_path):
|
||||
"""(root, path) -> (size, mtime_ns, sha256) from a previous walk."""
|
||||
cache = {}
|
||||
try:
|
||||
with open(jsonl_path, encoding="utf-8") as fh:
|
||||
for line in fh:
|
||||
try:
|
||||
rec = json.loads(line)
|
||||
except ValueError:
|
||||
continue
|
||||
if rec.get("kind") == KIND_FILE and rec.get("sha256"):
|
||||
cache[(rec["root"], rec["path"])] = (
|
||||
rec.get("size"), rec.get("mtime_ns"), rec["sha256"])
|
||||
except OSError:
|
||||
pass
|
||||
return cache
|
||||
|
||||
|
||||
def _atomic_write(path, write_fn):
|
||||
d = os.path.dirname(path) or "."
|
||||
fd, tmp = tempfile.mkstemp(prefix=".tmp-", dir=d)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8", newline="\n") as fh:
|
||||
write_fn(fh)
|
||||
os.replace(tmp, path)
|
||||
except BaseException:
|
||||
try:
|
||||
os.unlink(tmp)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
|
||||
# ── the walk ────────────────────────────────────────────────────────────────
|
||||
|
||||
class Walker:
|
||||
def __init__(self, manifest, hash_cache=None, do_hash=True):
|
||||
self.m = manifest
|
||||
# .agent/ is always excluded: it is the agent's own output, and indexing
|
||||
# it would feed the catalog back into itself
|
||||
self.excluder = Excluder(list(manifest.scope.exclude) + [AGENT_DIR])
|
||||
# …and on every OS regardless of case: a pre-existing ".Agent" on a
|
||||
# case-insensitive macOS volume IS the agent dir
|
||||
self.excluder._segment.append(re.compile("^" + re.escape(AGENT_DIR) + "$", re.IGNORECASE))
|
||||
self.cache = hash_cache or {}
|
||||
self.do_hash = do_hash
|
||||
self.entries = []
|
||||
self.counts = {k: 0 for k in (
|
||||
"files", "dirs", "symlinks", "other", "excluded", "errors",
|
||||
"too_large", "hashed", "hash_reused")}
|
||||
self.bytes = 0
|
||||
|
||||
# records ---------------------------------------------------------------
|
||||
def _rec(self, root_idx, rel, kind, **extra):
|
||||
rec = {"root": root_idx, "path": rel, "kind": kind}
|
||||
rec.update(extra)
|
||||
self.entries.append(rec)
|
||||
return rec
|
||||
|
||||
def _error(self, root_idx, rel, exc, kind=KIND_OTHER):
|
||||
self.counts["errors"] += 1
|
||||
msg = "%s: %s" % (type(exc).__name__, getattr(exc, "strerror", None) or exc)
|
||||
return self._rec(root_idx, rel, kind, error=msg)
|
||||
|
||||
# per-entry --------------------------------------------------------------
|
||||
def _file(self, root_idx, rel, entry_path, st):
|
||||
size = st.st_size
|
||||
mtime_ns = st.st_mtime_ns
|
||||
ext = os.path.splitext(rel)[1].lower()
|
||||
mime, err = _sniff(entry_path, ext)
|
||||
rec = {"size": size, "mtime": _utc_iso(mtime_ns), "mtime_ns": mtime_ns,
|
||||
"mime": mime, "sha256": None}
|
||||
if err is not None:
|
||||
self.counts["errors"] += 1
|
||||
rec["error"] = err
|
||||
self._rec(root_idx, rel, KIND_FILE, **rec)
|
||||
return
|
||||
self.counts["files"] += 1
|
||||
self.bytes += size
|
||||
if size > self.m.scope.max_file_bytes:
|
||||
self.counts["too_large"] += 1
|
||||
rec["skip"] = SKIP_TOO_LARGE
|
||||
elif self.do_hash:
|
||||
cached = self.cache.get((root_idx, rel))
|
||||
if cached and cached[0] == size and cached[1] == mtime_ns:
|
||||
rec["sha256"] = cached[2]
|
||||
self.counts["hash_reused"] += 1
|
||||
else:
|
||||
try:
|
||||
rec["sha256"] = _sha256(entry_path)
|
||||
self.counts["hashed"] += 1
|
||||
except OSError as exc:
|
||||
self.counts["errors"] += 1
|
||||
rec["error"] = "%s: %s" % (type(exc).__name__, exc.strerror or exc)
|
||||
self._rec(root_idx, rel, KIND_FILE, **rec)
|
||||
|
||||
def _link(self, root_idx, rel, entry_path):
|
||||
self.counts["symlinks"] += 1
|
||||
try:
|
||||
target = os.readlink(entry_path)
|
||||
except OSError as exc:
|
||||
target = None
|
||||
err = "%s: %s" % (type(exc).__name__, exc.strerror or exc)
|
||||
self._rec(root_idx, rel, KIND_SYMLINK, target=None, error=err)
|
||||
self.counts["errors"] += 1
|
||||
return
|
||||
if isinstance(target, bytes):
|
||||
target = os.fsdecode(target)
|
||||
self._rec(root_idx, rel, KIND_SYMLINK, target=target.replace(os.sep, "/"))
|
||||
|
||||
# recursion --------------------------------------------------------------
|
||||
def _walk_dir(self, root_idx, root, dirpath):
|
||||
try:
|
||||
it = os.scandir(dirpath)
|
||||
except OSError as exc:
|
||||
self._error(root_idx, _relposix(root, dirpath), exc, KIND_DIR)
|
||||
return
|
||||
with it:
|
||||
try:
|
||||
entries = list(it)
|
||||
except OSError as exc:
|
||||
self._error(root_idx, _relposix(root, dirpath), exc, KIND_DIR)
|
||||
return
|
||||
for entry in entries:
|
||||
rel = _relposix(root, entry.path)
|
||||
if self.excluder.matches(rel):
|
||||
self.counts["excluded"] += 1
|
||||
continue
|
||||
follow = self.m.scope.follow_symlinks
|
||||
if _is_link(entry) and not follow:
|
||||
self._link(root_idx, rel, entry.path)
|
||||
continue
|
||||
try:
|
||||
is_dir = entry.is_dir(follow_symlinks=follow)
|
||||
except OSError as exc:
|
||||
self._error(root_idx, rel, exc)
|
||||
continue
|
||||
if is_dir:
|
||||
self.counts["dirs"] += 1
|
||||
try:
|
||||
st = entry.stat(follow_symlinks=follow)
|
||||
self._rec(root_idx, rel, KIND_DIR,
|
||||
mtime=_utc_iso(st.st_mtime_ns), mtime_ns=st.st_mtime_ns)
|
||||
except OSError as exc:
|
||||
self._error(root_idx, rel, exc, KIND_DIR)
|
||||
continue
|
||||
self._walk_dir(root_idx, root, entry.path)
|
||||
continue
|
||||
try:
|
||||
is_file = entry.is_file(follow_symlinks=follow)
|
||||
st = entry.stat(follow_symlinks=follow)
|
||||
except OSError as exc:
|
||||
self._error(root_idx, rel, exc)
|
||||
continue
|
||||
if is_file:
|
||||
self._file(root_idx, rel, entry.path, st)
|
||||
else:
|
||||
self.counts["other"] += 1 # socket, fifo, device
|
||||
self._rec(root_idx, rel, KIND_OTHER)
|
||||
|
||||
def run(self):
|
||||
roots = []
|
||||
for idx, raw in enumerate(self.m.scope.roots):
|
||||
resolved = os.path.abspath(os.path.expanduser(raw))
|
||||
roots.append({"index": idx, "path": raw, "resolved": resolved})
|
||||
if not os.path.isdir(resolved):
|
||||
self.counts["errors"] += 1
|
||||
roots[-1]["error"] = "not a directory or not present on this machine"
|
||||
continue
|
||||
self._walk_dir(idx, resolved, resolved)
|
||||
self.entries.sort(key=lambda r: (r["root"], r["path"]))
|
||||
return roots
|
||||
|
||||
|
||||
def walk(manifest, out_dir, do_hash=True, now=None):
|
||||
"""Run the walk stage for a validated manifest; write walk.jsonl and
|
||||
walk.summary.json into out_dir. Returns the summary dict."""
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
jsonl = os.path.join(out_dir, "walk.jsonl")
|
||||
summary_path = os.path.join(out_dir, "walk.summary.json")
|
||||
cache = load_hash_cache(jsonl) if do_hash else {}
|
||||
w = Walker(manifest, hash_cache=cache, do_hash=do_hash)
|
||||
roots = w.run()
|
||||
|
||||
def _write_jsonl(fh):
|
||||
for rec in w.entries:
|
||||
fh.write(json.dumps(rec, ensure_ascii=False, sort_keys=True))
|
||||
fh.write("\n")
|
||||
|
||||
_atomic_write(jsonl, _write_jsonl)
|
||||
walked_at = now or _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
summary = {
|
||||
"freeagent": 1,
|
||||
"stage": "walk",
|
||||
"manifest_id": manifest.id,
|
||||
"manifest_digest": "sha256:" + signing_digest(manifest.raw),
|
||||
"walked_at": walked_at,
|
||||
"platform": {"system": platform.system(), "python": platform.python_version()},
|
||||
"roots": roots,
|
||||
"entries": len(w.entries),
|
||||
"bytes": w.bytes,
|
||||
"counts": w.counts,
|
||||
"hashing": bool(do_hash),
|
||||
}
|
||||
_atomic_write(summary_path, lambda fh: fh.write(
|
||||
json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True) + "\n"))
|
||||
return summary
|
||||
@@ -0,0 +1,423 @@
|
||||
"""The manifest gate's certificate and signature checks, in the standard library.
|
||||
|
||||
The gate used to shell out to `openssl` for three things: read the signing
|
||||
cert's CN, check that the cert chains to FREEAGENT_CA_ROOT, and verify the
|
||||
ECDSA P-256 / SHA-256 signature over the manifest payload. Windows ships no
|
||||
`openssl`, so on a Windows machine the agent refused every manifest and never
|
||||
ran — its bastion reported `missing-openssl` and the user's context store was
|
||||
never indexed.
|
||||
|
||||
This module does exactly those three checks with nothing but the standard
|
||||
library, so the gate runs wherever Python does:
|
||||
|
||||
* a minimal DER reader for the X.509 fields the gate needs (subject and
|
||||
issuer names, validity, the EC public key, basicConstraints, and the
|
||||
signed TBS bytes with their signature);
|
||||
* ECDSA verification over NIST P-256 (verification only — no secret is ever
|
||||
held here, so constant-time arithmetic is not a concern).
|
||||
|
||||
Only what the managed agent's certificates use is accepted: EC P-256 keys and
|
||||
ecdsa-with-SHA256 signatures. Anything else is refused by name, never guessed.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import datetime as _dt
|
||||
import hashlib
|
||||
import re
|
||||
|
||||
# ── P-256 (FIPS 186-4, D.1.2.3) ─────────────────────────────────────────────
|
||||
_P = 0xffffffff00000001000000000000000000000000ffffffffffffffffffffffff
|
||||
_A = _P - 3
|
||||
_B = 0x5ac635d8aa3a93e7b3ebbd55769886bc651d06b0cc53b0f63bce3c3e27d2604b
|
||||
_N = 0xffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551
|
||||
_G = (0x6b17d1f2e12c4247f8bce6e563a440f277037d812deb33a0f4a13945d898c296,
|
||||
0x4fe342e2fe1a7f9b8ee7eb4a7c0f9e162bce33576b315ececbb6406837bf51f5)
|
||||
|
||||
# OIDs, as their DER content bytes.
|
||||
_OID_CN = bytes.fromhex("550403") # 2.5.4.3
|
||||
_OID_EC_PUBKEY = bytes.fromhex("2a8648ce3d0201") # 1.2.840.10045.2.1
|
||||
_OID_P256 = bytes.fromhex("2a8648ce3d030107") # 1.2.840.10045.3.1.7
|
||||
_OID_ECDSA_SHA256 = bytes.fromhex("2a8648ce3d040302") # 1.2.840.10045.4.3.2
|
||||
_OID_BASIC_CONSTRAINTS = bytes.fromhex("551d13") # 2.5.29.19
|
||||
_OID_KEY_USAGE = bytes.fromhex("551d0f") # 2.5.29.15
|
||||
# Critical extensions the gate may meet and still decide correctly: those it
|
||||
# reads (above) and the ones Go's x509 writes on the managed install's cert.
|
||||
_UNDERSTOOD_EXTENSIONS = frozenset({
|
||||
_OID_BASIC_CONSTRAINTS, _OID_KEY_USAGE,
|
||||
bytes.fromhex("551d25"), # extKeyUsage
|
||||
bytes.fromhex("551d11"), # subjectAltName
|
||||
bytes.fromhex("551d0e"), # subjectKeyIdentifier
|
||||
bytes.fromhex("551d23"), # authorityKeyIdentifier
|
||||
})
|
||||
_KU_KEY_CERT_SIGN = 0x0400 # keyUsage bit 5, as a big-endian 16-bit field
|
||||
|
||||
|
||||
# The directory-string types a name value may carry, and how each decodes.
|
||||
_NAME_STRING_CODECS = {0x0c: "utf-8", 0x13: "ascii", 0x16: "ascii",
|
||||
0x14: "latin-1", 0x1e: "utf-16-be"}
|
||||
|
||||
|
||||
class CertError(Exception):
|
||||
"""A certificate or signature the gate cannot accept; the message says why."""
|
||||
|
||||
|
||||
# ── DER ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _tlv(buf, off):
|
||||
"""(tag, content_start, content_end, next_off) of the element at `off`."""
|
||||
if off + 2 > len(buf):
|
||||
raise CertError("truncated DER")
|
||||
tag = buf[off]
|
||||
ln = buf[off + 1]
|
||||
off += 2
|
||||
if ln & 0x80:
|
||||
n = ln & 0x7f
|
||||
if n == 0 or n > 4 or off + n > len(buf):
|
||||
raise CertError("unsupported DER length")
|
||||
ln = int.from_bytes(buf[off:off + n], "big")
|
||||
# DER: the long form only for lengths >= 128, with no leading zero
|
||||
# byte — the one encoding a length has, as openssl requires.
|
||||
if ln < 0x80 or buf[off] == 0:
|
||||
raise CertError("non-minimal DER length")
|
||||
off += n
|
||||
end = off + ln
|
||||
if end > len(buf):
|
||||
raise CertError("truncated DER")
|
||||
return tag, off, end, end
|
||||
|
||||
|
||||
def _children(buf, start, end):
|
||||
out = []
|
||||
off = start
|
||||
while off < end:
|
||||
tag, cs, ce, off = _tlv(buf, off)
|
||||
out.append((tag, cs, ce))
|
||||
return out
|
||||
|
||||
|
||||
def _expect(el, tag, what):
|
||||
if el[0] != tag:
|
||||
raise CertError("%s: unexpected DER tag 0x%02x" % (what, el[0]))
|
||||
return el
|
||||
|
||||
|
||||
def _name_parts(buf, start, end):
|
||||
"""[(oid_bytes, value_str)] of an X.501 Name (SEQUENCE OF SET OF ATV)."""
|
||||
parts = []
|
||||
for rdn in _children(buf, start, end):
|
||||
_expect(rdn, 0x31, "name RDN")
|
||||
for atv in _children(buf, rdn[1], rdn[2]):
|
||||
_expect(atv, 0x30, "name attribute")
|
||||
kids = _children(buf, atv[1], atv[2])
|
||||
if len(kids) != 2 or kids[0][0] != 0x06:
|
||||
raise CertError("malformed name attribute")
|
||||
raw = bytes(buf[kids[1][1]:kids[1][2]])
|
||||
codec = _NAME_STRING_CODECS.get(kids[1][0])
|
||||
if codec is None:
|
||||
raise CertError("name attribute has string type 0x%02x" % kids[1][0])
|
||||
try:
|
||||
val = raw.decode(codec)
|
||||
except UnicodeDecodeError:
|
||||
raise CertError("name attribute is not text")
|
||||
parts.append((bytes(buf[kids[0][1]:kids[0][2]]), val))
|
||||
return parts
|
||||
|
||||
|
||||
def _time(buf, el):
|
||||
raw = buf[el[1]:el[2]].decode("ascii", "replace")
|
||||
if el[0] == 0x17: # UTCTime YYMMDDHHMMSSZ
|
||||
m = re.fullmatch(r"(\d{2})(\d{2})(\d{2})(\d{2})(\d{2})(\d{2})Z", raw)
|
||||
if not m:
|
||||
raise CertError("unsupported UTCTime %r" % raw)
|
||||
y = int(m.group(1))
|
||||
y += 2000 if y < 50 else 1900
|
||||
rest = [int(g) for g in m.groups()[1:]]
|
||||
elif el[0] == 0x18: # GeneralizedTime YYYYMMDDHHMMSSZ
|
||||
m = re.fullmatch(r"(\d{4})(\d{2})(\d{2})(\d{2})(\d{2})(\d{2})Z", raw)
|
||||
if not m:
|
||||
raise CertError("unsupported GeneralizedTime %r" % raw)
|
||||
y = int(m.group(1))
|
||||
rest = [int(g) for g in m.groups()[1:]]
|
||||
else:
|
||||
raise CertError("validity time has tag 0x%02x" % el[0])
|
||||
return _dt.datetime(y, *rest, tzinfo=_dt.timezone.utc)
|
||||
|
||||
|
||||
class Cert:
|
||||
"""The fields of one X.509 certificate the gate reads."""
|
||||
|
||||
def __init__(self, der):
|
||||
self.der = bytes(der)
|
||||
buf = self.der
|
||||
top = _tlv(buf, 0)
|
||||
if top[0] != 0x30 or top[3] != len(buf):
|
||||
raise CertError("not a single DER certificate")
|
||||
parts = _children(buf, top[1], top[2])
|
||||
if len(parts) != 3:
|
||||
raise CertError("a certificate is exactly tbsCertificate, algorithm, signature")
|
||||
tbs, sig_alg, sig_val = parts
|
||||
_expect(tbs, 0x30, "tbsCertificate")
|
||||
# The signed bytes are the WHOLE tbsCertificate element, header included.
|
||||
tbs_hdr_start = top[1]
|
||||
self.tbs_bytes = buf[tbs_hdr_start:tbs[2]]
|
||||
self.sig_alg = self._alg_oid(buf, sig_alg)
|
||||
_expect(sig_val, 0x03, "signatureValue")
|
||||
if sig_val[2] - sig_val[1] < 2 or buf[sig_val[1]] != 0:
|
||||
raise CertError("signatureValue is empty or has unused bits")
|
||||
self.signature = buf[sig_val[1] + 1:sig_val[2]]
|
||||
|
||||
f = _children(buf, tbs[1], tbs[2])
|
||||
i = 0
|
||||
if f and f[0][0] == 0xa0: # [0] EXPLICIT version
|
||||
i = 1
|
||||
# serial, signature alg, issuer, validity, subject, spki, then optional
|
||||
# [1]/[2] unique ids and [3] extensions.
|
||||
if len(f) < i + 6:
|
||||
raise CertError("tbsCertificate is missing fields")
|
||||
self.tbs_sig_alg = self._alg_oid(buf, f[i + 1])
|
||||
issuer, validity, subject, spki = f[i + 2], f[i + 3], f[i + 4], f[i + 5]
|
||||
_expect(issuer, 0x30, "issuer")
|
||||
_expect(subject, 0x30, "subject")
|
||||
self.issuer_der = buf[issuer[1]:issuer[2]]
|
||||
self.subject_der = buf[subject[1]:subject[2]]
|
||||
self.subject = _name_parts(buf, subject[1], subject[2])
|
||||
v = _children(buf, validity[1], validity[2])
|
||||
if len(v) != 2:
|
||||
raise CertError("malformed validity")
|
||||
self.not_before, self.not_after = _time(buf, v[0]), _time(buf, v[1])
|
||||
self.public_key = self._ec_point(buf, spki)
|
||||
self.is_ca = False
|
||||
self.key_usage = None # None = no keyUsage extension
|
||||
seen_ext = set()
|
||||
for ext_wrap in f[i + 6:]:
|
||||
if ext_wrap[0] != 0xa3:
|
||||
continue
|
||||
exts = _children(buf, ext_wrap[1], ext_wrap[2])
|
||||
if len(exts) != 1:
|
||||
raise CertError("malformed extensions")
|
||||
for ext in _children(buf, exts[0][1], exts[0][2]):
|
||||
k = _children(buf, ext[1], ext[2])
|
||||
if len(k) not in (2, 3) or k[0][0] != 0x06 or k[-1][0] != 0x04:
|
||||
raise CertError("malformed extension")
|
||||
oid = bytes(buf[k[0][1]:k[0][2]])
|
||||
if oid in seen_ext:
|
||||
raise CertError("duplicate extension")
|
||||
seen_ext.add(oid)
|
||||
critical = len(k) == 3 and k[1][0] == 0x01 and buf[k[1][1]:k[1][2]] != b"\x00"
|
||||
octets = k[-1]
|
||||
if oid == _OID_BASIC_CONSTRAINTS:
|
||||
bc = _tlv(buf, octets[1])
|
||||
inner = _children(buf, bc[1], bc[2])
|
||||
self.is_ca = bool(inner and inner[0][0] == 0x01
|
||||
and buf[inner[0][1]:inner[0][2]] != b"\x00")
|
||||
elif oid == _OID_KEY_USAGE:
|
||||
ku = _tlv(buf, octets[1])
|
||||
bits_ = bytes(buf[ku[1]:ku[2]])
|
||||
self.key_usage = int.from_bytes(bits_[1:3].ljust(2, b"\x00"), "big") if bits_ else 0
|
||||
elif critical and oid not in _UNDERSTOOD_EXTENSIONS:
|
||||
# openssl refuses a certificate carrying a critical
|
||||
# extension it does not understand; so does the gate.
|
||||
raise CertError("certificate has an unknown critical extension")
|
||||
|
||||
@staticmethod
|
||||
def _alg_oid(buf, el):
|
||||
_expect(el, 0x30, "algorithm")
|
||||
k = _children(buf, el[1], el[2])
|
||||
if not k or k[0][0] != 0x06:
|
||||
raise CertError("malformed algorithm identifier")
|
||||
return bytes(buf[k[0][1]:k[0][2]])
|
||||
|
||||
@staticmethod
|
||||
def _ec_point(buf, spki):
|
||||
_expect(spki, 0x30, "subjectPublicKeyInfo")
|
||||
alg, bits = (_children(buf, spki[1], spki[2]) + [None, None])[:2]
|
||||
if bits is None:
|
||||
raise CertError("malformed subjectPublicKeyInfo")
|
||||
a = _children(buf, alg[1], alg[2])
|
||||
if (len(a) != 2 or buf[a[0][1]:a[0][2]] != _OID_EC_PUBKEY
|
||||
or buf[a[1][1]:a[1][2]] != _OID_P256):
|
||||
raise CertError("signing key is not an EC P-256 key — the gate verifies "
|
||||
"ecdsa-p256-sha256 only")
|
||||
_expect(bits, 0x03, "public key")
|
||||
if bits[2] - bits[1] < 1 or buf[bits[1]] != 0:
|
||||
raise CertError("public key BITSTRING is empty or has unused bits")
|
||||
pt = buf[bits[1] + 1:bits[2]]
|
||||
if len(pt) != 65 or pt[0] != 0x04:
|
||||
raise CertError("EC public key is not an uncompressed P-256 point")
|
||||
q = (int.from_bytes(pt[1:33], "big"), int.from_bytes(pt[33:], "big"))
|
||||
if not _on_curve(q):
|
||||
raise CertError("EC public key is not on P-256")
|
||||
return q
|
||||
|
||||
def cn(self):
|
||||
cns = [v for oid, v in self.subject if oid == _OID_CN]
|
||||
if len(cns) != 1:
|
||||
# exactly one CN: a second one is how a cert could carry the
|
||||
# owner's name next to its real one
|
||||
raise CertError("signing cert must carry exactly one CN, has %d" % len(cns))
|
||||
return cns[0]
|
||||
|
||||
def valid_at(self, when):
|
||||
return self.not_before <= when <= self.not_after
|
||||
|
||||
|
||||
def load_pem_certs(text):
|
||||
"""Every CERTIFICATE block in a PEM text, as Cert objects (in order)."""
|
||||
if isinstance(text, bytes):
|
||||
text = text.decode("ascii", "replace")
|
||||
blocks = re.findall(r"-----BEGIN CERTIFICATE-----(.*?)-----END CERTIFICATE-----", text, re.S)
|
||||
if not blocks:
|
||||
raise CertError("no PEM certificate found")
|
||||
out = []
|
||||
for b in blocks:
|
||||
try:
|
||||
out.append(Cert(base64.b64decode("".join(b.split()), validate=True)))
|
||||
except CertError:
|
||||
raise
|
||||
except (ValueError, TypeError, IndexError, OverflowError) as exc:
|
||||
raise CertError("unreadable PEM certificate: %s" % exc)
|
||||
return out
|
||||
|
||||
|
||||
def load_root_file(text):
|
||||
"""The trust anchors in a CA file: every certificate the gate can use.
|
||||
One the gate cannot use (another key type, say) is skipped, not fatal —
|
||||
openssl would skip it too; an empty result is refused."""
|
||||
if isinstance(text, bytes):
|
||||
text = text.decode("ascii", "replace")
|
||||
blocks = re.findall(r"-----BEGIN CERTIFICATE-----(.*?)-----END CERTIFICATE-----", text, re.S)
|
||||
roots = []
|
||||
for b in blocks:
|
||||
try:
|
||||
roots.append(Cert(base64.b64decode("".join(b.split()), validate=True)))
|
||||
except (CertError, ValueError, TypeError, IndexError, OverflowError):
|
||||
continue
|
||||
if not roots:
|
||||
raise CertError("FREEAGENT_CA_ROOT holds no usable P-256 certificate")
|
||||
return roots
|
||||
|
||||
|
||||
def load_cert_file(path):
|
||||
with open(path, "rb") as fh:
|
||||
certs = load_pem_certs(fh.read())
|
||||
return certs[0]
|
||||
|
||||
|
||||
# ── P-256 arithmetic (affine; verification only) ────────────────────────────
|
||||
|
||||
def _on_curve(pt):
|
||||
x, y = pt
|
||||
if not (0 <= x < _P and 0 <= y < _P):
|
||||
return False
|
||||
return (y * y - (x * x * x + _A * x + _B)) % _P == 0
|
||||
|
||||
|
||||
def _add(p1, p2):
|
||||
if p1 is None:
|
||||
return p2
|
||||
if p2 is None:
|
||||
return p1
|
||||
x1, y1 = p1
|
||||
x2, y2 = p2
|
||||
if x1 == x2:
|
||||
if (y1 + y2) % _P == 0:
|
||||
return None
|
||||
lam = (3 * x1 * x1 + _A) * pow(2 * y1, -1, _P) % _P
|
||||
else:
|
||||
lam = (y2 - y1) * pow(x2 - x1, -1, _P) % _P
|
||||
x3 = (lam * lam - x1 - x2) % _P
|
||||
return (x3, (lam * (x1 - x3) - y1) % _P)
|
||||
|
||||
|
||||
def _mul(k, pt):
|
||||
acc = None
|
||||
while k:
|
||||
if k & 1:
|
||||
acc = _add(acc, pt)
|
||||
pt = _add(pt, pt)
|
||||
k >>= 1
|
||||
return acc
|
||||
|
||||
|
||||
def _der_sig(sig):
|
||||
top = _tlv(sig, 0)
|
||||
if top[0] != 0x30 or top[3] != len(sig):
|
||||
raise CertError("signature is not a DER ECDSA signature")
|
||||
k = _children(sig, top[1], top[2])
|
||||
if len(k) != 2 or k[0][0] != 0x02 or k[1][0] != 0x02:
|
||||
raise CertError("signature is not a DER ECDSA signature")
|
||||
vals = []
|
||||
for _, cs, ce in k:
|
||||
raw = sig[cs:ce]
|
||||
# DER INTEGER: non-empty, positive here, and minimal — no leading
|
||||
# zero byte unless the next byte's top bit needs it. openssl refuses
|
||||
# a padded signature; so does the gate (one encoding per signature).
|
||||
if not raw or raw[0] & 0x80 or (len(raw) > 1 and raw[0] == 0 and not raw[1] & 0x80):
|
||||
raise CertError("signature integers are not minimal DER")
|
||||
vals.append(int.from_bytes(raw, "big"))
|
||||
return vals[0], vals[1]
|
||||
|
||||
|
||||
def ecdsa_p256_sha256_verify(pub, message, der_signature):
|
||||
"""True iff `der_signature` is a valid ECDSA P-256 / SHA-256 signature by
|
||||
`pub` over `message`. Malformed input is simply False."""
|
||||
try:
|
||||
r, s = _der_sig(bytes(der_signature))
|
||||
except CertError:
|
||||
return False
|
||||
if not (1 <= r < _N and 1 <= s < _N) or not _on_curve(pub):
|
||||
return False
|
||||
e = int.from_bytes(hashlib.sha256(message).digest(), "big")
|
||||
w = pow(s, -1, _N)
|
||||
pt = _add(_mul(e * w % _N, _G), _mul(r * w % _N, pub))
|
||||
return pt is not None and pt[0] % _N == r
|
||||
|
||||
|
||||
# ── the chain ───────────────────────────────────────────────────────────────
|
||||
|
||||
def _signed_by(child, issuer):
|
||||
if child.sig_alg != _OID_ECDSA_SHA256 or child.tbs_sig_alg != _OID_ECDSA_SHA256:
|
||||
raise CertError("certificate is not signed with ecdsa-with-SHA256")
|
||||
return ecdsa_p256_sha256_verify(issuer.public_key, child.tbs_bytes, child.signature)
|
||||
|
||||
|
||||
def verify_chain(leaf, roots, now=None):
|
||||
"""Raise CertError unless `leaf` is one of `roots` or is signed directly by
|
||||
one of them. For the managed install — whose cert IS its root — this is
|
||||
the answer `openssl verify -CAfile root leaf` gave. For a separate issuer
|
||||
it enforces validity, CA basicConstraints and keyUsage keyCertSign, but it
|
||||
treats every cert in the file as a trust anchor and does not build longer
|
||||
chains or check pathLen / AKID: a one-level PKI only."""
|
||||
now = now or _dt.datetime.now(_dt.timezone.utc)
|
||||
if not leaf.valid_at(now):
|
||||
raise CertError("signing cert is outside its validity window (%s — %s)"
|
||||
% (leaf.not_before.isoformat(), leaf.not_after.isoformat()))
|
||||
for root in roots:
|
||||
if root.der == leaf.der:
|
||||
# The managed install's cert is its own root. A trust anchor is
|
||||
# trusted as given; its self-signature is still required to hold.
|
||||
if not _signed_by(leaf, leaf):
|
||||
raise CertError("self-signed root does not verify")
|
||||
return root
|
||||
why = "signing cert does not chain to FREEAGENT_CA_ROOT"
|
||||
for root in roots:
|
||||
if root.subject_der != leaf.issuer_der:
|
||||
continue
|
||||
# Every same-named root is tried: a name is not an identity.
|
||||
if not root.valid_at(now):
|
||||
why = "the CA root is outside its validity window"
|
||||
continue
|
||||
if not root.is_ca:
|
||||
why = "the issuing certificate is not a CA"
|
||||
continue
|
||||
if root.key_usage is not None and not root.key_usage & _KU_KEY_CERT_SIGN:
|
||||
why = "the issuing certificate may not sign certificates (keyUsage)"
|
||||
continue
|
||||
try:
|
||||
if _signed_by(leaf, root):
|
||||
return root
|
||||
except CertError as exc:
|
||||
why = str(exc)
|
||||
continue
|
||||
why = "signing cert's signature does not verify against the CA root"
|
||||
raise CertError(why)
|
||||
Reference in New Issue
Block a user