maskflow-agent 0.1.0__py3-none-any.whl

maskflow_agent/__init__.py

@@ -0,0 +1,14 @@
+"""MaskFlow guard for AI coding agents.
+
+Claude Code hooks mask PII in every tool result before the model sees it and
+restore real values only inside allowlisted commands, at run time, in a
+wrapper process -- so the model, the transcript on disk and the permission
+prompt only ever see placeholders.
+
+Import-light on purpose: every hook invocation starts a fresh Python
+process, so this package's top level and the hook client (`hooks.py`,
+`protocol.py`, `config.py`, `paths.py`) use the standard library only. The
+masking engine is imported by the per-session daemon (`daemon.py`) alone.
+"""
+
+__version__ = "0.1.0"

maskflow_agent/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+raise SystemExit(main())

maskflow_agent/cli.py

@@ -0,0 +1,76 @@
+"""`maskflow-agent` command line. Stdlib only at import time -- `hook` and
+`exec` run on every tool call, so heavy modules load only in the
+subcommands that need them (`daemon`, `doctor`)."""
+
+from __future__ import annotations
+
+import sys
+
+from . import __version__
+
+
+def main(argv: list[str] | None = None) -> int:
+    argv = sys.argv[1:] if argv is None else argv
+    # Fast paths first: no argparse tree for the per-tool-call commands.
+    if argv[:1] == ["hook"] and len(argv) == 2:
+        from . import hooks
+
+        if argv[1] not in hooks.HANDLERS:
+            sys.stderr.write(f"maskflow-agent: unknown hook {argv[1]!r}\n")
+            return 2
+        return hooks.main(argv[1])
+    if argv[:1] == ["exec"]:
+        from . import execwrap
+
+        return execwrap.main(argv[1:])
+    if argv[:1] == ["daemon"]:
+        from . import daemon
+
+        return daemon.main(argv[1:])
+
+    import argparse
+    from pathlib import Path
+
+    parser = argparse.ArgumentParser(
+        prog="maskflow-agent",
+        description="MaskFlow PII guard for AI coding agents (Claude Code hooks).",
+    )
+    parser.add_argument("--version", action="version", version=f"maskflow-agent {__version__}")
+    sub = parser.add_subparsers(dest="cmd", required=True)
+    for name, help_text in (
+        ("install", "add the guard's hooks to Claude Code settings"),
+        ("uninstall", "remove the guard's hooks from Claude Code settings"),
+    ):
+        p = sub.add_parser(name, help=help_text)
+        p.add_argument(
+            "--scope",
+            choices=("project", "local", "user"),
+            default="project",
+            help="project: .claude/settings.json (default, commit it to share with "
+            "your team) | local: .claude/settings.local.json | user: ~/.claude/settings.json",
+        )
+    sub.add_parser("doctor", help="verify the guard end to end with synthetic canaries")
+    sub.add_parser("hook", help="(internal) run one Claude Code hook")
+    sub.add_parser("exec", help="(internal) run a command with placeholders restored")
+    sub.add_parser("daemon", help="(internal) per-session masking daemon")
+    args = parser.parse_args(argv)
+
+    cwd = Path.cwd()
+    if args.cmd in ("install", "uninstall"):
+        from . import install
+
+        path = install.settings_path(args.scope, cwd)
+        if args.cmd == "install":
+            install.install(path)
+            print(f"maskflow-agent hooks installed in {path}")
+            print("Start a new Claude Code session, then run `maskflow-agent doctor`.")
+        else:
+            removed = install.uninstall(path)
+            print(f"{'removed from' if removed else 'no maskflow-agent hooks in'} {path}")
+        return 0
+    if args.cmd == "doctor":
+        from . import doctor
+
+        return doctor.run(cwd)
+    parser.error(f"use `maskflow-agent {args.cmd}` only through Claude Code hooks")
+    return 2

maskflow_agent/config.py

