keyfence 0.2.0__py3-none-any.whl

keyfence/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.2.0"

keyfence/__main__.py

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

keyfence/addon.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+import json
+import logging
+import os
+import sys
+import time
+from pathlib import Path
+
+_PKG_PARENT = str(Path(__file__).resolve().parent.parent)
+if _PKG_PARENT not in sys.path:
+    sys.path.insert(0, _PKG_PARENT)
+
+from mitmproxy import http  # noqa: E402
+
+from keyfence.config import Config  # noqa: E402
+from keyfence.detectors import Finding, scan  # noqa: E402
+from keyfence.notice import add_notice  # noqa: E402
+from keyfence.streaming import SSERestorer, restore  # noqa: E402
+from keyfence.vault import Vault  # noqa: E402
+
+log = logging.getLogger("keyfence")
+ENV_VAULT_VAR = "KEYFENCE_ENV_VAULT"
+MAPPING_KEY = "keyfence_mapping"
+STREAMED_KEY = "keyfence_streamed"
+AUDIT_PREVIEW_LIMIT = 50
+PLACEHOLDER_ID_LENGTH = 10
+
+
+class KeyFence:
+    def __init__(self):
+        self.config = Config.load()
+        self.vault = self._load_vault()
+        self._vault_mtime = self._mtime(self.vault.path)
+        self.stats = {"scanned": 0, "findings": 0, "blocked": 0, "errors": 0, "canaries": 0}
+
+    @staticmethod
+    def _mtime(path: Path) -> float:
+        try:
+            return Path(path).stat().st_mtime
+        except OSError:
+            return 0.0
+
+    def _load_vault(self) -> Vault:
+        vault = Vault()
+        vault.ensure_saved()
+        extra = os.environ.get(ENV_VAULT_VAR)
+        if extra and Path(extra).exists():
+            vault.merge(Vault(path=Path(extra)))
+        return vault
+
+    def _maybe_reload_vault(self) -> None:
+        mtime = self._mtime(self.vault.path)
+        if mtime != self._vault_mtime:
+            self.vault = self._load_vault()
+            self._vault_mtime = mtime
+            log.info("vault reloaded: %d secret(s)", self.vault.count())
+
+    def load(self, loader):
+        log.info("mode=%s | %d hosts monitored | %d rules | vault with %d secret(s)",
+                 self.config.mode, len(self.config.hosts),
+                 len(self.config.scan.rules), self.vault.count())
+
+    def request(self, flow: http.HTTPFlow) -> None:
+        host = flow.request.pretty_host
+        if not self.config.host_matches(host):
+            return
+        try:
+            self._inspect(flow, host)
+        except Exception as exc:
+            self.stats["errors"] += 1
+            log.error("detector failure, failing closed for %s: %r", host, exc)
+            flow.response = self._blocked_response(
+                f"keyfence internal error ({type(exc).__name__}); request not sent.")
+
+    def _inspect(self, flow: http.HTTPFlow, host: str) -> None:
+        text = flow.request.get_text(strict=False)
+        if not text:
+            return
+        self._maybe_reload_vault()
+        self.stats["scanned"] += 1
+        findings = scan(text, vault=self.vault, config=self.config.scan)
+        if not findings:
+            return
+
+        self.stats["findings"] += len(findings)
+        self._audit(flow, findings)
+        tripped = sorted({self.vault.canary_label(f.value) for f in findings if f.kind == "canary"})
+        if tripped:
+            self.stats["canaries"] += len(tripped)
+            log.warning("CANARY tripped -> %s: %s was read and sent", host, ", ".join(tripped))
+
+        if self.config.mode == "block":
+            self.stats["blocked"] += 1
+            kinds = sorted({f.kind for f in findings})
+            flow.response = self._blocked_response(
+                f"Request blocked by keyfence: {len(findings)} secret(s) detected "
+                f"({', '.join(kinds)}). Nothing was sent to the provider.")
+            log.warning("BLOCKED -> %s: %d secret(s)", host, len(findings))
+            return
+
+        mapping: dict[str, str] = {}
+        new_text = text
+        for f in sorted(findings, key=lambda f: f.start, reverse=True):
+            if self.config.mode == "placeholder":
+                token = self._placeholder(f.value, mapping)
+                mapping[token] = f.value
+            else:
+                token = f"[REDACTED:{f.kind}]"
+            new_text = new_text[:f.start] + token + new_text[f.end:]
+
+        if self.config.notice:
+            new_text = add_notice(new_text, host)
+        flow.request.set_text(new_text)
+        if mapping:
+            flow.metadata[MAPPING_KEY] = mapping
+            flow.request.headers["accept-encoding"] = "identity"
+        log.warning("%s -> %s: %d secret(s) removed from request",
+                    self.config.mode.upper(), host, len(findings))
+
+    def _placeholder(self, value: str, mapping: dict[str, str]) -> str:
+        digest = self.vault.placeholder_digest(value)
+        length = PLACEHOLDER_ID_LENGTH
+        token = f"<<SECRET_{digest[:length]}>>"
+        while token in mapping and mapping[token] != value:
+            length += 2
+            token = f"<<SECRET_{digest[:length]}>>"
+        return token
+
+    @staticmethod
+    def _blocked_response(message: str) -> http.Response:
+        return http.Response.make(
+            403,
+            json.dumps({"error": {"type": "keyfence_blocked", "message": message}},
+                       ensure_ascii=False),
+            {"Content-Type": "application/json"},
+        )
+
+    def responseheaders(self, flow: http.HTTPFlow) -> None:
+        if flow.response is None:
+            return
+        mapping = flow.metadata.get(MAPPING_KEY)
+        if not mapping:
+            flow.response.stream = True
+            return
+        content_type = flow.response.headers.get("content-type", "")
+        encoding = flow.response.headers.get("content-encoding", "identity")
+        if "text/event-stream" in content_type and encoding in ("identity", ""):
+            flow.response.stream = SSERestorer(mapping).feed
+            flow.metadata[STREAMED_KEY] = True
+
+    def response(self, flow: http.HTTPFlow) -> None:
+        mapping = flow.metadata.get(MAPPING_KEY)
+        if not mapping or flow.metadata.get(STREAMED_KEY) or flow.response is None:
+            return
+        text = flow.response.get_text(strict=False)
+        if not text:
+            return
+        restored = restore(text, mapping)
+        if restored != text:
+            flow.response.set_text(restored)
+            log.info("placeholders restored in response")
+
+    def _audit_finding(self, f: Finding) -> dict:
+        entry = {"kind": f.kind, "preview": f.masked, "key": f.key}
+        if f.kind == "canary":
+            entry["label"] = self.vault.canary_label(f.value)
+        return entry
+
+    def _audit(self, flow: http.HTTPFlow, findings: list[Finding]) -> None:
+        try:
+            path = Path(self.config.audit_log)
+            path.parent.mkdir(parents=True, exist_ok=True)
+            entry = {
+                "ts": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
+                "host": flow.request.pretty_host,
+                "path": flow.request.path.split("?")[0],
+                "mode": self.config.mode,
+                "count": len(findings),
+                "findings": [self._audit_finding(f) for f in findings[:AUDIT_PREVIEW_LIMIT]],
+            }
+            with path.open("a") as fh:
+                fh.write(json.dumps(entry, ensure_ascii=False) + "\n")
+        except OSError as exc:
+            log.warning("could not write audit log: %s", exc)
+
+
+addons = [KeyFence()]

