154 lines
5.0 KiB
Python
154 lines
5.0 KiB
Python
"""PII 识别引擎：内置正则规则 + 校验算法，支持自定义规则注册。
|
||
|
||
输出统一为 {"type", "text", "page", "bbox"} 列表。
|
||
坐标来源是"词"粒度（extractor 的 words），因此匹配策略：
|
||
1) 先把整页 words 按行拼接（y 中心相近视为同线），
|
||
2) 在行文本上跑规则，命中后把命中的字符区间映射回覆盖的 word bbox，
|
||
3) 多个 word 命中时 bbox 取并集。
|
||
"""
|
||
|
||
import re
|
||
from dataclasses import dataclass, field
|
||
from typing import Callable, List, Optional
|
||
|
||
|
||
# ---------- 校验算法 ----------
|
||
|
||
def _luhn_ok(digits: str) -> bool:
|
||
if not digits.isdigit() or len(digits) < 12:
|
||
return False
|
||
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
|
||
|
||
|
||
_ID_WEIGHTS = [7, 9, 10, 5, 8, 4, 2, 1, 6, 3, 7, 9, 10, 5, 8, 4, 2]
|
||
_ID_CODES = "10X98765432"
|
||
|
||
|
||
def _idcard_ok(s: str) -> bool:
|
||
if len(s) != 18:
|
||
return False
|
||
try:
|
||
total = sum(int(s[i]) * _ID_WEIGHTS[i] for i in range(17))
|
||
except ValueError:
|
||
return False
|
||
return _ID_CODES[total % 11] == s[17].upper()
|
||
|
||
|
||
def _usc_ok(s: str) -> bool:
|
||
"""统一社会信用代码校验（GB 32100-2015）。"""
|
||
if len(s) != 18:
|
||
return False
|
||
chars = "0123456789ABCDEFGHJKLMNPQRTUWXY"
|
||
weights = [1, 3, 9, 27, 19, 26, 16, 17, 20, 29, 25, 13, 8, 24, 10, 30, 28]
|
||
try:
|
||
total = sum(chars.index(s[i]) * weights[i] for i in range(17))
|
||
except ValueError:
|
||
return False
|
||
return chars[(31 - total % 31) % 31] == s[17]
|
||
|
||
|
||
# ---------- 规则定义 ----------
|
||
|
||
@dataclass
|
||
class Rule:
|
||
name: str
|
||
pattern: "re.Pattern"
|
||
validator: Optional[Callable[[str], bool]] = field(default=None)
|
||
|
||
|
||
def r(pattern: str) -> "re.Pattern":
|
||
return re.compile(pattern)
|
||
|
||
|
||
# 左右边界：仅对 ASCII 字母/数字要求非邻接，允许紧贴中文
|
||
_L, _R = r"(?<![0-9A-Za-z])", r"(?![0-9A-Za-z])"
|
||
|
||
|
||
def _b(core: str) -> str:
|
||
return _L + core + _R
|
||
|
||
|
||
BUILTIN_RULES: List[Rule] = [
|
||
Rule("id_card", r(_b(r"[1-9]\d{5}(?:19|20)\d{2}(?:0[1-9]|1[0-2])(?:[0-2]\d|3[01])\d{3}[\dXx]")), _idcard_ok),
|
||
Rule("phone", r(_b(r"1[3-9]\d{9}"))),
|
||
Rule("bank_card", r(_b(r"[3-6]\d{15,18}")), _luhn_ok),
|
||
Rule("usc_code", r(_b(r"[159Y][1239]\d{6}[0-9A-HJ-NPQRTUWXY]{2}\d{6}[0-9A-HJ-NPQRTUWXY]")), _usc_ok),
|
||
Rule("email", r(_b(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}"))),
|
||
]
|
||
|
||
|
||
# ---------- 检测器 ----------
|
||
|
||
class PIIDetector:
|
||
def __init__(self, rules: Optional[List[Rule]] = None):
|
||
self.rules = list(rules) if rules else list(BUILTIN_RULES)
|
||
|
||
def register(self, rule: Rule) -> None:
|
||
"""注册自定义规则。"""
|
||
self.rules.append(rule)
|
||
|
||
def _match(self, text: str):
|
||
for rule in self.rules:
|
||
for m in rule.pattern.finditer(text):
|
||
if rule.validator is None or rule.validator(m.group()):
|
||
yield rule.name, m
|
||
|
||
def scan_text(self, text: str) -> List[dict]:
|
||
"""纯文本检测（无坐标），返回 [{type, text}]。"""
|
||
return [{"type": t, "text": m.group()} for t, m in self._match(text)]
|
||
|
||
# -- 行重组：把 words 按基线聚成行 --
|
||
|
||
@staticmethod
|
||
def _group_lines(words, y_tol=4.0):
|
||
lines = []
|
||
for w in sorted(words, key=lambda w: (round(w["bbox"][1], 1), w["bbox"][0])):
|
||
placed = False
|
||
for line in lines:
|
||
ly = sum(w2["bbox"][1] + w2["bbox"][3] for w2 in line) / (2 * len(line))
|
||
wy = (w["bbox"][1] + w["bbox"][3]) / 2
|
||
if abs(ly - wy) <= y_tol:
|
||
line.append(w)
|
||
placed = True
|
||
break
|
||
if not placed:
|
||
lines.append([w])
|
||
for line in lines:
|
||
line.sort(key=lambda w: w["bbox"][0])
|
||
return lines
|
||
|
||
def scan_page_words(self, words: List[dict], page: int) -> List[dict]:
|
||
"""对单页 words（来自 extractor）做 PII 检测，返回带 bbox 命中。"""
|
||
hits: List[dict] = []
|
||
for line in self._group_lines(words):
|
||
text = "".join(w["text"] for w in line)
|
||
# 记录每个 word 在拼接文本中的起止偏移
|
||
offsets, pos = [], 0
|
||
for w in line:
|
||
offsets.append((pos, pos + len(w["text"])))
|
||
pos += len(w["text"])
|
||
for ptype, m in self._match(text):
|
||
s, e = m.span()
|
||
covered = [w for w, (os_, oe) in zip(line, offsets) if oe > s and os_ < e]
|
||
if not covered:
|
||
continue
|
||
x0 = min(w["bbox"][0] for w in covered)
|
||
y0 = min(w["bbox"][1] for w in covered)
|
||
x1 = max(w["bbox"][2] for w in covered)
|
||
y1 = max(w["bbox"][3] for w in covered)
|
||
hits.append({"type": ptype, "text": m.group(), "page": page, "bbox": (x0, y0, x1, y1)})
|
||
return hits
|
||
|
||
|
||
def default_detector() -> PIIDetector:
|
||
return PIIDetector()
|