From fa5523a37d189c69511a10850ccb5c666ad5bf3f Mon Sep 17 00:00:00 2001 From: yunmei Date: Tue, 4 Aug 2026 17:15:02 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=AE=8C=E6=88=90=20issue=20#48=20DLP?= =?UTF-8?q?=20=E6=95=8F=E6=84=9F=E6=95=B0=E6=8D=AE=E6=8B=A6=E6=88=AA?= =?UTF-8?q?=EF=BC=88=E7=9B=AE=E6=A0=87=20100%=20=E6=8B=A6=E6=88=AA?= =?UTF-8?q?=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- core/llm-gateway/README.md | 76 ++++ core/llm-gateway/__init__.py | 34 ++ core/llm-gateway/config/dlp.template.yaml | 51 +++ core/llm-gateway/dlp.py | 446 ++++++++++++++++++++++ core/llm-gateway/tests/_bootstrap.py | 16 + core/llm-gateway/tests/test_dlp.py | 215 +++++++++++ 6 files changed, 838 insertions(+) create mode 100644 core/llm-gateway/README.md create mode 100644 core/llm-gateway/__init__.py create mode 100644 core/llm-gateway/config/dlp.template.yaml create mode 100644 core/llm-gateway/dlp.py create mode 100644 core/llm-gateway/tests/_bootstrap.py create mode 100644 core/llm-gateway/tests/test_dlp.py diff --git a/core/llm-gateway/README.md b/core/llm-gateway/README.md new file mode 100644 index 0000000..2625af3 --- /dev/null +++ b/core/llm-gateway/README.md @@ -0,0 +1,76 @@ +# iAOP-Core · LLM 网关（LLM Gateway） + +对应 PRD 5.4「④ LLM 网关 + RAG」与 EPIC #6（Issue #48 等子任务）： +本地 70B（敏感/核心）+ 云端 API（脱敏/通用）**混合**，安全分级路由， +**敏感数据本地闭环，仅脱敏/公开内容可走云端 API**（数据不出厂）。 + +## 模块结构 + +``` +core/llm-gateway/ +├── __init__.py 包入口（导出 DLP 引擎 API） +├── dlp.py DLP 敏感数据拦截引擎（Issue #48） +├── config/ +│ └── dlp.template.yaml 模板 DLP 规则资产（ti-cl4 示例，换行业只改它） +└── tests/ + ├── _bootstrap.py 测试引导（目录含连字符，挂载包名 llm_gateway） + └── test_dlp.py DLP 引擎单元测试 +``` + +## DLP 敏感数据拦截（Issue #48） + +出站防线：任何要发往**云端**的内容（用户 query / RAG 检索到的 context / +模型输出）先经 `DlpEngine.check_outbound()` 检查，**命中任一 block 规则 +即整包拦截**（fail-closed），并落结构化审计日志（全量）。 + +```python +from llm_gateway.dlp import DlpEngine + +engine = DlpEngine.from_template_config("config/dlp.template.yaml") +result = engine.check_outbound({ + "query": user_query, # 用户提问 + "context": rag_context, # RAG 命中文档片段 + "output": llm_answer, # 模型生成结果 +}) +if result.blocked: + # 不发往云端：转本地处理 / 转人工确认 / 拒绝（PRD 5.4 异常时转人工） + engine.drain_audit() # 取走审计记录（对接外部审计管道） +``` + +### 语义（目标 DLP 拦截率 100%） + +- **命中即拦截**：`decision="block"`、`reason="hit"`； +- **未配置即拦截**：引擎无任何规则（含内置保底）且 `fail_closed=True` 时 + `reason="no_rules"`——未配置 = 不安全，保守拒绝出站； +- **审计不落明文敏感内容**：命中明细只记录规则名 / 类别 / 所属 part / + 位置与 `` 脱敏占位。 + +### 规则来源 + +1. **内置保底**（`DLP_DEFAULT_RULES`，内核自带）：通用 PII 与凭证 + （身份证、手机号、邮箱、IPv4、云访问密钥 AK、密钥键值对）+ 示例工艺词； +2. **模板资产**（`config/dlp.template.yaml`）：行业敏感资产（工艺参数 / + 配方 / 关键设备 / 人员岗位），同名规则覆盖内置实现模板定制。 + +规则两种匹配方式：`keyword`（大小写不敏感子串）与 `regex`（正则）。 + +### 验收对照（PRD 5.4 / 父 Issue #6） + +| 验收项 | 实现 | +|--------|------| +| DLP 拦截率 100% | fail-closed：命中即 block；未配置规则也 block | +| 敏感数据本地闭环 | 出站（云端通道）必经 check_outbound 检查 | +| 敏感词/DLP 规则为配置点 | `config/dlp.template.yaml`，换行业只改资产 | +| 审计日志全量（NFR 安全-数据） | 每次检查落一条结构化审计，可 drain 对接管道 | + +## 运行测试 + +```bash +cd core/llm-gateway +python -m unittest discover -s tests -v +``` + +## 后续子任务（EPIC #6 拆分，待扩展） + +- 敏感度路由（准确率 ≥ 96.5%）与路由准确率评估脚本（Issue #49）； +- Prompt 版本管理与幻觉/事实性校验中间件（Issue #47）。 diff --git a/core/llm-gateway/__init__.py b/core/llm-gateway/__init__.py new file mode 100644 index 0000000..b714c0c --- /dev/null +++ b/core/llm-gateway/__init__.py @@ -0,0 +1,34 @@ +# -*- coding: utf-8 -*- +"""iAOP-Core · LLM 网关（LLM Gateway）—— 混合 LLM 的安全出站防线。 + +对应 PRD 5.4「④ LLM 网关 + RAG」与 EPIC #6： +本地 70B（敏感/核心）+ 云端 API（脱敏/通用）混合，敏感数据**本地闭环**。 + +当前子模块（Issue #48）： +- dlp DLP 敏感数据拦截引擎：出站内容（query / RAG context / 模型输出） + 发往云端前做敏感规则检查，命中即拦截（目标 100% 拦截），全量审计。 + +后续子任务（EPIC #6 拆分，将在本包扩展）： +- 敏感度路由（router）、Prompt 版本管理与幻觉校验中间件、路由准确率评估。 + +测试：`python -m unittest discover -s tests -v`（在 core/llm-gateway 目录下执行）。 +""" +__version__ = "0.1.0" + +from .dlp import ( + DLP_DEFAULT_RULES, + DlpEngine, + DlpHit, + DlpResult, + DlpRule, + DlpRuleKind, +) + +__all__ = [ + "DlpRuleKind", + "DlpRule", + "DlpHit", + "DlpResult", + "DlpEngine", + "DLP_DEFAULT_RULES", +] diff --git a/core/llm-gateway/config/dlp.template.yaml b/core/llm-gateway/config/dlp.template.yaml new file mode 100644 index 0000000..0629952 --- /dev/null +++ b/core/llm-gateway/config/dlp.template.yaml @@ -0,0 +1,51 @@ +# -*- coding: utf-8 -*- +# 模板 DLP 规则资产示例：ti-cl4（氯化车间/海绵钛，Template-Ti 一期）。 +# +# 说明： +# - 这是「敏感词/DLP 规则」配置点（PRD 5.4）：换行业只改本文件，内核零改动； +# - mode: fail-closed —— 命中任一规则即拦截出站（目标 DLP 拦截率 100%）； +# - kind: keyword 大小写不敏感子串匹配；regex 正则匹配； +# - 通用 PII/凭证规则由内核内置保底（DLP_DEFAULT_RULES），无需重复配置； +# 本文件只补充**行业敏感资产**（工艺参数 / 配方 / 关键设备 / 人员等）。 +template: ti-cl4 +version: 1.0.0 +mode: fail-closed +rules: + # ---- 工艺敏感参数（Ti-1 氯化车间） ---- + - name: proc_cl2_flow + category: process-parameter + kind: keyword + pattern: 氯气流量 + description: 氯气流量参数（工艺敏感） + - name: proc_furnace_temp + category: process-parameter + kind: keyword + pattern: 炉温 + description: 炉温参数（工艺敏感） + - name: proc_feeding_ratio + category: process-parameter + kind: keyword + pattern: 加料比 + description: 加料配比参数（工艺敏感） + - name: proc_ti_purity + category: process-parameter + kind: keyword + pattern: 钛纯度 + description: 产品质量指标（钛纯度） + - name: proc_formula_detail + category: process-parameter + kind: regex + pattern: '配比\s*[:：]?\s*[\d.]+%?' + description: 配方配比数值（结构化敏感数据） + # ---- 设备 / 工艺关键点（内部敏感命名） ---- + - name: proc_furnace_id + category: process-asset + kind: regex + pattern: '沸腾氯化炉\s*[A-Z]?\d+' + description: 关键设备编号（沸腾氯化炉） + # ---- 人员 / 排班（内部敏感） ---- + - name: staff_contact + category: personnel + kind: keyword + pattern: 值班长 + description: 岗位/人员信息（内部敏感） diff --git a/core/llm-gateway/dlp.py b/core/llm-gateway/dlp.py new file mode 100644 index 0000000..dd9963f --- /dev/null +++ b/core/llm-gateway/dlp.py @@ -0,0 +1,446 @@ +# -*- coding: utf-8 -*- +"""iAOP-Core · LLM 网关 —— DLP 敏感数据拦截（Issue #48，EPIC #6 子任务）。 + +对应 PRD 5.4「④ LLM 网关 + RAG」： +- 数据不出厂：敏感数据**本地闭环**，仅脱敏/公开内容可走云端 API； +- 验收：**DLP 拦截率 100%**，敏感词/DLP 规则作为模板配置点； +- 用户操作流程：用户提问 → 路由判断敏感级 → 本地/云端生成 → + RAG 溯源校验 → 返回答案（出站前经 DLP 检查）。 + +本模块实现 **出站 DLP 防线**：任何要发往**云端**的内容 +（用户 query / RAG 检索到的 context / 模型输出）先经 `DlpEngine.check_outbound` +检查，命中任一 block 规则即**整包拦截**（fail-closed，目标 100% 拦截）， +并落**结构化审计日志**（全量，PRD NFR 安全-数据）。 + +规则来源： +- 模板资产 `config/dlp.template.yaml`（行业敏感词 + 通用 PII 正则）， + 换行业只改该资产，内核零改动（对齐 rag-kb 模板化思想）； +- `DLP_DEFAULT_RULES` 内置保底规则：即使未加载任何配置， + 通用 PII（身份证/手机号/邮箱/IP/访问密钥等）也默认生效。 + +测试：`python -m unittest discover -s tests -v`（在 core/llm-gateway 目录下执行）。 +""" +from __future__ import annotations + +import re +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Callable, Dict, List, Optional, Tuple + +# --------------------------------------------------------------------------- +# 轻量 YAML 子集解析（零第三方依赖，递归下降）：map / list / scalar / 注释。 +# 与 rag-kb/templating.py 同款（模块内自持一份，保持模块零耦合）。 +# 足以解析 `config/dlp.template.yaml` 这类模板资产。 +# --------------------------------------------------------------------------- + + +def _parse_scalar(text: str) -> str: + """去掉标量两侧引号与行内注释（`key: value # comment`）。""" + t = text.split(" #", 1)[0].strip() + if len(t) >= 2 and t[0] == t[-1] and t[0] in ("'", '"'): + return t[1:-1] + return t + + +def _strip_comments(lines: List[str]) -> List[Tuple[str, int]]: + """剔除空行与整行注释，保留行号（1 起）用于报错定位。""" + out = [] + for i, ln in enumerate(lines): + s = ln.strip() + if not s or s.startswith("#"): + continue + out.append((ln, i + 1)) + return out + + +def _parse_node(lines: List[Tuple[str, int]], i: int, indent: int): + """递归解析从 lines[i] 开始、缩进为 `indent` 的一个节点。 + + 返回 `(value, next_i)`：value 为 dict / list / str，next_i 为下一个 + 未消费行的下标。 + """ + text, no = lines[i] + lead = len(text) - len(text.lstrip(" ")) + + # ---- list 节点：`- item` 或 `- key: val`（map 项） ---- + if text.lstrip(" ").startswith("- "): + items: List[object] = [] + while i < len(lines): + t, no2 = lines[i] + stripped = t.lstrip(" ") + if not stripped.startswith("- "): + break + lead_j = len(t) - len(t.lstrip(" ")) + if lead_j != indent: + break + item_text = stripped[2:].strip() + if not item_text: + raise ValueError(f"dlp.yaml 第 {no2} 行：list 项为空") + if ":" in item_text: + map_indent = len(t) - len(t.lstrip(" ")) + 2 + lines[i] = (" " * map_indent + item_text, no2) + v, i = _parse_node(lines, i, map_indent) + items.append(v) + else: + items.append(_parse_scalar(item_text)) + i += 1 + return items, i + + # ---- map 节点：`key: value` / `key:`（嵌套值） ---- + result: Dict[str, object] = {} + while i < len(lines): + t, no = lines[i] + lead_j = len(t) - len(t.lstrip(" ")) + if lead_j < indent or t.lstrip(" ").startswith("- "): + break + if lead_j > indent: + raise ValueError(f"dlp.yaml 第 {no} 行缩进异常（期望 {indent}，实际 {lead_j}）") + if ":" not in t: + raise ValueError(f"dlp.yaml 第 {no} 行不是合法键值对：{t!r}") + key, _, rest = t.partition(":") + key = key.strip() + rest = rest.strip() + if rest: + result[key] = _parse_scalar(rest) + i += 1 + continue + if i + 1 >= len(lines): + raise ValueError(f"dlp.yaml 第 {no} 行 {key!r} 缺少值") + sub_indent = len(lines[i + 1][0]) - len(lines[i + 1][0].lstrip(" ")) + if sub_indent <= indent: + raise ValueError(f"dlp.yaml 第 {no} 行 {key!r} 缺少值（无嵌套内容）") + v, i = _parse_node(lines, i + 1, sub_indent) + result[key] = v + return result, i + + +def _load_yaml_text(text: str) -> Dict[str, object]: + """解析 YAML 子集 → 嵌套 dict/list。顶层必须为 map。""" + lines = _strip_comments(text.splitlines()) + if not lines: + return {} + top_indent = len(lines[0][0]) - len(lines[0][0].lstrip(" ")) + value, next_i = _parse_node(lines, 0, top_indent) + if not isinstance(value, dict): + raise ValueError("dlp.yaml 顶层必须是 map") + if next_i < len(lines): + raise ValueError( + f"dlp.yaml 第 {lines[next_i][1]} 行：顶层存在多个节点（缩进不一致）" + ) + return value + + +# --------------------------------------------------------------------------- +# 规则模型 +# --------------------------------------------------------------------------- + + +class DlpRuleKind: + """规则匹配方式：`keyword` 子串匹配 / `regex` 正则匹配。""" + + KEYWORD = "keyword" + REGEX = "regex" + + +@dataclass(frozen=True) +class DlpRule: + """一条 DLP 敏感数据规则（来自模板资产或内置保底规则）。 + + - `name` 规则唯一名（审计用）； + - `category` 敏感类别（pii / process-keyword / credential 等）； + - `kind` keyword 子串匹配 或 regex 正则匹配； + - `pattern` 关键词原文，或正则表达式； + - `description` 规则说明（人读）。 + """ + + name: str + category: str + pattern: str + kind: str = DlpRuleKind.KEYWORD + description: str = "" + + def _compiled(self) -> "re.Pattern[str]": + if self.kind == DlpRuleKind.REGEX: + return re.compile(self.pattern, re.IGNORECASE) + return re.compile(re.escape(self.pattern), re.IGNORECASE) + + def find(self, text: str) -> List[Tuple[str, int, int]]: + """返回文本中所有命中片段 `(matched, start, end)`（不重叠）。""" + out: List[Tuple[str, int, int]] = [] + if not text: + return out + if self.kind == DlpRuleKind.REGEX: + for m in self._compiled().finditer(text): + out.append((m.group(0), m.start(), m.end())) + else: + # 大小写不敏感子串匹配（中文无大小写，直接命中） + lower = text.lower() + needle = self.pattern.lower() + start = 0 + while True: + pos = lower.find(needle, start) + if pos < 0: + break + out.append((text[pos:pos + len(self.pattern)], pos, pos + len(self.pattern))) + start = pos + max(len(self.pattern), 1) + return out + + +@dataclass(frozen=True) +class DlpHit: + """一次敏感命中：哪条规则、哪段文本、位于何处（可选所属出站 part）。""" + + rule: str + category: str + matched: str + start: int + end: int + part: str = "" + + def to_dict(self) -> Dict[str, object]: + # 审计只记录脱敏预览（`` 占位）与所属 part，不落明文敏感内容 + return { + "rule": self.rule, + "category": self.category, + "part": self.part, + "matched": f"<{self.category}>", + "span": [self.start, self.end], + } + + +@dataclass(frozen=True) +class DlpResult: + """一次出站检查的结论。 + + - `decision` allow（放行）| block（拦截）； + - `reason` 拦截原因（`hit` 命中敏感规则 / `no_rules` 未配置规则且 fail-closed）； + - `hits` 命中明细； + - `audit` 本次检查的结构化审计记录（JSON 序列化友好）。 + """ + + decision: str + reason: str + hits: List[DlpHit] + audit: Dict[str, object] + + @property + def blocked(self) -> bool: + return self.decision == "block" + + +# --------------------------------------------------------------------------- +# 内置保底规则：未加载任何模板配置时，通用 PII 也默认生效（目标 100% 拦截） +# --------------------------------------------------------------------------- + +DLP_DEFAULT_RULES: List[DlpRule] = [ + # ---- PII / 凭证（通用，任何行业模板都适用） ---- + DlpRule(name="pii_id_card", category="pii", + kind=DlpRuleKind.REGEX, pattern=r"\b\d{17}[\dXx]\b", + description="身份证号（18 位）"), + DlpRule(name="pii_mobile", category="pii", + kind=DlpRuleKind.REGEX, pattern=r"\b1[3-9]\d{9}\b", + description="中国大陆手机号"), + DlpRule(name="pii_email", category="pii", + kind=DlpRuleKind.REGEX, pattern=r"[\w.+-]+@[\w-]+\.[\w.-]+", + description="邮箱地址"), + DlpRule(name="pii_ipv4", category="pii", + kind=DlpRuleKind.REGEX, pattern=r"\b(?:\d{1,3}\.){3}\d{1,3}\b", + description="IPv4 地址"), + DlpRule(name="credential_ak", category="credential", + kind=DlpRuleKind.REGEX, pattern=r"\b(?:AKIA|LTAI)[0-9A-Z]{16,}\b", + description="云访问密钥 ID（AWS AKIA / 阿里云 LTAI 样式）"), + DlpRule(name="credential_kv", category="credential", + kind=DlpRuleKind.REGEX, + pattern=r"(?:access[_-]?key|secret|password|token)\s*[:=]\s*\S+", + description="密钥/口令键值对"), + # ---- 工艺敏感词示例（行业模板会按需增删，见 config/dlp.template.yaml） ---- + DlpRule(name="process_formula", category="process-keyword", + kind=DlpRuleKind.KEYWORD, pattern="配方", + description="配方类敏感词"), +] + + +class DlpEngine: + """出站 DLP 检查引擎（fail-closed：命中即拦截，未配置规则即拦截）。 + + 用法： + ```python + engine = DlpEngine.from_template_config("config/dlp.template.yaml") + result = engine.check_outbound({ + "query": user_query, + "context": rag_context, + "output": llm_answer, + }) + if result.blocked: + # 不发往云端：转本地处理 / 转人工 / 拒绝（PRD 5.4） + ... + ``` + """ + + def __init__(self, + rules: Optional[List[DlpRule]] = None, + fail_closed: bool = True, + template: str = "builtin", + version: str = "1.0.0", + include_builtins: bool = True, + audit_hook: Optional[Callable[[Dict[str, object]], None]] = None): + # 内置保底规则 + 模板规则（模板同名规则覆盖内置，实现模板定制） + merged: Dict[str, DlpRule] = {} + if include_builtins: + merged.update({r.name: r for r in DLP_DEFAULT_RULES}) + for r in rules or []: + merged[r.name] = r + self._rules: List[DlpRule] = list(merged.values()) + self.fail_closed = fail_closed + self.template = template + self.version = version + self.audit_hook = audit_hook + self._audit_buffer: List[Dict[str, object]] = [] + + # ------------------------------------------------------------------ + # 构造 + # ------------------------------------------------------------------ + + @classmethod + def from_template_config(cls, path: str, + audit_hook: Optional[Callable[[Dict[str, object]], None]] = None + ) -> "DlpEngine": + """从模板 DLP 规则资产加载引擎（期望结构见 `config/dlp.template.yaml`）。""" + with open(path, "r", encoding="utf-8") as fh: + data = _load_yaml_text(fh.read()) + + template = str(data.get("template", "")).strip() + version = str(data.get("version", "1.0.0")).strip() or "1.0.0" + mode = str(data.get("mode", "fail-closed")).strip().lower() + fail_closed = mode != "fail-open" + + rules: List[DlpRule] = [] + raw_rules = data.get("rules") or [] + if not isinstance(raw_rules, list): + raise ValueError("dlp.yaml rules 必须是 list") + for item in raw_rules: + if not isinstance(item, dict): + raise ValueError(f"dlp.yaml rules 项必须是 map，实际 {item!r}") + name = str(item.get("name", "")).strip() + if not name: + raise ValueError("dlp.yaml 规则缺少 name") + kind = str(item.get("kind", "keyword")).strip().lower() + if kind not in (DlpRuleKind.KEYWORD, DlpRuleKind.REGEX): + raise ValueError(f"dlp.yaml 规则 {name!r} 的 kind 非法：{kind!r}") + pattern = str(item.get("pattern", "")).strip() + if not pattern: + raise ValueError(f"dlp.yaml 规则 {name!r} 缺少 pattern") + rules.append(DlpRule( + name=name, + category=str(item.get("category", "custom")).strip() or "custom", + kind=kind, + pattern=pattern, + description=str(item.get("description", "")).strip(), + )) + return cls(rules=rules, fail_closed=fail_closed, + template=template, version=version, audit_hook=audit_hook) + # ------------------------------------------------------------------ + # 检查 + # ------------------------------------------------------------------ + + def inspect(self, text: str) -> List[DlpHit]: + """对单段文本做敏感规则扫描，返回全部命中（按文本位置排序）。""" + hits: List[DlpHit] = [] + for rule in self._rules: + for matched, start, end in rule.find(text or ""): + hits.append(DlpHit( + rule=rule.name, category=rule.category, + matched=matched, start=start, end=end, + )) + hits.sort(key=lambda h: (h.start, h.end)) + return hits + + def mask(self, text: str) -> str: + """把命中片段替换为 `` 占位（供"脱敏后可出站"的后续流程）。 + + 注意：本任务（Issue #48）的出站策略为**命中即拦截**，不自动放行 + 脱敏内容；mask 仅作为审计预览与后续脱敏通道的辅助能力。 + """ + hits = self.inspect(text) + if not hits: + return text + out = [] + prev = 0 + for h in hits: + out.append(text[prev:h.start]) + out.append(f"<{h.category}>") + prev = h.end + out.append(text[prev:]) + return "".join(out) + + def check_outbound(self, parts: Dict[str, str]) -> DlpResult: + """出站检查：对多段内容（query / context / output 等）统一扫描。 + + **fail-closed 语义（目标 100% 拦截）**： + - 命中任一规则 → block（`reason="hit"`）； + - 引擎无任何规则（连内置保底都没有）且 fail_closed → block + （`reason="no_rules"`，未配置 = 不安全，保守拒绝出站）。 + 每检查一次落一条结构化审计记录（含命中规则/类别，不落明文敏感内容）。 + """ + all_hits: List[DlpHit] = [] + for part, text in (parts or {}).items(): + for h in self.inspect(text): + # 带上所属 part，便于审计定位哪段内容触发拦截 + all_hits.append(DlpHit( + rule=h.rule, category=h.category, matched=h.matched, + start=h.start, end=h.end, part=part, + )) + all_hits.sort(key=lambda h: (h.start, h.end)) + + if all_hits: + decision, reason = "block", "hit" + elif self.fail_closed and not self._rules: + decision, reason = "block", "no_rules" + else: + decision, reason = "allow", "clean" + + audit: Dict[str, object] = { + "audit_id": uuid.uuid4().hex, + "ts": datetime.now(timezone.utc).isoformat(timespec="seconds"), + "channel": "outbound", + "template": self.template, + "dlp_version": self.version, + "decision": decision, + "reason": reason, + "parts": sorted(parts.keys()) if parts else [], + "hit_count": len(all_hits), + "hits": [h.to_dict() for h in all_hits], + } + self._audit_buffer.append(audit) + if self.audit_hook is not None: + self.audit_hook(audit) + + return DlpResult(decision=decision, reason=reason, + hits=all_hits, audit=audit) + + def drain_audit(self) -> List[Dict[str, object]]: + """取走并清空已产生的审计记录（对接外部审计管道/日志）。""" + buf, self._audit_buffer = self._audit_buffer, [] + return buf + + @property + def rule_count(self) -> int: + return len(self._rules) + + def rule_names(self) -> List[str]: + return [r.name for r in self._rules] + + def __repr__(self) -> str: # pragma: no cover - 调试辅助 + return (f"DlpEngine(template={self.template!r}, version={self.version!r}, " + f"rules={len(self._rules)}, fail_closed={self.fail_closed})") + + +__all__ = [ + "DlpRuleKind", + "DlpRule", + "DlpHit", + "DlpResult", + "DlpEngine", + "DLP_DEFAULT_RULES", +] diff --git a/core/llm-gateway/tests/_bootstrap.py b/core/llm-gateway/tests/_bootstrap.py new file mode 100644 index 0000000..e0acea0 --- /dev/null +++ b/core/llm-gateway/tests/_bootstrap.py @@ -0,0 +1,16 @@ +# -*- coding: utf-8 -*- +"""测试引导：把 `core/llm-gateway` 以包名 `llm_gateway` 挂载到 sys.modules。 + +目录名 `llm-gateway` 含连字符，无法直接以包名 import；挂载后模块内相对导入 +（`from .dlp import ...`）在 unittest 发现机制下可正常解析。 +""" +import os +import sys +import types + +LLM_GW_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, LLM_GW_DIR) +if "llm_gateway" not in sys.modules: + pkg = types.ModuleType("llm_gateway") + pkg.__path__ = [LLM_GW_DIR] + sys.modules["llm_gateway"] = pkg diff --git a/core/llm-gateway/tests/test_dlp.py b/core/llm-gateway/tests/test_dlp.py new file mode 100644 index 0000000..de8bd38 --- /dev/null +++ b/core/llm-gateway/tests/test_dlp.py @@ -0,0 +1,215 @@ +# -*- coding: utf-8 -*- +"""DLP 敏感数据拦截引擎（dlp）单元测试：规则匹配 / 出站拦截 / 审计 / 模板加载。""" +import os +import sys +import unittest + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import _bootstrap # noqa: F401 + +from llm_gateway.dlp import ( # noqa: E402 + DLP_DEFAULT_RULES, + DlpEngine, + DlpRule, + DlpRuleKind, +) + +CONFIG_PATH = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "config", "dlp.template.yaml", +) + + +class BuiltinRulesTest(unittest.TestCase): + """内置保底规则（无任何配置时通用 PII 也默认拦截）。""" + + def setUp(self): + self.engine = DlpEngine() + + def test_default_rules_present(self): + self.assertGreater(len(DLP_DEFAULT_RULES), 0) + self.assertGreaterEqual(self.engine.rule_count, len(DLP_DEFAULT_RULES)) + + def test_id_card_blocked(self): + result = self.engine.check_outbound({"query": "员工身份证 110101199003071234 入职"}) + self.assertTrue(result.blocked) + self.assertEqual(result.reason, "hit") + self.assertTrue(any(h.category == "pii" for h in result.hits)) + + def test_mobile_blocked(self): + result = self.engine.check_outbound({"query": "联系 13812345678"}) + self.assertTrue(result.blocked) + self.assertIn("pii_mobile", {h.rule for h in result.hits}) + + def test_cloud_ak_blocked(self): + result = self.engine.check_outbound({"output": "ak = AKIAIOSFODNN7EXAMPLE"}) + self.assertTrue(result.blocked) + self.assertIn("credential_ak", {h.rule for h in result.hits}) + + def test_clean_content_allowed(self): + result = self.engine.check_outbound({ + "query": "今天炉况怎么样", + "context": "炉温控制在 850-950°C 范围内。", + "output": "当前运行平稳。", + }) + self.assertFalse(result.blocked) + self.assertEqual(result.decision, "allow") + + +class KeywordRuleTest(unittest.TestCase): + def test_cn_keyword_hit(self): + engine = DlpEngine(rules=[ + DlpRule(name="proc_formula", category="process-keyword", + kind=DlpRuleKind.KEYWORD, pattern="配方"), + ]) + result = engine.check_outbound({"output": "已按配方调整投料"}) + self.assertTrue(result.blocked) + self.assertIn("proc_formula", {h.rule for h in result.hits}) + hit = next(h for h in result.hits if h.rule == "proc_formula") + self.assertEqual(hit.matched, "配方") + + def test_en_keyword_case_insensitive(self): + engine = DlpEngine(rules=[ + DlpRule(name="secret_word", category="credential", + kind=DlpRuleKind.KEYWORD, pattern="SecretKey"), + ]) + self.assertTrue(engine.check_outbound({"output": "the secretkey is hidden"}).blocked) + self.assertTrue(engine.check_outbound({"output": "the SECRETKEY is hidden"}).blocked) + + def test_no_match_allowed(self): + engine = DlpEngine(rules=[ + DlpRule(name="proc_formula", category="process-keyword", + kind=DlpRuleKind.KEYWORD, pattern="配方"), + ]) + result = engine.check_outbound({"output": "炉温正常"}) + self.assertFalse(result.blocked) + + +class RegexRuleTest(unittest.TestCase): + def test_structured_formula_ratio(self): + engine = DlpEngine(rules=[ + DlpRule(name="proc_ratio", category="process-parameter", + kind=DlpRuleKind.REGEX, pattern=r"配比\s*[:：]?\s*[\d.]+%?"), + ]) + result = engine.check_outbound({"output": "当前配比：0.35 保持不变"}) + self.assertTrue(result.blocked) + self.assertEqual(result.hits[0].matched, "配比：0.35") + + +class OutboundCheckTest(unittest.TestCase): + """出站检查：多 part 聚合 / 命中定位 / 拦截原因。""" + + def setUp(self): + self.engine = DlpEngine() + + def test_hit_in_context_blocks_whole_packet(self): + result = self.engine.check_outbound({ + "query": "请问这个员工怎么样", + "context": "档案：13900000000", + "output": "该员工表现良好。", + }) + self.assertTrue(result.blocked) + # 命中定位在 context part + self.assertEqual(result.hits[0].part, "context") + + def test_empty_parts_allowed(self): + self.assertFalse(self.engine.check_outbound({}).blocked) + self.assertFalse(self.engine.check_outbound({"query": ""}).blocked) + + +class FailClosedTest(unittest.TestCase): + def test_no_rules_fail_closed_blocks(self): + engine = DlpEngine(rules=[], include_builtins=False, fail_closed=True) + result = engine.check_outbound({"query": "任意内容"}) + self.assertTrue(result.blocked) + self.assertEqual(result.reason, "no_rules") + + def test_fail_open_without_rules_allows(self): + engine = DlpEngine(rules=[], include_builtins=False, fail_closed=False) + self.assertFalse(engine.check_outbound({"query": "任意内容"}).blocked) + + +class MaskTest(unittest.TestCase): + def test_mask_redacts_hits(self): + engine = DlpEngine() + masked = engine.mask("联系 13812345678 或 13900000000") + self.assertNotIn("13812345678", masked) + self.assertNotIn("13900000000", masked) + self.assertIn("", masked) + + def test_mask_clean_text_unchanged(self): + engine = DlpEngine() + text = "炉温正常" + self.assertEqual(engine.mask(text), text) + + +class AuditTest(unittest.TestCase): + def test_audit_record_fields(self): + engine = DlpEngine() + result = engine.check_outbound({"query": "身份证 110101199003071234"}) + audit = result.audit + self.assertEqual(audit["channel"], "outbound") + self.assertEqual(audit["decision"], "block") + self.assertEqual(audit["reason"], "hit") + self.assertGreaterEqual(audit["hit_count"], 1) + # 审计不落明文敏感内容，只记脱敏占位 + serialized = str(audit) + self.assertNotIn("110101199003071234", serialized) + self.assertIn("", serialized) + + def test_audit_buffer_and_hook(self): + seen = [] + engine = DlpEngine(audit_hook=seen.append) + engine.check_outbound({"query": "13812345678"}) + engine.check_outbound({"query": "炉温正常"}) + self.assertEqual(len(seen), 2) + self.assertEqual(len(engine.drain_audit()), 2) + self.assertEqual(engine.drain_audit(), []) + + +class TemplateConfigTest(unittest.TestCase): + """从真实模板资产（config/dlp.template.yaml）加载。""" + + def test_load_example_assets(self): + engine = DlpEngine.from_template_config(CONFIG_PATH) + self.assertEqual(engine.template, "ti-cl4") + self.assertEqual(engine.version, "1.0.0") + self.assertTrue(engine.fail_closed) + # 模板规则 + 内置保底规则 + names = set(engine.rule_names()) + self.assertIn("proc_cl2_flow", names) + self.assertIn("proc_furnace_temp", names) + self.assertIn("pii_mobile", names) # 内置保底仍在 + + def test_template_industry_keyword_blocks(self): + engine = DlpEngine.from_template_config(CONFIG_PATH) + result = engine.check_outbound({"output": "今日氯气流量为 5.2 t/h"}) + self.assertTrue(result.blocked) + self.assertIn("proc_cl2_flow", {h.rule for h in result.hits}) + + def test_template_regex_blocks(self): + engine = DlpEngine.from_template_config(CONFIG_PATH) + result = engine.check_outbound({"output": "配比：0.42 已确认"}) + self.assertTrue(result.blocked) + self.assertIn("proc_formula_detail", {h.rule for h in result.hits}) + + def test_template_rule_overrides_builtin(self): + # 同名规则（pii_mobile）在模板中自定义 pattern 后覆盖内置 + engine = DlpEngine.from_template_config(CONFIG_PATH) + result = engine.check_outbound({"query": "联系 13812345678"}) + # 内置手机号规则仍在（模板未覆盖），应照常命中 + self.assertTrue(result.blocked) + + +class TemplateOverrideTest(unittest.TestCase): + def test_same_name_overrides_builtin(self): + engine = DlpEngine(rules=[ + DlpRule(name="pii_mobile", category="pii", + kind=DlpRuleKind.REGEX, pattern=r"\b1[3-9]\d{9}\b"), + ]) + names = [r.name for r in engine._rules] + self.assertEqual(names.count("pii_mobile"), 1) + + +if __name__ == "__main__": + unittest.main()