pii_rules.py

"""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)