keyfence/cli.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+import argparse
+import getpass
+import json
+import os
+import re
+import secrets
+import subprocess
+import sys
+from collections import Counter
+from pathlib import Path
+
+from . import runner
+from .config import Config
+from .detectors import scan
+from .importer import default_paths, env_values, import_files
+from .vault import DEFAULT_DIR, Vault
+
+
+def cmd_add_secret(_args) -> int:
+    vault = Vault()
+    print("Paste the secret to protect (hidden; only its hash is stored):")
+    value = getpass.getpass("> ")
+    if not value.strip():
+        print("Nothing entered, aborting.")
+        return 1
+    if not vault.add(value):
+        print(f"Secret too short (minimum {vault.min_length} characters); "
+              "shorter values would cause too many false positives.")
+        return 1
+    print(f"OK. Vault now holds {vault.count()} secret(s) (hashes only, in {vault.path}).")
+    return 0
+
+
+def cmd_import(args) -> int:
+    vault = Vault()
+    paths = [Path(p).expanduser() for p in args.paths] or default_paths()
+    missing = [p for p in paths if not p.is_file()]
+    for p in missing:
+        print(f"skip: {p} (not a file)")
+    paths = [p for p in paths if p.is_file()]
+    if not paths and not args.env:
+        print("Nothing to import. Pass file paths, or use --env to import from the environment.")
+        return 1
+    total = 0
+    for path, added in import_files(vault, paths, everything=args.all):
+        total += added
+        print(f"{path}: {added} new secret(s)")
+    if args.env:
+        added = vault.add_many(env_values(os.environ, vault.min_length, everything=args.all))
+        total += added
+        print(f"environment: {added} new secret(s)")
+    print(f"Done. {total} new secret(s); vault now holds {vault.count()} (hashes only).")
+    return 0
+
+
+def cmd_canary(args) -> int:
+    path = Path(args.file)
+    name = args.name
+    existing = path.read_text(errors="replace") if path.exists() else ""
+    if re.search(rf"^\s*(?:export\s+)?{re.escape(name)}\s*=", existing, re.MULTILINE):
+        print(f"{path} already defines {name}. Pick another name with --name.")
+        return 1
+    value = secrets.token_urlsafe(24)
+    vault = Vault()
+    vault.add_canary(value, str(path.resolve()))
+    prefix = "" if not existing or existing.endswith("\n") else "\n"
+    with path.open("a") as fh:
+        fh.write(f"{prefix}{name}={value}\n")
+    print(f"Canary planted in {path} as {name} and registered in the vault (hash only).")
+    print("If it ever shows up in a request, keyfence logs a 'canary' detection with this file's path.")
+    return 0
+
+
+def cmd_scan(args) -> int:
+    if args.file:
+        text = Path(args.file).read_text(errors="replace")
+    elif args.text:
+        text = args.text
+    else:
+        text = sys.stdin.read()
+    cfg = Config.load()
+    findings = scan(text, vault=Vault(), config=cfg.scan)
+    if not findings:
+        print("No secrets detected.")
+        return 0
+    print(f"{len(findings)} secret(s) detected:")
+    for f in findings:
+        where = f" in \"{f.key}\"" if f.key else ""
+        print(f"  - [{f.kind}] {f.masked} (offset {f.start}-{f.end}{where})")
+    return 2
+
+
+def cmd_run(args) -> int:
+    print(f"Starting keyfence on http://127.0.0.1:{args.port}")
+    print("Point your tools at it, e.g.:")
+    print(f"  export HTTPS_PROXY=http://127.0.0.1:{args.port}")
+    print(f"  export HTTP_PROXY=http://127.0.0.1:{args.port}")
+    print("or run them through it directly: keyfence exec -- <command>")
+    print("(Ctrl+C to stop)\n")
+    try:
+        return subprocess.call(runner.proxy_command(args.port))
+    except FileNotFoundError:
+        print("mitmdump not found. Install it with: pip install mitmproxy")
+        return 1
+
+
+def cmd_exec(args) -> int:
+    command = list(args.argv)
+    if command and command[0] == "--":
+        command = command[1:]
+    if not command:
+        print("usage: keyfence exec [-p PORT] [--all-env] -- <command> [args...]")
+        return 1
+    return runner.run(command, args.port, everything=args.all_env)
+
+
+def cmd_status(_args) -> int:
+    cfg = Config.load()
+    vault = Vault()
+    print(f"Mode:            {cfg.mode}")
+    print(f"Hosts:           {len(cfg.hosts)} monitored"
+          + (" (intercepting ALL hosts)" if cfg.intercept_all_hosts else ""))
+    print(f"Vault:           {vault.count()} secret(s), {vault.canary_count()} canary(ies) in {vault.path}")
+    print(f"Rules:           {len(cfg.scan.rules)} gitleaks rules"
+          + ("" if cfg.scan.gitleaks else " (disabled)"))
+    print(f"Entropy:         {'on' if cfg.scan.entropy_enabled else 'off'}"
+          f" (min_len={cfg.scan.entropy_min_length}, threshold={cfg.scan.entropy_threshold})")
+    print(f"Audit log:       {cfg.audit_log}")
+    log = Path(cfg.audit_log)
+    if log.exists():
+        lines = log.read_text().strip().splitlines()[-5:]
+        if lines:
+            print("\nRecent detections:")
+            for line in lines:
+                try:
+                    e = json.loads(line)
+                    counts = Counter(f["kind"] for f in e["findings"])
+                    kinds = ", ".join(
+                        f"{kind} x{n}" if n > 1 else kind for kind, n in counts.most_common())
+                    total = e.get("count", len(e["findings"]))
+                    suffix = f"  {total} total" if total > len(e["findings"]) else ""
+                    print(f"  {e['ts']}  {e['host']}  [{kinds}]{suffix}  ({e['mode']})")
+                except (json.JSONDecodeError, KeyError):
+                    pass
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="keyfence",
+        description="Local proxy that keeps your secrets out of LLM requests.")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    sub.add_parser("add-secret", help="register one of your secrets in the vault (hash only)")
+
+    p_import = sub.add_parser(
+        "import", help="register secrets found in .env files, credential stores or the environment")
+    p_import.add_argument("paths", nargs="*", help="files to read (default: .env* and common credential files)")
+    p_import.add_argument("--env", action="store_true", help="also import values from environment variables")
+    p_import.add_argument("--all", action="store_true", help="import every value, not only secret-looking ones")
+
+    p_canary = sub.add_parser(
+        "canary", help="plant a fake secret in a file; keyfence reports if a tool ever sends it")
+    p_canary.add_argument("file", nargs="?", default=".env", help="file to append to (default: .env)")
+    p_canary.add_argument("--name", default="INTERNAL_API_TOKEN", help="variable name to use")
+
+    p_scan = sub.add_parser("scan", help="test detection on text, a file or stdin")
+    p_scan.add_argument("text", nargs="?", help="text to scan")
+    p_scan.add_argument("-f", "--file", help="file to scan")
+
+    p_run = sub.add_parser("run", help="start the proxy")
+    p_run.add_argument("-p", "--port", type=int, default=8888)
+
+    p_exec = sub.add_parser("exec", help="run a command with the proxy already wired in")
+    p_exec.add_argument("-p", "--port", type=int, default=8888)
+    p_exec.add_argument("--all-env", action="store_true",
+                        help="treat every environment variable value as a secret, not only secret-looking names")
+    p_exec.add_argument("argv", nargs=argparse.REMAINDER, metavar="command")
+
+    sub.add_parser("status", help="show configuration and recent detections")
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    DEFAULT_DIR.mkdir(parents=True, exist_ok=True)
+    return {
+        "add-secret": cmd_add_secret,
+        "import": cmd_import,
+        "canary": cmd_canary,
+        "scan": cmd_scan,
+        "run": cmd_run,
+        "exec": cmd_exec,
+        "status": cmd_status,
+    }[args.command](args)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