@@ -0,0 +1,200 @@
+"""Guard configuration and the unmask allowlist. Stdlib only.
+
+Config is read from the first of `<cwd>/.claude/maskflow-agent.json`,
+`$MASKFLOW_AGENT_HOME/config.json`; missing files mean defaults:
+
+    {
+      "unmask_allow": ["psql", "mysql", "kubectl exec"],   // replaces defaults
+      "unmask_allow_extra": ["my-db-cli"],                 // added to defaults
+      "prompt_mode": "warn"                                // or "block" / "off"
+    }
+
+Placeholders are restored only when EVERY simple command in the Bash
+command line starts with an allowlisted program. Anything the check can't
+read with certainty -- command substitution, backticks, process
+substitution, unbalanced quotes -- is treated as not allowlisted, so the
+placeholder text is passed through literally.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import shlex
+from pathlib import Path
+from typing import NamedTuple
+
+from . import paths
+
+DEFAULT_UNMASK_ALLOW: tuple[str, ...] = (
+    "psql",
+    "mysql",
+    "mariadb",
+    "sqlite3",
+    "duckdb",
+    "mongosh",
+    "mongo",
+    "redis-cli",
+    "clickhouse-client",
+    "cqlsh",
+    "kubectl exec",
+)
+
+PROMPT_MODES = ("warn", "block", "off")
+
+_CONTROL = {";", "&&", "||", "|", "&", "\n"}
+_UNREADABLE = ("$(", "`", "<(", ">(")
+
+
+class AgentConfig(NamedTuple):
+    # NamedTuple, not a dataclass: dataclasses imports inspect (~10 ms), and
+    # this module loads on every tool call.
+    unmask_allow: tuple[str, ...] = DEFAULT_UNMASK_ALLOW
+    prompt_mode: str = "warn"
+    sources: tuple[str, ...] = ()
+
+
+def _candidates(cwd: str | None) -> list[Path]:
+    found = []
+    if cwd:
+        found.append(Path(cwd) / ".claude" / "maskflow-agent.json")
+    found.append(paths.home() / "config.json")
+    return found
+
+
+def load(cwd: str | None = None) -> AgentConfig:
+    for candidate in _candidates(cwd):
+        if not candidate.is_file():
+            continue
+        try:
+            raw = json.loads(candidate.read_text(encoding="utf-8"))
+        except (OSError, json.JSONDecodeError):
+            continue
+        if not isinstance(raw, dict):
+            continue
+        allow = raw.get("unmask_allow", DEFAULT_UNMASK_ALLOW)
+        extra = raw.get("unmask_allow_extra", ())
+        mode = raw.get("prompt_mode", "warn")
+        allow_t = tuple(str(a) for a in allow if isinstance(a, str) and a.strip())
+        extra_t = tuple(str(a) for a in extra if isinstance(a, str) and a.strip())
+        return AgentConfig(
+            unmask_allow=tuple(dict.fromkeys((*allow_t, *extra_t))),
+            prompt_mode=mode if mode in PROMPT_MODES else "warn",
+            sources=(str(candidate),),
+        )
+    return AgentConfig()
+
+
+def _simple_commands(command: str) -> list[list[str]] | None:
+    """Split a command line into simple commands (argv lists), or None when
+    it can't be read with certainty."""
+    if any(marker in command for marker in _UNREADABLE):
+        return None
+    try:
+        lexer = shlex.shlex(command, posix=True, punctuation_chars=";&|")
+        lexer.whitespace_split = True
+        lexer.commenters = ""
+        tokens = list(lexer)
+    except ValueError:
+        return None
+    commands: list[list[str]] = [[]]
+    for token in tokens:
+        if token in _CONTROL or set(token) <= set(";&|"):
+            commands.append([])
+        else:
+            commands[-1].append(token)
+    return [c for c in commands if c]
+
+
+def _program(argv: list[str]) -> list[str]:
+    """Drop leading VAR=value assignments and a leading `env`/`sudo`-free
+    path prefix from argv[0]."""
+    i = 0
+    while i < len(argv) and "=" in argv[i] and not argv[i].startswith("="):
+        name = argv[i].split("=", 1)[0]
+        if not name.replace("_", "").isalnum():
+            break
+        i += 1
+    rest = argv[i:]
+    if rest:
+        rest = [os.path.basename(rest[0]), *rest[1:]]
+    return rest
+
+
+def is_unmask_allowed(command: str, allow: tuple[str, ...]) -> bool:
+    commands = _simple_commands(command)
+    if not commands:
+        return False
+    rules = [rule.split() for rule in allow]
+    for argv in commands:
+        program = _program(argv)
+        if not program:
+            return False
+        if not any(program[: len(rule)] == rule for rule in rules if rule):
+            return False
+    return True
+
+
+# -- the user's own Bash permission rules ---------------------------------------
+# Claude Code evaluates permission rules against the command a PreToolUse hook
+# returns. Once a command is wrapped in `maskflow-agent exec`, a user rule such
+# as `Bash(psql:*)` in "deny" would no longer match it -- so before wrapping,
+# the hook checks the ORIGINAL command against the user's deny/ask rules
+# itself and keeps their decision.
+def _settings_files(cwd: str | None) -> list[Path]:
+    files = [Path.home() / ".claude" / "settings.json"]
+    if cwd:
+        files += [
+            Path(cwd) / ".claude" / "settings.json",
+            Path(cwd) / ".claude" / "settings.local.json",
+        ]
+    return files
+
+
+def _bash_rules(cwd: str | None, kind: str) -> list[str]:
+    rules: list[str] = []
+    for file in _settings_files(cwd):
+        try:
+            data = json.loads(file.read_text(encoding="utf-8"))
+        except (OSError, json.JSONDecodeError):
+            continue
+        perms = data.get("permissions") if isinstance(data, dict) else None
+        entries = perms.get(kind, []) if isinstance(perms, dict) else []
+        rules += [e for e in entries if isinstance(e, str) and e.startswith("Bash")]
+    return rules
+
+
+def _rule_matches(rule: str, command_text: str) -> bool:
+    if rule == "Bash" or rule == "Bash(*)":
+        return True
+    if not (rule.startswith("Bash(") and rule.endswith(")")):
+        return False
+    inner = rule[5:-1].strip()
+    if inner.endswith(":*"):
+        prefix = inner[:-2]
+        return command_text == prefix or command_text.startswith(prefix + " ")
+    if inner.endswith("*"):
+        return command_text.startswith(inner[:-1])
+    return command_text == inner
+
+
+def user_bash_decision(command: str, cwd: str | None) -> str | None:
+    """What the user's own rules say about the ORIGINAL command: "deny" or
+    "ask" if any deny/ask rule matches any of its simple commands, "allow"
+    only if allow rules cover EVERY simple command, else None (no rule --
+    Claude Code's normal flow decides). Unreadable commands never get
+    "allow"."""
+    commands = _simple_commands(command)
+    texts = [" ".join(_program(argv)) for argv in commands] if commands else None
+    for kind in ("deny", "ask"):
+        rules = _bash_rules(cwd, kind)
+        if not rules:
+            continue
+        if texts is None:
+            return kind
+        if any(_rule_matches(rule, text) for rule in rules for text in texts):
+            return kind
+    allow = _bash_rules(cwd, "allow")
+    if texts and allow and all(any(_rule_matches(r, t) for r in allow) for t in texts):
+        return "allow"
+    return None

