qiurui
cbfe4a23dc
- Pipeline: analyze(), redact(), analyze_image() methods - API: /analyze, /redact, /analyze/frame, /analyze/frame/base64 endpoints - Detectors: regex, NER, face (RKNN NPU) - Privacy frame route added for KVM-Privacy Hub integration
437 lines
16 KiB
Python
437 lines
16 KiB
Python
# tests/test_integration_real.py
|
|
"""真实数据集成测试：生成含 PII 的真实文档，验证检测 + 遮罩全链路。
|
|
|
|
格式覆盖：
|
|
PDF（文字层）/ DOCX / XLSX / PNG 图像遮罩
|
|
保密文件拦截 / 多类型混合检测
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import os
|
|
import sys
|
|
import textwrap
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import docx as _docx
|
|
import numpy as np
|
|
import openpyxl
|
|
import pytest
|
|
from pdfminer.high_level import extract_text
|
|
from reportlab.lib.pagesizes import A4
|
|
from reportlab.pdfbase import pdfmetrics
|
|
from reportlab.pdfbase.ttfonts import TTFont
|
|
from reportlab.pdfgen import canvas
|
|
|
|
from info_privacy.pipeline import PrivacyPipeline
|
|
from info_privacy.models import EntityType
|
|
|
|
# ─── 字体注册 ─────────────────────────────────────────────────────────────────
|
|
|
|
_FONT_PATH = "/usr/share/fonts/truetype/arphic/uming.ttc"
|
|
_FONT_NAME = "UMing"
|
|
|
|
if os.path.exists(_FONT_PATH):
|
|
pdfmetrics.registerFont(TTFont(_FONT_NAME, _FONT_PATH, subfontIndex=0))
|
|
_HAS_FONT = True
|
|
else:
|
|
_HAS_FONT = False
|
|
|
|
_needs_font = pytest.mark.skipif(not _HAS_FONT, reason="uming.ttc 字体不存在")
|
|
|
|
# ─── 测试数据常量 ─────────────────────────────────────────────────────────────
|
|
|
|
ID_CARD = "110101199001011234"
|
|
PHONE = "13812345678"
|
|
BANK = "6222021234567890123"
|
|
EMAIL = "zhang.wei@example.com"
|
|
PLATE = "京A12345"
|
|
NAME_CTX = "申请人：张伟，"
|
|
ADDR_CTX = "地址：北京市朝阳区建国路88号"
|
|
|
|
PII_PARAGRAPH = (
|
|
f"申请人：张伟，身份证号：{ID_CARD}\n"
|
|
f"联系电话：{PHONE}，邮箱：{EMAIL}\n"
|
|
f"银行卡：{BANK}，车牌：{PLATE}\n"
|
|
f"{ADDR_CTX}"
|
|
)
|
|
|
|
# ─── 工具函数 ──────────────────────────────────────────────────────────────────
|
|
|
|
def _make_pdf(path: str, lines: list[str]) -> None:
|
|
"""用 reportlab 生成含中文文字层的 PDF。"""
|
|
c = canvas.Canvas(path, pagesize=A4)
|
|
c.setFont(_FONT_NAME, 14)
|
|
y = 750
|
|
for line in lines:
|
|
c.drawString(50, y, line)
|
|
y -= 30
|
|
c.save()
|
|
|
|
|
|
def _pdf_text(path: str) -> str:
|
|
return extract_text(path)
|
|
|
|
|
|
def _docx_text(path: str) -> str:
|
|
doc = _docx.Document(path)
|
|
return "\n".join(p.text for p in doc.paragraphs)
|
|
|
|
|
|
def _xlsx_text(path: str) -> str:
|
|
wb = openpyxl.load_workbook(path, data_only=True)
|
|
parts = []
|
|
for ws in wb.worksheets:
|
|
for row in ws.iter_rows(values_only=True):
|
|
parts.extend(str(c) for c in row if c is not None)
|
|
return " ".join(parts)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def pipeline():
|
|
return PrivacyPipeline()
|
|
|
|
|
|
# ─── PDF（文字层）测试 ────────────────────────────────────────────────────────
|
|
|
|
@_needs_font
|
|
def test_pdf_analyze_detects_id_card(tmp_path, pipeline):
|
|
pdf = str(tmp_path / "pii.pdf")
|
|
_make_pdf(pdf, [
|
|
f"申请人：张伟",
|
|
f"身份证号：{ID_CARD}",
|
|
f"联系电话：{PHONE}",
|
|
f"邮箱：{EMAIL}",
|
|
])
|
|
report = pipeline.analyze(pdf)
|
|
assert report.blocked is False
|
|
types = {e.type for e in report.entities}
|
|
assert EntityType.ID_CARD in types, f"未检测到身份证，实体：{types}"
|
|
|
|
|
|
@_needs_font
|
|
def test_pdf_analyze_detects_phone_and_email(tmp_path, pipeline):
|
|
pdf = str(tmp_path / "pii.pdf")
|
|
_make_pdf(pdf, [f"电话：{PHONE}", f"邮箱：{EMAIL}"])
|
|
report = pipeline.analyze(pdf)
|
|
types = {e.type for e in report.entities}
|
|
assert EntityType.PHONE in types, f"未检测到手机号，实体：{types}"
|
|
assert EntityType.EMAIL in types, f"未检测到邮箱，实体：{types}"
|
|
|
|
|
|
@_needs_font
|
|
def test_pdf_analyze_detects_bank_and_plate(tmp_path, pipeline):
|
|
pdf = str(tmp_path / "pii.pdf")
|
|
_make_pdf(pdf, [f"银行卡：{BANK}", f"车牌：{PLATE}"])
|
|
report = pipeline.analyze(pdf)
|
|
types = {e.type for e in report.entities}
|
|
assert EntityType.BANK_CARD in types, f"未检测到银行卡，实体：{types}"
|
|
assert EntityType.LICENSE_PLATE in types, f"未检测到车牌，实体：{types}"
|
|
|
|
|
|
@_needs_font
|
|
def test_pdf_redact_removes_id_card(tmp_path, pipeline):
|
|
pdf = str(tmp_path / "pii.pdf")
|
|
out = str(tmp_path / "redacted.pdf")
|
|
_make_pdf(pdf, [f"身份证号：{ID_CARD}", f"电话：{PHONE}"])
|
|
pipeline.redact(pdf, redact_types=["id_card"], out_path=out)
|
|
result_text = _pdf_text(out)
|
|
assert ID_CARD not in result_text, "身份证号未被遮罩"
|
|
assert "电话" in result_text, "无关内容被误删"
|
|
|
|
|
|
@_needs_font
|
|
def test_pdf_redact_multi_type(tmp_path, pipeline):
|
|
pdf = str(tmp_path / "pii.pdf")
|
|
out = str(tmp_path / "redacted.pdf")
|
|
_make_pdf(pdf, [f"身份证：{ID_CARD}", f"电话：{PHONE}", f"邮箱：{EMAIL}"])
|
|
pipeline.redact(pdf, redact_types=["id_card", "phone", "email"], out_path=out)
|
|
result_text = _pdf_text(out)
|
|
assert ID_CARD not in result_text
|
|
assert PHONE not in result_text
|
|
assert EMAIL not in result_text
|
|
|
|
|
|
@_needs_font
|
|
def test_pdf_classified_blocked(tmp_path, pipeline):
|
|
pdf = str(tmp_path / "secret.pdf")
|
|
_make_pdf(pdf, ["【机密】本文件仅供内部使用", "禁止外传"])
|
|
report = pipeline.analyze(pdf)
|
|
assert report.blocked is True, "保密文件未被拦截"
|
|
assert report.classification.value == "classified"
|
|
|
|
|
|
@_needs_font
|
|
def test_pdf_normal_no_false_positive(tmp_path, pipeline):
|
|
pdf = str(tmp_path / "normal.pdf")
|
|
_make_pdf(pdf, ["第一季度销售报告", "本季度营业额同比增长12%"])
|
|
report = pipeline.analyze(pdf)
|
|
assert report.blocked is False
|
|
# 不应有高置信度 PII
|
|
high_risk = [e for e in report.entities if e.security_level == "high"]
|
|
assert len(high_risk) == 0, f"正常文档误报高风险实体：{[e.type for e in high_risk]}"
|
|
|
|
|
|
# ─── DOCX 测试 ────────────────────────────────────────────────────────────────
|
|
|
|
@pytest.fixture
|
|
def pii_docx(tmp_path):
|
|
doc = _docx.Document()
|
|
doc.add_paragraph(f"申请人：张伟 身份证：{ID_CARD}")
|
|
doc.add_paragraph(f"联系电话：{PHONE} 邮箱：{EMAIL}")
|
|
doc.add_paragraph(f"银行卡号：{BANK}")
|
|
doc.add_paragraph(f"住址：北京市朝阳区建国路88号")
|
|
p = str(tmp_path / "pii.docx")
|
|
doc.save(p)
|
|
return p
|
|
|
|
|
|
def test_docx_analyze_detects_id_card(pii_docx, pipeline):
|
|
report = pipeline.analyze(pii_docx)
|
|
types = {e.type for e in report.entities}
|
|
assert EntityType.ID_CARD in types, f"未检测到身份证，实体：{types}"
|
|
|
|
|
|
def test_docx_analyze_detects_phone(pii_docx, pipeline):
|
|
report = pipeline.analyze(pii_docx)
|
|
types = {e.type for e in report.entities}
|
|
assert EntityType.PHONE in types, f"未检测到手机号，实体：{types}"
|
|
|
|
|
|
def test_docx_analyze_detects_name(pii_docx, pipeline):
|
|
report = pipeline.analyze(pii_docx)
|
|
types = {e.type for e in report.entities}
|
|
assert EntityType.NAME in types, f"未检测到姓名，实体：{types}"
|
|
|
|
|
|
def test_docx_analyze_detects_address(pii_docx, pipeline):
|
|
report = pipeline.analyze(pii_docx)
|
|
types = {e.type for e in report.entities}
|
|
assert EntityType.ADDRESS in types, f"未检测到地址，实体：{types}"
|
|
|
|
|
|
def test_docx_summary_counts(pii_docx, pipeline):
|
|
report = pipeline.analyze(pii_docx)
|
|
assert report.summary.get("id_card", 0) >= 1
|
|
assert report.summary.get("phone", 0) >= 1
|
|
|
|
|
|
def test_docx_redact_id_card(pii_docx, tmp_path, pipeline):
|
|
out = str(tmp_path / "out.docx")
|
|
pipeline.redact(pii_docx, redact_types=["id_card"], out_path=out)
|
|
text = _docx_text(out)
|
|
assert ID_CARD not in text, "身份证号未被遮罩"
|
|
assert "张伟" in text, "姓名被误删"
|
|
assert PHONE in text, "手机号被误删"
|
|
|
|
|
|
def test_docx_redact_multi_type(pii_docx, tmp_path, pipeline):
|
|
out = str(tmp_path / "out.docx")
|
|
pipeline.redact(pii_docx, redact_types=["id_card", "phone", "bank_card"], out_path=out)
|
|
text = _docx_text(out)
|
|
assert ID_CARD not in text
|
|
assert PHONE not in text
|
|
assert BANK not in text
|
|
|
|
|
|
def test_docx_redact_preserves_unspecified(pii_docx, tmp_path, pipeline):
|
|
"""未指定的类型不应被遮罩。"""
|
|
out = str(tmp_path / "out.docx")
|
|
pipeline.redact(pii_docx, redact_types=["id_card"], out_path=out)
|
|
text = _docx_text(out)
|
|
assert EMAIL in text, "邮箱不应被遮罩"
|
|
assert BANK in text, "银行卡不应被遮罩"
|
|
|
|
|
|
def test_docx_classified_blocked(tmp_path, pipeline):
|
|
doc = _docx.Document()
|
|
doc.add_paragraph("【绝密】本文件属于国家机密")
|
|
p = str(tmp_path / "secret.docx")
|
|
doc.save(p)
|
|
report = pipeline.analyze(p)
|
|
assert report.blocked is True
|
|
assert len(report.entities) == 0, "拦截文件不应返回实体列表"
|
|
|
|
|
|
def test_docx_meta_fields(pii_docx, pipeline):
|
|
report = pipeline.analyze(pii_docx)
|
|
assert report.doc_meta["format"] == "docx"
|
|
assert report.doc_meta["has_text_layer"] is True
|
|
assert report.doc_meta["redact_strategy"] == "text_replace"
|
|
|
|
|
|
# ─── XLSX 测试 ────────────────────────────────────────────────────────────────
|
|
|
|
@pytest.fixture
|
|
def pii_xlsx(tmp_path):
|
|
wb = openpyxl.Workbook()
|
|
ws = wb.active
|
|
ws.title = "人员信息"
|
|
ws.append(["姓名", "身份证号", "手机号", "邮箱", "银行卡"])
|
|
ws.append(["张伟", ID_CARD, PHONE, EMAIL, BANK])
|
|
ws.append(["李娜", "310101199505051234", "13900001111", "li@test.com", "6222021111111111111"])
|
|
p = str(tmp_path / "pii.xlsx")
|
|
wb.save(p)
|
|
return p
|
|
|
|
|
|
def test_xlsx_analyze_detects_pii(pii_xlsx, pipeline):
|
|
report = pipeline.analyze(pii_xlsx)
|
|
types = {e.type for e in report.entities}
|
|
assert EntityType.ID_CARD in types
|
|
assert EntityType.PHONE in types
|
|
|
|
|
|
def test_xlsx_analyze_two_id_cards(pii_xlsx, pipeline):
|
|
report = pipeline.analyze(pii_xlsx)
|
|
id_count = sum(1 for e in report.entities if e.type == EntityType.ID_CARD)
|
|
assert id_count >= 2, f"应检测到至少 2 个身份证，实际：{id_count}"
|
|
|
|
|
|
def test_xlsx_redact_removes_id_cards(pii_xlsx, tmp_path, pipeline):
|
|
out = str(tmp_path / "out.xlsx")
|
|
pipeline.redact(pii_xlsx, redact_types=["id_card"], out_path=out)
|
|
text = _xlsx_text(out)
|
|
assert ID_CARD not in text, "身份证号未被遮罩"
|
|
assert "310101199505051234" not in text, "第二个身份证号未被遮罩"
|
|
|
|
|
|
def test_xlsx_redact_preserves_headers(pii_xlsx, tmp_path, pipeline):
|
|
out = str(tmp_path / "out.xlsx")
|
|
pipeline.redact(pii_xlsx, redact_types=["id_card"], out_path=out)
|
|
wb = openpyxl.load_workbook(out, data_only=True)
|
|
ws = wb.active
|
|
headers = [cell.value for cell in ws[1]]
|
|
assert "姓名" in headers, "列头被误删"
|
|
assert "身份证号" in headers, "列头被误删"
|
|
|
|
|
|
def test_xlsx_meta_fields(pii_xlsx, pipeline):
|
|
report = pipeline.analyze(pii_xlsx)
|
|
assert report.doc_meta["format"] == "xlsx"
|
|
assert report.doc_meta["redact_strategy"] == "text_replace"
|
|
|
|
|
|
# ─── PNG 图像遮罩测试（不依赖 RKNN，直接测试 ImageRedactor）────────────────────
|
|
|
|
@pytest.fixture
|
|
def pii_png(tmp_path):
|
|
"""生成包含文字的 PNG（用于测试图像遮罩，OCR 阶段跳过）。"""
|
|
img = np.ones((400, 800, 3), dtype=np.uint8) * 240
|
|
cv2.putText(img, f"ID: {ID_CARD}", (20, 80),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0), 2)
|
|
cv2.putText(img, f"TEL: {PHONE}", (20, 160),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0), 2)
|
|
cv2.rectangle(img, (200, 220), (400, 320), (180, 200, 220), -1) # 模拟人脸区域
|
|
p = str(tmp_path / "pii.png")
|
|
cv2.imwrite(p, img)
|
|
return p
|
|
|
|
|
|
def test_png_image_redactor_black_box(pii_png, tmp_path):
|
|
"""ImageRedactor 对指定 bbox 覆盖纯黑矩形。"""
|
|
from info_privacy.redactors.image_redactor import ImageRedactor
|
|
from info_privacy.models import Entity, EntityType
|
|
|
|
out = str(tmp_path / "out.png")
|
|
entity = Entity(
|
|
id="f1", type=EntityType.FACE, value=None,
|
|
page=1, bbox=[200.0, 220.0, 400.0, 320.0],
|
|
layer="image", security_level="high",
|
|
)
|
|
ImageRedactor().redact(pii_png, [entity], out)
|
|
|
|
result = cv2.imread(out)
|
|
region = result[220:320, 200:400]
|
|
assert region.max() == 0, "遮罩区域应为纯黑"
|
|
|
|
|
|
def test_png_image_redactor_preserves_size(pii_png, tmp_path):
|
|
"""遮罩后图像尺寸不变。"""
|
|
from info_privacy.redactors.image_redactor import ImageRedactor
|
|
from info_privacy.models import Entity, EntityType
|
|
|
|
out = str(tmp_path / "out.png")
|
|
entity = Entity(
|
|
id="f2", type=EntityType.FACE, value=None,
|
|
page=1, bbox=[20.0, 40.0, 300.0, 100.0],
|
|
layer="image", security_level="high",
|
|
)
|
|
ImageRedactor().redact(pii_png, [entity], out)
|
|
orig = cv2.imread(pii_png)
|
|
result = cv2.imread(out)
|
|
assert orig.shape == result.shape
|
|
|
|
|
|
def test_png_image_redactor_outside_bbox_intact(pii_png, tmp_path):
|
|
"""遮罩区域外的像素不应被修改。"""
|
|
from info_privacy.redactors.image_redactor import ImageRedactor
|
|
from info_privacy.models import Entity, EntityType
|
|
|
|
out = str(tmp_path / "out.png")
|
|
# 只遮罩右下角
|
|
entity = Entity(
|
|
id="f3", type=EntityType.FACE, value=None,
|
|
page=1, bbox=[600.0, 300.0, 780.0, 380.0],
|
|
layer="image", security_level="high",
|
|
)
|
|
ImageRedactor().redact(pii_png, [entity], out)
|
|
orig = cv2.imread(pii_png)
|
|
result = cv2.imread(out)
|
|
# 左上角应与原图一致
|
|
np.testing.assert_array_equal(
|
|
orig[0:50, 0:50], result[0:50, 0:50],
|
|
err_msg="非遮罩区域像素被修改",
|
|
)
|
|
|
|
|
|
def test_png_image_redactor_multiple_entities(pii_png, tmp_path):
|
|
"""多个 entity 都被遮罩。"""
|
|
from info_privacy.redactors.image_redactor import ImageRedactor
|
|
from info_privacy.models import Entity, EntityType
|
|
|
|
out = str(tmp_path / "out.png")
|
|
entities = [
|
|
Entity(id="e1", type=EntityType.FACE, value=None,
|
|
page=1, bbox=[20.0, 40.0, 400.0, 110.0], layer="image", security_level="high"),
|
|
Entity(id="e2", type=EntityType.FACE, value=None,
|
|
page=1, bbox=[200.0, 220.0, 400.0, 320.0], layer="image", security_level="high"),
|
|
]
|
|
ImageRedactor().redact(pii_png, entities, out)
|
|
result = cv2.imread(out)
|
|
r1 = result[40:110, 20:400]
|
|
r2 = result[220:320, 200:400]
|
|
assert r1.max() == 0, "第一个遮罩区域未被覆盖"
|
|
assert r2.max() == 0, "第二个遮罩区域未被覆盖"
|
|
|
|
|
|
# ─── 综合 summary / warning 测试 ─────────────────────────────────────────────
|
|
|
|
def test_sensitive_partial_warning_many_entities(tmp_path, pipeline):
|
|
"""大量高危实体触发 sensitive_partial 警告。"""
|
|
doc = _docx.Document()
|
|
for i in range(7):
|
|
n = int(ID_CARD[:17]) + i
|
|
doc.add_paragraph(f"身份证：{n}X")
|
|
p = str(tmp_path / "many.docx")
|
|
doc.save(p)
|
|
report = pipeline.analyze(p)
|
|
if report.warning:
|
|
# 有警告时，分类应为 sensitive_partial
|
|
assert report.classification.value == "sensitive_partial"
|
|
|
|
|
|
def test_redact_returns_security_report(tmp_path, pipeline):
|
|
"""redact() 返回字典包含 redacted 和 strategy_used 字段。"""
|
|
doc = _docx.Document()
|
|
doc.add_paragraph(f"身份证：{ID_CARD}")
|
|
p = str(tmp_path / "t.docx")
|
|
doc.save(p)
|
|
out = str(tmp_path / "out.docx")
|
|
report = pipeline.redact(p, redact_types=["id_card"], out_path=out)
|
|
assert "redacted" in report
|
|
assert "strategy_used" in report
|
|
assert report["strategy_used"] == "text_replace"
|
|
assert report["security_guarantee"] == "byte_level"
|