Add/update pii_rules.py

pii_rules.py

@@ -0,0 +1,101 @@
+"""Deterministic PII layer: regex + checksum/structural validators.
+
+Mirrors the "deterministic layer" described in Rampart's whitepaper, which is the
+*system of record* for classes the neural model is weak on (cards, SSNs) and for
+classes whose structure lives in punctuation (email, URL, IP). Each detector
+returns character spans `(start, end, label)` over the ORIGINAL text, so they can
+be unioned with the model's spans — the deterministic layer taking precedence.
+"""
+import re
+from typing import List, Tuple
+
+Span = Tuple[int, int, str]
+
+EMAIL_RE = re.compile(r"\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}\b")
+URL_RE = re.compile(r"\b(?:https?://|www\.)[^\s<>()]+", re.IGNORECASE)
+IPV4_RE = re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b")
+IPV6_RE = re.compile(r"\b(?:[A-Fa-f0-9]{1,4}:){2,7}[A-Fa-f0-9]{1,4}\b")
+MAC_RE = re.compile(r"\b(?:[A-Fa-f0-9]{2}:){5}[A-Fa-f0-9]{2}\b")
+SSN_RE = re.compile(r"\b(\d{3})[-\s](\d{2})[-\s](\d{4})\b")
+# A run of 13-19 digits, allowing single space/hyphen separators (card shape).
+CARD_RE = re.compile(r"\b(?:\d[ -]?){12,18}\d\b")
+
+
+def _luhn_ok(digits: str) -> bool:
+    if not (13 <= len(digits) <= 19):
+        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
+
+
+def _ssn_ok(area: str, group: str, serial: str) -> bool:
+    a = int(area)
+    if a == 0 or a == 666 or a >= 900:      # reserved / invalid SSA areas
+        return False
+    if int(group) == 0 or int(serial) == 0:
+        return False
+    return True
+
+
+def _ipv4_ok(s: str) -> bool:
+    parts = s.split(".")
+    return len(parts) == 4 and all(p.isdigit() and 0 <= int(p) <= 255 for p in parts)
+
+
+def detect_rules(text: str) -> List[Span]:
+    spans: List[Span] = []
+
+    for m in EMAIL_RE.finditer(text):
+        spans.append((m.start(), m.end(), "EMAIL"))
+    for m in URL_RE.finditer(text):
+        spans.append((m.start(), m.end(), "URL"))
+    for m in MAC_RE.finditer(text):
+        spans.append((m.start(), m.end(), "IP_ADDRESS"))
+    for m in IPV6_RE.finditer(text):
+        spans.append((m.start(), m.end(), "IP_ADDRESS"))
+    for m in IPV4_RE.finditer(text):
+        if _ipv4_ok(m.group(0)):
+            spans.append((m.start(), m.end(), "IP_ADDRESS"))
+    for m in SSN_RE.finditer(text):
+        if _ssn_ok(m.group(1), m.group(2), m.group(3)):
+            spans.append((m.start(), m.end(), "SSN"))
+    for m in CARD_RE.finditer(text):
+        digits = re.sub(r"\D", "", m.group(0))
+        if _luhn_ok(digits):
+            spans.append((m.start(), m.end(), "CREDIT_CARD"))
+
+    return spans
+
+
+def _overlaps(a: Span, b: Span) -> bool:
+    return a[0] < b[1] and b[0] < a[1]
+
+
+def union(model_spans: List[Span], rule_spans: List[Span]) -> List[Span]:
+    """Union the two layers; the deterministic layer wins on any overlap."""
+    final = list(rule_spans)
+    for ms in model_spans:
+        if not any(_overlaps(ms, rs) for rs in rule_spans):
+            final.append(ms)
+    final.sort(key=lambda s: s[0])
+    return final
+
+
+def redact(text: str, spans: List[Span]) -> str:
+    out, last = [], 0
+    for s, e, label in sorted(spans, key=lambda x: x[0]):
+        if s < last:
+            continue
+        out.append(text[last:s])
+        out.append(f"[{label}]")
+        last = e
+    out.append(text[last:])
+    return "".join(out)