vault-engine 0.1.0__py3-none-any.whl

vaultengine/cli.py

@@ -0,0 +1,269 @@
+"""Command-line interface.
+
+    vault-engine scrub  INFILE [-o OUT] [--map MAP] [--report REPORT] ...
+    vault-engine rehydrate INFILE --map MAP [-o OUT]
+    vault-engine providers
+    vault-engine models
+    vault-engine version
+
+Use ``-`` for INFILE/OUT to read stdin / write stdout.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from typing import List, Optional
+
+from . import __version__, formats
+from .config import POLICIES, Config
+from .mapping import Vault
+from .pipeline import deidentify, rehydrate
+from .providers import available
+from .providers.ollama import DEFAULT_ENDPOINT, installed_models
+
+EXIT_OK = 0
+EXIT_ERROR = 1
+EXIT_DEGRADED = 3   # ran, but a requested protection layer did not (under-redacted)
+
+DEFAULT_CLIP_MAP = os.path.join(os.path.expanduser("~"), ".vault-engine",
+                                "clip.map.json")
+
+
+def _read(path: str) -> str:
+    if path == "-":
+        return sys.stdin.read()
+    with open(path, encoding="utf-8") as fh:
+        return fh.read()
+
+
+def _write(path: str, data: str) -> None:
+    if path == "-":
+        sys.stdout.write(data)
+        return
+    with open(path, "w", encoding="utf-8") as fh:
+        fh.write(data)
+
+
+def _default_out(infile: str) -> str:
+    if infile == "-":
+        return "-"
+    stem, ext = os.path.splitext(infile)
+    return f"{stem}.scrubbed{ext or '.txt'}"
+
+
+def _config_from(args: argparse.Namespace) -> Config:
+    overrides = {
+        "provider": args.provider, "model": args.model,
+        "endpoint": args.endpoint, "api_key": args.api_key,
+        "policy": args.policy, "locale": args.locale,
+        "use_llm": (False if args.no_llm else None),
+        "critic": (False if args.no_critic else None),
+    }
+    return Config.load(path=args.config,
+                       overrides={k: v for k, v in overrides.items()
+                                  if v is not None})
+
+
+def cmd_scrub(args: argparse.Namespace) -> int:
+    text = _read(args.infile)
+    config = _config_from(args)
+    segs = formats.segment(text, args.format)
+    result = deidentify(text, config, segments=segs)
+
+    out = args.out or _default_out(args.infile)
+    _write(out, result.text)
+
+    # reverse map — local only, never for the cloud
+    if not args.one_way:
+        map_path = args.map or (
+            "vault.map.json" if out == "-" else out + ".map.json")
+        if map_path != "-":
+            result.vault.save(map_path)
+        else:
+            sys.stdout.write("\n" + json.dumps(result.map, ensure_ascii=False))
+
+    if args.report:
+        _write(args.report, result.report.to_markdown())
+
+    # always tell the human what happened (on stderr so stdout stays clean)
+    print(result.report.summary(), file=sys.stderr)
+    for w in result.report.warnings:
+        print("  ⚠️ " + w, file=sys.stderr)
+    if not args.one_way and out != "-":
+        print(f"  脱敏文本 → {out}", file=sys.stderr)
+        print(f"  反向映射 → {map_path}  （本地保存，切勿出云/入库）", file=sys.stderr)
+
+    if not result.safe and not args.allow_degraded:
+        print("✗ 请求的模型检测层未运行；输出可能脱敏不足。确认无碍可加 "
+              "--allow-degraded。", file=sys.stderr)
+        return EXIT_DEGRADED
+    return EXIT_OK
+
+
+def cmd_rehydrate(args: argparse.Namespace) -> int:
+    payload = _read(args.infile)
+    vault = Vault.load(args.map)
+    try:                                   # structured reply: rehydrate in place
+        obj = json.loads(payload)
+        restored = json.dumps(rehydrate(obj, vault), ensure_ascii=False, indent=2)
+    except json.JSONDecodeError:           # plain text reply
+        restored = rehydrate(payload, vault)
+    _write(args.out or "-", restored)
+    if args.out and args.out != "-":
+        print(f"已回填真实身份 → {args.out}", file=sys.stderr)
+    return EXIT_OK
+
+
+def cmd_clip(args: argparse.Namespace) -> int:
+    """Scrub (or rehydrate) the clipboard in place — the paste-into-ChatGPT hook."""
+    from . import clipboard
+    map_path = args.map or DEFAULT_CLIP_MAP
+    try:
+        text = clipboard.read_clipboard()
+    except clipboard.ClipboardError as exc:
+        print(f"✗ {exc}", file=sys.stderr)
+        return EXIT_ERROR
+    if not text.strip():
+        print("剪贴板为空。", file=sys.stderr)
+        return EXIT_ERROR
+
+    if args.rehydrate:
+        vault = Vault.load(map_path)
+        try:
+            restored = json.dumps(rehydrate(json.loads(text), vault),
+                                  ensure_ascii=False, indent=2)
+        except json.JSONDecodeError:
+            restored = rehydrate(text, vault)
+        clipboard.write_clipboard(restored)
+        print("✓ 剪贴板已用本地映射还原为真实身份。", file=sys.stderr)
+        return EXIT_OK
+
+    config = _config_from(args)
+    segs = formats.segment(text, args.format)
+    result = deidentify(text, config, segments=segs)
+    clipboard.write_clipboard(result.text)
+    if not args.one_way:
+        os.makedirs(os.path.dirname(map_path), exist_ok=True)
+        result.vault.save(map_path)
+
+    print(result.report.summary(), file=sys.stderr)
+    for w in result.report.warnings:
+        print("  ⚠️ " + w, file=sys.stderr)
+    print("✓ 剪贴板已脱敏，可直接粘贴给云端 AI。", file=sys.stderr)
+    if not args.one_way:
+        print(f"  反向映射 → {map_path}（本地保存，勿出云）", file=sys.stderr)
+        print("  云端回复贴回剪贴板后，用 `vault-engine clip --rehydrate` 还原。",
+              file=sys.stderr)
+    if not result.safe and not args.allow_degraded:
+        print("✗ 模型检测层未运行，可能脱敏不足（--allow-degraded 忽略）。", file=sys.stderr)
+        return EXIT_DEGRADED
+    return EXIT_OK
+
+
+def cmd_providers(_args: argparse.Namespace) -> int:
+    print("可用 provider：")
+    notes = {"ollama": "本地 Ollama（默认，原文不出本机）",
+             "openai-compat": "OpenAI 兼容端点（⚠️ 原文会出本机，默认不用）",
+             "null": "纯正则离线（无模型，仅兜底）"}
+    for name in available():
+        print(f"  - {name:14s} {notes.get(name, '')}")
+    return EXIT_OK
+
+
+def cmd_models(args: argparse.Namespace) -> int:
+    endpoint = args.endpoint or DEFAULT_ENDPOINT
+    models = installed_models(endpoint)
+    if not models:
+        print(f"未能从 {endpoint} 获取模型列表（Ollama 未运行？）", file=sys.stderr)
+        return EXIT_ERROR
+    print(f"{endpoint} 已安装模型：")
+    for m in models:
+        print(f"  - {m}")
+    return EXIT_OK
+
+
+def cmd_version(_args: argparse.Namespace) -> int:
+    print(f"vault-engine {__version__}")
+    return EXIT_OK
+
+
+def _add_model_flags(p: argparse.ArgumentParser) -> None:
+    p.add_argument("--provider", help="ollama(默认) / openai-compat / null")
+    p.add_argument("--model", help="模型 tag，默认 qwen3.6:27b")
+    p.add_argument("--endpoint", help="provider 端点 URL")
+    p.add_argument("--api-key", dest="api_key", help="远程端点的 API key")
+    p.add_argument("--config", help="JSON 配置文件路径（亦读 VAULT_* 环境变量）")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="vault-engine",
+        description="出云前的身份脱敏：本地检测 + 一致化假名 + 可逆回路（默认本地模型）。")
+    parser.add_argument("--version", action="version",
+                        version=f"vault-engine {__version__}")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    s = sub.add_parser("scrub", help="脱敏一个文件/标准输入")
+    s.add_argument("infile", help="输入文件（- 表示 stdin）")
+    s.add_argument("-o", "--out", help="脱敏输出（默认 <name>.scrubbed.<ext>；- 表示 stdout）")
+    s.add_argument("--map", help="反向映射保存路径（默认 <out>.map.json）")
+    s.add_argument("--report", help="把完整风险报告写到此路径")
+    s.add_argument("--format", choices=formats.FORMATS, default=formats.AUTO,
+                   help="plain / markdown(保护 ``` 代码块) / auto(默认)")
+    s.add_argument("--policy", choices=POLICIES, help="脱敏力度，默认 balanced")
+    s.add_argument("--locale", help="zh(默认) / en：正则集与 prompt 语言")
+    s.add_argument("--no-llm", action="store_true", help="仅正则、不调模型（离线兜底）")
+    s.add_argument("--no-critic", action="store_true", help="跳过残留风险复审")
+    s.add_argument("--one-way", action="store_true",
+                   help="单向发布模式：不产出反向映射（不可回填）")
+    s.add_argument("--allow-degraded", action="store_true",
+                   help="即便模型层未运行也以 0 退出")
+    _add_model_flags(s)
+    s.set_defaults(func=cmd_scrub)
+
+    r = sub.add_parser("rehydrate", help="用映射把云端回复里的代号还原成真实身份")
+    r.add_argument("infile", help="云端回复（- 表示 stdin）")
+    r.add_argument("--map", required=True, help="scrub 时产出的 *.map.json")
+    r.add_argument("-o", "--out", help="输出路径（默认 stdout）")
+    r.set_defaults(func=cmd_rehydrate)
+
+    c = sub.add_parser("clip", help="脱敏剪贴板内容（贴进 ChatGPT 前一键洗）")
+    c.add_argument("--rehydrate", action="store_true",
+                   help="反向：用映射把剪贴板里的代号还原成真实身份")
+    c.add_argument("--map", help=f"映射路径（默认 {DEFAULT_CLIP_MAP}）")
+    c.add_argument("--format", choices=formats.FORMATS, default=formats.AUTO)
+    c.add_argument("--policy", choices=POLICIES, help="脱敏力度，默认 balanced")
+    c.add_argument("--locale", help="zh(默认) / en")
+    c.add_argument("--no-llm", action="store_true", help="仅正则、不调模型")
+    c.add_argument("--no-critic", action="store_true", help="跳过残留风险复审")
+    c.add_argument("--one-way", action="store_true", help="不产出反向映射")
+    c.add_argument("--allow-degraded", action="store_true",
+                   help="即便模型层未运行也以 0 退出")
+    _add_model_flags(c)
+    c.set_defaults(func=cmd_clip)
+
+    sub.add_parser("providers", help="列出可用大模型 provider").set_defaults(
+        func=cmd_providers)
+
+    m = sub.add_parser("models", help="列出本地 Ollama 已安装模型")
+    m.add_argument("--endpoint", help="Ollama 端点，默认 http://localhost:11434")
+    m.set_defaults(func=cmd_models)
+
+    sub.add_parser("version", help="打印版本").set_defaults(func=cmd_version)
+    return parser
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    args = build_parser().parse_args(argv)
+    try:
+        return args.func(args)
+    except (OSError, ValueError) as exc:
+        print(f"✗ {exc}", file=sys.stderr)
+        return EXIT_ERROR
+
+
+if __name__ == "__main__":
+    sys.exit(main())

vaultengine/clipboard.py

@@ -0,0 +1,71 @@
+"""Cross-platform clipboard access through OS tools — stdlib subprocess, no deps.
+
+macOS uses pbcopy/pbpaste; Windows uses PowerShell; Linux uses wl-clipboard,
+xclip, or xsel (whichever is installed). Raises :class:`ClipboardError` with an
+actionable message when no tool is available.
+"""
+
+from __future__ import annotations
+
+import shutil
+import subprocess
+import sys
+
+_TIMEOUT = 10
+
+
+class ClipboardError(Exception):
+    pass
+
+
+def _candidates():
+    """Return (paste_commands, copy_commands) for the current platform."""
+    if sys.platform == "darwin":
+        return [["pbpaste"]], [["pbcopy"]]
+    if sys.platform == "win32":
+        return ([["powershell", "-noprofile", "-command", "Get-Clipboard"]],
+                [["powershell", "-noprofile", "-command", "$input | Set-Clipboard"]])
+    paste, copy = [], []
+    if shutil.which("wl-paste"):
+        paste.append(["wl-paste", "--no-newline"])
+    if shutil.which("wl-copy"):
+        copy.append(["wl-copy"])
+    if shutil.which("xclip"):
+        paste.append(["xclip", "-selection", "clipboard", "-o"])
+        copy.append(["xclip", "-selection", "clipboard", "-i"])
+    if shutil.which("xsel"):
+        paste.append(["xsel", "--clipboard", "--output"])
+        copy.append(["xsel", "--clipboard", "--input"])
+    return paste, copy
+
+
+def _pick(cmds):
+    for cmd in cmds:
+        if sys.platform == "win32" or shutil.which(cmd[0]):
+            return cmd
+    return None
+
+
+def read_clipboard() -> str:
+    cmd = _pick(_candidates()[0])
+    if not cmd:
+        raise ClipboardError(
+            "找不到剪贴板读取工具（macOS: pbpaste；Linux: 装 xclip / xsel / wl-clipboard）")
+    try:
+        proc = subprocess.run(cmd, capture_output=True, encoding="utf-8",
+                              timeout=_TIMEOUT)
+    except (OSError, subprocess.SubprocessError) as exc:
+        raise ClipboardError(f"读取剪贴板失败：{exc}") from exc
+    return proc.stdout or ""
+
+
+def write_clipboard(text: str) -> None:
+    cmd = _pick(_candidates()[1])
+    if not cmd:
+        raise ClipboardError(
+            "找不到剪贴板写入工具（macOS: pbcopy；Linux: 装 xclip / xsel / wl-clipboard）")
+    try:
+        subprocess.run(cmd, input=text, encoding="utf-8", timeout=_TIMEOUT,
+                       check=True)
+    except (OSError, subprocess.SubprocessError) as exc:
+        raise ClipboardError(f"写入剪贴板失败：{exc}") from exc

vaultengine/config.py

@@ -0,0 +1,96 @@
+"""Runtime configuration.
+
+Resolution order (later wins):
+    dataclass defaults  <  JSON config file  <  VAULT_* environment  <  CLI flags
+
+Everything is plain stdlib so the package stays dependency-free.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import asdict, dataclass, fields
+from typing import Any, Dict, Optional
+
+# Redaction policies — see README "脱敏力度 / aggressiveness".
+POLICY_BALANCED = "balanced"  # typed reversible tokens; coarse structure kept
+POLICY_MAX = "max"            # opaque tokens, dates coarsened — type hidden too
+POLICY_LIGHT = "light"        # only persons + explicit PII; context untouched
+POLICIES = (POLICY_BALANCED, POLICY_MAX, POLICY_LIGHT)
+
+DEFAULT_MODEL = "qwen3.6:27b"   # the de-identification model, runtime-swappable
+DEFAULT_PROVIDER = "ollama"
+
+
+@dataclass
+class Config:
+    provider: str = DEFAULT_PROVIDER
+    model: str = DEFAULT_MODEL
+    endpoint: str = ""          # provider default used when empty
+    api_key: str = ""           # only for remote (openai-compat) providers
+    policy: str = POLICY_BALANCED
+    locale: str = "zh"          # selects regex sets / prompt language
+
+    use_llm: bool = True        # False => deterministic detectors only (offline)
+    critic: bool = True         # residual-risk second LLM pass
+
+    num_ctx: int = 8192         # context window hint passed to local models
+    timeout: int = 300          # seconds per model call
+    chunk_chars: int = 6000     # long inputs are split for detection
+    chunk_overlap: int = 400    # overlap so entities on a boundary aren't lost
+
+    def validate(self) -> "Config":
+        if self.policy not in POLICIES:
+            raise ValueError(
+                f"未知 policy {self.policy!r}（可选 {POLICIES}）")
+        if self.chunk_overlap >= self.chunk_chars:
+            raise ValueError("chunk_overlap 必须小于 chunk_chars")
+        return self
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+    # -- layered construction ------------------------------------------------
+    @classmethod
+    def _field_names(cls) -> set:
+        return {f.name for f in fields(cls)}
+
+    @classmethod
+    def load(cls, path: Optional[str] = None,
+             overrides: Optional[Dict[str, Any]] = None) -> "Config":
+        """Build a Config from (optional) file + environment + overrides."""
+        data: Dict[str, Any] = {}
+
+        path = path or os.environ.get("VAULT_CONFIG")
+        if path and os.path.isfile(path):
+            with open(path, encoding="utf-8") as fh:
+                file_data = json.load(fh)
+            data.update({k: v for k, v in file_data.items()
+                         if k in cls._field_names()})
+
+        data.update(cls._from_env())
+
+        if overrides:
+            data.update({k: v for k, v in overrides.items()
+                         if k in cls._field_names() and v is not None})
+
+        return cls(**data).validate()
+
+    @classmethod
+    def _from_env(cls) -> Dict[str, Any]:
+        names = cls._field_names()
+        bool_fields = {"use_llm", "critic"}
+        int_fields = {"num_ctx", "timeout", "chunk_chars", "chunk_overlap"}
+        out: Dict[str, Any] = {}
+        for name in names:
+            raw = os.environ.get("VAULT_" + name.upper())
+            if raw is None:
+                continue
+            if name in bool_fields:
+                out[name] = raw.strip().lower() in ("1", "true", "yes", "on")
+            elif name in int_fields:
+                out[name] = int(raw)
+            else:
+                out[name] = raw
+        return out

