216 lines
8.3 KiB
Python
216 lines
8.3 KiB
Python
# -*- 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("<pii>", 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("<pii>", 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()
|