maskflow_agent/daemon.py

@@ -0,0 +1,361 @@
+"""Per-session masking daemon.
+
+One process per Claude Code session, listening on a 0600 unix socket. It
+owns the session's `maskflow.Session` -- the only place real values live,
+in memory, never on disk -- and serialises every request through one lock,
+so hooks firing in parallel for parallel tool calls always see one
+consistent value <-> placeholder mapping.
+
+Requests and responses are single JSON lines (see protocol.py). Nothing in
+a request is ever logged; the log carries lifecycle events and counts only.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import logging
+import os
+import re
+import secrets
+import signal
+import socket
+import socketserver
+import sys
+import threading
+import time
+from collections import Counter
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+from . import __version__, paths
+
+_TOKEN_RE = re.compile(r"<([A-Z][A-Z_]*?)_\d+(?:_[0-9a-f]+)?>")
+_MAX_REQUEST = 256 * 1024 * 1024
+DEFAULT_IDLE_SECONDS = 6 * 3600
+TICKET_TTL = 600.0
+_NOISY_PROMPT_TYPES = frozenset({"PERSON_NAME", "INDIAN_ADDRESS", "PERSON", "LOCATION"})
+
+log = logging.getLogger("maskflow_agent.daemon")
+
+
+class _KnownValues:
+    """Every value this session has already masked, matched literally in
+    later text whatever its context -- a bare mobile in a `grep` hit is
+    masked because the same number was identified earlier under a `mobile`
+    header. Aho-Corasick (linear time), rebuilt only when the mapping grows.
+    A match must not continue an alphanumeric run on either side, so a known
+    number inside a longer number is left alone."""
+
+    def __init__(self) -> None:
+        self._size = -1
+        self._automaton: Any = None
+
+    def _rebuild(self, mapping: Any) -> None:
+        import ahocorasick
+
+        automaton = ahocorasick.Automaton()
+        for token in mapping:
+            entry = mapping[token]
+            original = getattr(entry, "original", None)
+            if entry.reversible and isinstance(original, str) and len(original) >= 4:
+                automaton.add_word(original, (original, token))
+        self._automaton = automaton if len(automaton) else None
+        if self._automaton is not None:
+            self._automaton.make_automaton()
+        self._size = len(mapping)
+
+    def apply(self, text: str, mapping: Any) -> str:
+        if len(mapping) != self._size:
+            self._rebuild(mapping)
+        if self._automaton is None or not text:
+            return text
+        hits: list[tuple[int, int, str]] = []
+        for end, (original, token) in self._automaton.iter(text):
+            start = end - len(original) + 1
+            before = text[start - 1] if start > 0 else ""
+            after = text[end + 1] if end + 1 < len(text) else ""
+            if (original[0].isalnum() and before.isalnum()) or (
+                original[-1].isalnum() and after.isalnum()
+            ):
+                continue
+            hits.append((start, end + 1, token))
+        if not hits:
+            return text
+        # Longest-first, then left-most, non-overlapping.
+        hits.sort(key=lambda h: (-(h[1] - h[0]), h[0]))
+        taken: list[tuple[int, int, str]] = []
+        for hit in hits:
+            if all(hit[1] <= t[0] or hit[0] >= t[1] for t in taken):
+                taken.append(hit)
+        taken.sort()
+        out, cursor = [], 0
+        for start, end, token in taken:
+            out.append(text[cursor:start])
+            out.append(token)
+            cursor = end
+        out.append(text[cursor:])
+        return "".join(out)
+
+
+class _State:
+    def __init__(self, session_id: str, patterns_only: bool) -> None:
+        # maskflow_core's config loader rejects unknown MASKFLOW_* variables;
+        # ours are consumed already (paths.home() is cached).
+        paths.home()
+        for name in [n for n in os.environ if n.startswith("MASKFLOW_AGENT_")]:
+            del os.environ[name]
+        import maskflow  # heavy: engine + packs + (lazily) spaCy
+
+        self.session_id = session_id
+        self.ner = not patterns_only and _spacy_model_available()
+        if not patterns_only and not self.ner:
+            log.warning(
+                "spaCy model en_core_web_sm not installed: masking IDs/numbers only, "
+                "not names/addresses (python -m spacy download en_core_web_sm)"
+            )
+        self.session = maskflow.Session(ttl_seconds=None, patterns_only=not self.ner)
+        self.lock = threading.Lock()
+        self.last_seen = time.monotonic()
+        self.started = time.time()
+        self.masked_occurrences: Counter[str] = Counter()
+        self.prompt_findings: Counter[str] = Counter()
+        self.unmask_calls = 0
+        self.tool_results = 0
+        self.withheld = 0
+        self.known = _KnownValues()
+        # Single-use run-time-unmask tickets: sha256(command) -> (ticket, expiry).
+        self.tickets: dict[str, tuple[str, float]] = {}
+
+    # -- helpers -----------------------------------------------------------
+    def tokens(self) -> set[str]:
+        return set(self.session.mapping)
+
+    def mask_text(self, text: str) -> str:
+        masked = self.known.apply(self.session.mask(text), self.session.mapping)
+        if masked != text:
+            for match in _TOKEN_RE.finditer(masked):
+                if match.group(0) in self.session.mapping:
+                    self.masked_occurrences[match.group(1)] += 1
+        return masked
+
+    def mask_value(self, value: Any, depth: int = 0) -> Any:
+        if depth > 64:
+            raise ValueError("value nested too deeply")
+        if isinstance(value, str):
+            return self.mask_text(value)
+        if isinstance(value, list):
+            return [self.mask_value(v, depth + 1) for v in value]
+        if isinstance(value, dict):
+            return {k: self.mask_value(v, depth + 1) for k, v in value.items()}
+        return value
+
+    def distinct_values(self) -> dict[str, int]:
+        counts: Counter[str] = Counter()
+        for entry in self.session.mapping.values():
+            counts[str(entry.entity_type)] += 1
+        return dict(counts)
+
+    def summary(self) -> dict[str, Any]:
+        return {
+            "session_id": self.session_id,
+            "maskflow_agent": __version__,
+            "ner": self.ner,
+            "started_at": self.started,
+            "ended_at": time.time(),
+            "tool_results_masked": self.tool_results,
+            "masked_occurrences_by_type": dict(self.masked_occurrences),
+            "distinct_values_by_type": self.distinct_values(),
+            "prompt_findings_by_type": dict(self.prompt_findings),
+            "run_time_unmask_calls": self.unmask_calls,
+        }
+
+
+def _digest(command: str) -> str:
+    return hashlib.sha256(command.encode("utf-8")).hexdigest()
+
+
+def _spacy_model_available() -> bool:
+    try:
+        from maskflow_core.ner import _get_nlp
+
+        _get_nlp()
+    except (OSError, ImportError):
+        return False
+    return True
+
+
+def _handle(state: _State, request: dict[str, Any], stop: Callable[[], None]) -> dict[str, Any]:
+    op = request.get("op")
+    if op == "ping":
+        return {"ok": True, "version": __version__, "pid": os.getpid(), "ner": state.ner}
+    if op == "mask_value":
+        before = len(state.session.mapping)
+        masked = state.mask_value(request.get("value"))
+        state.tool_results += 1
+        return {
+            "ok": True,
+            "value": masked,
+            "changed": masked != request.get("value"),
+            "new_values": len(state.session.mapping) - before,
+        }
+    if op == "issue_ticket":
+        # Only the PreToolUse hook asks for these, for the exact command it
+        # has just checked against the allowlist and the user's rules.
+        command = str(request.get("command", ""))
+        ticket = secrets.token_hex(16)
+        now = time.monotonic()
+        state.tickets = {k: v for k, v in state.tickets.items() if v[1] > now}
+        state.tickets[_digest(command)] = (ticket, now + TICKET_TTL)
+        return {"ok": True, "ticket": ticket}
+    if op == "unmask_text":
+        # Never without a ticket: neither a model-crafted `maskflow-agent exec`
+        # call nor a direct socket client can recover a value.
+        text = str(request.get("text", ""))
+        issued = state.tickets.pop(_digest(text), None)
+        if (
+            issued is None
+            or issued[1] < time.monotonic()
+            or not secrets.compare_digest(issued[0], str(request.get("ticket", "")))
+        ):
+            return {"ok": False, "error": "no valid ticket"}
+        state.unmask_calls += 1
+        return {"ok": True, "text": state.session.unmask(text)}
+    if op == "tokens_in":
+        text = str(request.get("text", ""))
+        known = state.tokens()
+        found = sorted({m.group(0) for m in _TOKEN_RE.finditer(text)} & known)
+        return {"ok": True, "tokens": found}
+    if op == "scan":
+        # Prompt warnings use the precise pattern recognizers only: NER
+        # name/address guesses on short instructions ("SELECT", "PAN") are
+        # noise the user would learn to ignore.
+        from maskflow_core.detection import detect_patterns_only
+
+        counts = Counter(
+            str(s.entity_type)
+            for s in detect_patterns_only(str(request.get("text", "")))
+            if str(s.entity_type) not in _NOISY_PROMPT_TYPES
+        )
+        state.prompt_findings.update(counts)
+        return {"ok": True, "counts": dict(counts)}
+    if op == "reserve":
+        tokens = [t for t in request.get("tokens", []) if isinstance(t, str)]
+        if tokens:
+            # Session.mask() reserves every placeholder look-alike in its
+            # input, so values masked from now on never reuse these names.
+            state.session.mask(" ".join(sorted(set(tokens))))
+        return {"ok": True, "reserved": len(set(tokens))}
+    if op == "stats":
+        return {"ok": True, "summary": state.summary()}
+    if op == "end":
+        summary = state.summary()
+        try:
+            out = paths.summary_path(state.session_id)
+            out.write_text(json.dumps(summary, indent=2), encoding="utf-8")
+            os.chmod(out, 0o600)
+        except OSError:
+            log.warning("could not write session summary")
+        stop()
+        return {"ok": True, "summary": summary}
+    return {"ok": False, "error": "unknown op"}
+
+
+class _Server(socketserver.ThreadingUnixStreamServer):
+    daemon_threads = True
+    state: _State
+    stop_event: threading.Event
+
+
+class _RequestHandler(socketserver.StreamRequestHandler):
+    server: _Server
+
+    def handle(self) -> None:
+        raw = self.rfile.read(_MAX_REQUEST + 1)
+        if len(raw) > _MAX_REQUEST:
+            self._reply({"ok": False, "error": "request too large"})
+            return
+        try:
+            request = json.loads(raw.decode("utf-8"))
+            if not isinstance(request, dict):
+                raise ValueError
+        except (UnicodeDecodeError, ValueError):
+            self._reply({"ok": False, "error": "bad request"})
+            return
+        state = self.server.state
+        with state.lock:
+            state.last_seen = time.monotonic()
+            try:
+                response = _handle(state, request, self.server.stop_event.set)
+            except Exception as exc:  # noqa: BLE001 -- report type only, never content
+                log.error("request failed: %s", type(exc).__name__)
+                response = {"ok": False, "error": type(exc).__name__}
+        self._reply(response)
+
+    def _reply(self, response: dict[str, Any]) -> None:
+        self.wfile.write(json.dumps(response, ensure_ascii=False).encode("utf-8") + b"\n")
+
+
+def serve(session_id: str, idle_seconds: float, patterns_only: bool) -> int:
+    sock_path = paths.socket_path(session_id)
+    if sock_path.exists():
+        # Another live daemon owns this session: don't steal its socket.
+        try:
+            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe:
+                probe.settimeout(1.0)
+                probe.connect(str(sock_path))
+            log.info("daemon already running for this session")
+            return 0
+        except OSError:
+            sock_path.unlink(missing_ok=True)
+
+    state = _State(session_id, patterns_only)
+    # Load the engine (and spaCy, unless patterns_only) before accepting
+    # connections, so the first real request isn't the slow one.
+    state.session.mask("warm-up")
+
+    old_umask = os.umask(0o177)
+    try:
+        server = _Server(str(sock_path), _RequestHandler)
+    finally:
+        os.umask(old_umask)
+    server.state = state
+    server.stop_event = threading.Event()
+
+    def _watch() -> None:
+        while not server.stop_event.wait(5.0):
+            if time.monotonic() - state.last_seen > idle_seconds:
+                log.info("idle timeout")
+                server.stop_event.set()
+        server.shutdown()
+
+    def _on_signal(signum: int, _frame: object) -> None:
+        server.stop_event.set()
+
+    signal.signal(signal.SIGTERM, _on_signal)
+    signal.signal(signal.SIGINT, _on_signal)
+    threading.Thread(target=_watch, daemon=True).start()
+    log.info("daemon ready pid=%d", os.getpid())
+    try:
+        server.serve_forever(poll_interval=0.2)
+    finally:
+        server.server_close()
+        Path(sock_path).unlink(missing_ok=True)
+        log.info("daemon stopped")
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="maskflow-agent daemon")
+    parser.add_argument("--session", required=True)
+    parser.add_argument("--idle-seconds", type=float, default=DEFAULT_IDLE_SECONDS)
+    parser.add_argument("--patterns-only", action="store_true")
+    args = parser.parse_args(argv)
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s %(levelname)s %(message)s",
+        stream=sys.stderr,
+    )
+    return serve(args.session, args.idle_seconds, args.patterns_only)