vaultengine/detectors.py

@@ -0,0 +1,88 @@
+"""Deterministic detectors — the always-on, offline, no-LLM layer.
+
+These catch structured PII that an LLM may paraphrase past or silently miss:
+emails, phone numbers, URLs, handles, IP/MAC, government IDs, payment cards
+(Luhn-checked), crypto addresses. They run with or without a model and give the
+pipeline a reliable floor of coverage.
+
+Every detector returns offset-anchored :class:`~vaultengine.spans.Span` objects,
+so overlaps are resolved deterministically by ``spans.merge_overlapping``.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import List
+
+from .spans import (CAT_CONTACT, CAT_ID, Span)
+
+# (sub-kind, category, compiled pattern). Order is informational only — overlap
+# resolution is by offset/length, not declaration order.
+_PATTERNS = [
+    ("email", CAT_CONTACT,
+     re.compile(r"[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}")),
+    ("url", CAT_CONTACT,
+     re.compile(r"(?:https?://|www\.)[^\s<>()\[\]{}　，。]+",
+                re.IGNORECASE)),
+    ("cn-mobile", CAT_CONTACT, re.compile(r"(?<!\d)1[3-9]\d{9}(?!\d)")),
+    ("phone-e164", CAT_CONTACT,
+     re.compile(r"(?<![\w+])\+\d{1,3}[\s\-]?\d[\d\s\-]{6,}\d(?!\d)")),
+    ("handle", CAT_CONTACT, re.compile(r"(?<![\w@./])@[A-Za-z0-9_]{2,30}\b")),
+    ("ipv4", CAT_ID,
+     re.compile(r"(?<![\d.])(?:(?:25[0-5]|2[0-4]\d|1?\d?\d)\.){3}"
+                r"(?:25[0-5]|2[0-4]\d|1?\d?\d)(?![\d.])")),
+    ("ipv6", CAT_ID,
+     re.compile(r"(?<![:\w])(?:[A-Fa-f0-9]{1,4}:){2,7}[A-Fa-f0-9]{1,4}(?![:\w])")),
+    ("mac", CAT_ID,
+     re.compile(r"(?<![\w:])(?:[0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}(?![\w:])")),
+    ("cn-id", CAT_ID, re.compile(r"(?<!\d)\d{17}[\dXx](?!\d)")),
+    ("crypto-eth", CAT_ID, re.compile(r"\b0x[a-fA-F0-9]{40}\b")),
+    ("crypto-btc", CAT_ID,
+     re.compile(r"\b(?:bc1[a-z0-9]{25,39}|[13][a-km-zA-HJ-NP-Z1-9]{25,34})\b")),
+]
+
+# Payment cards need a checksum, so they get a dedicated pass.
+_CARD_CANDIDATE = re.compile(r"(?<![\d\-])(?:\d[ \-]?){13,19}(?![\d\-])")
+# Bare long digit runs (account-ish). Low confidence: only tokenized when the
+# policy already redacts the `id` category.
+_LONG_DIGITS = re.compile(r"(?<!\d)\d{9,}(?!\d)")
+
+
+def _luhn_ok(digits: str) -> bool:
+    total, alt = 0, False
+    for ch in reversed(digits):
+        d = ord(ch) - 48
+        if alt:
+            d *= 2
+            if d > 9:
+                d -= 9
+        total += d
+        alt = not alt
+    return total % 10 == 0
+
+
+def detect(text: str, locale: str = "zh") -> List[Span]:
+    """Run every deterministic detector over ``text``; return anchored Spans."""
+    if not text:
+        return []
+    spans: List[Span] = []
+
+    for kind, category, pattern in _PATTERNS:
+        for m in pattern.finditer(text):
+            spans.append(Span(surface=m.group(0), category=category,
+                              source=f"regex:{kind}", confidence=1.0,
+                              start=m.start(), end=m.end(), note=kind))
+
+    for m in _CARD_CANDIDATE.finditer(text):
+        digits = re.sub(r"[ \-]", "", m.group(0))
+        if 13 <= len(digits) <= 19 and _luhn_ok(digits):
+            spans.append(Span(surface=m.group(0), category=CAT_ID,
+                              source="regex:card", confidence=1.0,
+                              start=m.start(), end=m.end(), note="card"))
+
+    for m in _LONG_DIGITS.finditer(text):
+        spans.append(Span(surface=m.group(0), category=CAT_ID,
+                          source="regex:digits", confidence=0.5,
+                          start=m.start(), end=m.end(), note="digits"))
+
+    return spans

vaultengine/formats.py

@@ -0,0 +1,59 @@
+"""Input formats — decide which regions of a document get scrubbed.
+
+`plain` scrubs everything. `markdown` scrubs the prose but leaves fenced
+``` code blocks untouched — so an embedded JSON schema, code sample, or output
+template you include for the model is preserved verbatim instead of being
+mangled by pseudonymization.
+
+A "segment plan" is a list of ``(kind, content)`` where kind is ``'keep'`` or
+``'scrub'``; the concatenation of all contents always reproduces the input
+exactly.
+"""
+
+from __future__ import annotations
+
+from typing import List, Tuple
+
+Segment = Tuple[str, str]
+
+PLAIN = "plain"
+MARKDOWN = "markdown"
+AUTO = "auto"
+FORMATS = (PLAIN, MARKDOWN, AUTO)
+
+_FENCE = "```"
+
+
+def detect_format(text: str) -> str:
+    """A fenced code block is the signal to switch on `markdown` protection."""
+    return MARKDOWN if (text and _FENCE in text) else PLAIN
+
+
+def _split_fences(s: str) -> List[Segment]:
+    """Keep fenced ``` blocks verbatim; everything else is scrubbable."""
+    out: List[Segment] = []
+    i = 0
+    while True:
+        start = s.find(_FENCE, i)
+        if start == -1:
+            if s[i:]:
+                out.append(("scrub", s[i:]))
+            break
+        if s[i:start]:
+            out.append(("scrub", s[i:start]))
+        end = s.find(_FENCE, start + 3)
+        if end == -1:                      # unterminated fence: keep the rest
+            out.append(("keep", s[start:]))
+            break
+        end += 3
+        out.append(("keep", s[start:end]))
+        i = end
+    return out
+
+
+def segment(text: str, fmt: str = AUTO) -> List[Segment]:
+    if fmt == AUTO:
+        fmt = detect_format(text)
+    if fmt == MARKDOWN:
+        return _split_fences(text)
+    return [("scrub", text)]