keyfence/config.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import fnmatch
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import yaml
+
+from .detectors import ScanConfig
+from .rules import DEFAULT_DISABLED, load_rules
+from .vault import DEFAULT_DIR
+
+DEFAULT_AI_HOSTS = [
+    "api.openai.com",
+    "api.anthropic.com",
+    "generativelanguage.googleapis.com",
+    "api.mistral.ai",
+    "api.groq.com",
+    "api.cohere.com",
+    "api.together.xyz",
+    "api.deepseek.com",
+    "api.x.ai",
+    "openrouter.ai",
+    "api.perplexity.ai",
+    "api.fireworks.ai",
+    "bedrock-runtime.*.amazonaws.com",
+    "*.openai.azure.com",
+]
+
+MODES = ("block", "redact", "placeholder")
+
+
+@dataclass
+class Config:
+    mode: str = "redact"
+    hosts: list[str] = field(default_factory=lambda: list(DEFAULT_AI_HOSTS))
+    intercept_all_hosts: bool = False
+    notice: bool = True
+    scan: ScanConfig = field(default_factory=ScanConfig)
+    audit_log: Path = field(default_factory=lambda: DEFAULT_DIR / "audit.log")
+
+    @classmethod
+    def load(cls, path: str | os.PathLike | None = None) -> "Config":
+        cfg_path = Path(
+            path
+            or os.environ.get("KEYFENCE_CONFIG", DEFAULT_DIR / "config.yaml"))
+        cfg = cls()
+        if cfg_path.exists():
+            data = yaml.safe_load(cfg_path.read_text()) or {}
+            cfg.mode = data.get("mode", cfg.mode)
+            if data.get("hosts"):
+                cfg.hosts = list(data["hosts"])
+            if data.get("extra_hosts"):
+                cfg.hosts.extend(data["extra_hosts"])
+            cfg.intercept_all_hosts = bool(
+                data.get("intercept_all_hosts", cfg.intercept_all_hosts))
+            cfg.notice = bool(data.get("notice", cfg.notice))
+            scan = data.get("scan") or {}
+            cfg.scan = ScanConfig(
+                patterns_enabled=scan.get("patterns", True),
+                entropy_enabled=scan.get("entropy", True),
+                entropy_min_length=scan.get("entropy_min_length", 24),
+                entropy_threshold=scan.get("entropy_threshold", 4.5),
+                entropy_max_length=scan.get("entropy_max_length", 512),
+                allowlist=scan.get("allowlist") or [],
+                gitleaks=scan.get("gitleaks", True),
+                gitleaks_rules=scan.get("gitleaks_rules"),
+                gitleaks_disabled=list(scan.get("gitleaks_disabled", DEFAULT_DISABLED)),
+            )
+            if data.get("audit_log"):
+                cfg.audit_log = Path(data["audit_log"]).expanduser()
+        if cfg.mode not in MODES:
+            raise ValueError(f"invalid mode: {cfg.mode!r} (expected one of {', '.join(MODES)})")
+        if cfg.scan.gitleaks:
+            cfg.scan.rules = load_rules(cfg.scan.gitleaks_rules, cfg.scan.gitleaks_disabled)
+        return cfg
+
+    def host_matches(self, host: str) -> bool:
+        if self.intercept_all_hosts:
+            return True
+        host = host.lower()
+        for pattern in self.hosts:
+            p = pattern.lower()
+            if "*" in p:
+                if fnmatch.fnmatch(host, p):
+                    return True
+            elif host == p or host.endswith("." + p):
+                return True
+        return False