src/privacy/redactor.py

"""
PIIRedactor — wraps Microsoft Presidio to redact PII from complaint text locally.

Runs in-process before any text is sent to the Anthropic API or DL models.
Detected spans are replaced with <ENTITY_TYPE> placeholders.

Entity types detected:
    PERSON, PHONE_NUMBER, EMAIL_ADDRESS, CREDIT_CARD, IBAN_CODE,
    US_BANK_NUMBER, IN_AADHAAR, IN_PAN, IN_VEHICLE_REGISTRATION

Failure mode: fail-open — on any error the original text is returned unchanged
so the pipeline is never blocked.

Output dataclass:
    RedactionResult(redacted_text: str, pii_types_found: list[str], pii_redacted: bool)
"""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass, field
from typing import Optional

from presidio_analyzer import AnalyzerEngine, PatternRecognizer, Pattern
from presidio_analyzer.nlp_engine import NlpEngineProvider
from presidio_anonymizer import AnonymizerEngine
from presidio_anonymizer.entities import OperatorConfig

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Entity types to detect — order matters for operator dict construction only.
# ---------------------------------------------------------------------------
_ENTITY_TYPES: list[str] = [
    "PERSON",
    "PHONE_NUMBER",
    "EMAIL_ADDRESS",
    "CREDIT_CARD",
    "IBAN_CODE",
    "US_BANK_NUMBER",
    "IN_AADHAAR",
    "IN_PAN",
    "IN_VEHICLE_REGISTRATION",
]

# ---------------------------------------------------------------------------
# Complaint reference ID deny-list recognizer.
#
# INTENTIONALLY excluded from _ENTITY_TYPES — COMPLAINT_REF_ID spans are
# registered with high confidence (0.99) so Presidio's conflict-resolution
# logic reserves those character positions before US_BANK_NUMBER can claim
# them.  Because COMPLAINT_REF_ID is not in _ENTITY_TYPES, the anonymizer
# never replaces these spans; they pass through verbatim.
# ---------------------------------------------------------------------------

def _complaint_id_deny_recognizer() -> PatternRecognizer:
    """Protect REF-/TRN-/TXN-/OD-/BK-/CLM-/CASE-/CMP-/LN-/POL-/CR- IDs from redaction."""
    return PatternRecognizer(
        supported_entity="COMPLAINT_REF_ID",
        patterns=[
            Pattern(
                name="complaint_ref_id",
                regex=(
                    r"\b(?:REF|TRN|TXN|OD|BK|CLM|CASE|CMP|LN|POL|CR)"
                    r"[-]?[A-Z0-9][-A-Z0-9]*\b"
                ),
                score=0.99,
            )
        ],
    )


# ---------------------------------------------------------------------------
# Custom recognizers for Indian-specific identifiers not in Presidio's default
# set.  Patterns sourced from official format specifications.
# ---------------------------------------------------------------------------

def _aadhaar_recognizer() -> PatternRecognizer:
    """12-digit Aadhaar UID; first digit 2-9; optional spaces or hyphens."""
    return PatternRecognizer(
        supported_entity="IN_AADHAAR",
        patterns=[
            Pattern(
                name="aadhaar_spaced",
                regex=r"\b[2-9]\d{3}[ -]?\d{4}[ -]?\d{4}\b",
                score=0.85,
            )
        ],
        context=["aadhaar", "uid", "unique identification"],
    )


def _pan_recognizer() -> PatternRecognizer:
    """PAN card: 5 uppercase letters · 4 digits · 1 uppercase letter (e.g. ABCDE1234F).

    global_regex_flags omits re.IGNORECASE so only uppercase PANs match —
    lowercase input is not a valid PAN format per official spec.
    """
    return PatternRecognizer(
        supported_entity="IN_PAN",
        patterns=[
            Pattern(
                name="pan_card",
                regex=r"\b[A-Z]{5}[0-9]{4}[A-Z]\b",
                score=0.85,
            )
        ],
        context=["pan", "permanent account number", "income tax"],
        global_regex_flags=re.DOTALL | re.MULTILINE,
    )


def _vehicle_registration_recognizer() -> PatternRecognizer:
    """Indian vehicle registration number (e.g. MH01AB1234, DL3CAF0001)."""
    return PatternRecognizer(
        supported_entity="IN_VEHICLE_REGISTRATION",
        patterns=[
            Pattern(
                name="vehicle_reg",
                # State code (2 letters) · district (1-2 digits) · series (1-3 letters) · number (4 digits)
                regex=r"\b[A-Z]{2}[0-9]{1,2}[A-Z]{1,3}[0-9]{4}\b",
                score=0.75,
            )
        ],
        context=["vehicle", "registration", "number plate", "reg no"],
    )


def _self_introduced_name_recognizer() -> PatternRecognizer:
    """Catch first names that spaCy mislabels when they follow a self-introduction.

    spaCy en_core_web_lg frequently tags Indian first names (e.g. "Vikas") as
    ORG rather than PERSON, so they slip past the default PERSON recognizer and
    reach the API un-redacted. These high-precision context patterns force a
    PERSON span for a Capitalised token immediately following a self-identifying
    phrase. The leading word must be capitalised, so lowercase continuations such
    as "I am writing" or "this is regarding" are not matched.

    Emits the PERSON entity so the existing <PERSON> operator handles replacement.
    """
    name = r"[A-Z][a-z]+"
    phrases = [
        ("name_after_my_name_is", "my name is", 0.95),
        ("name_after_name_is", "name is", 0.90),
        ("name_after_myself", "myself", 0.85),
        ("name_after_i_am", "i am", 0.85),
        ("name_after_im", "i'm", 0.85),
        ("name_after_this_is", "this is", 0.80),
    ]
    return PatternRecognizer(
        supported_entity="PERSON",
        patterns=[
            # Fixed-width scoped-case lookbehind so only the name token is the span.
            Pattern(name=pname, regex=rf"(?<=(?i:{phrase}) ){name}", score=score)
            for pname, phrase, score in phrases
        ],
    )


# ---------------------------------------------------------------------------
# RedactionResult
# ---------------------------------------------------------------------------

@dataclass
class RedactionSpan:
    """A single piece of PII that was detected and replaced.

    *original* is the raw substring from the user's text; it is returned ONLY to
    the same user's browser (it never leaves for any third-party API) so the UI
    can show a side-by-side "what we protected" reveal.
    """
    entity_type: str
    original: str
    placeholder: str
    start: int
    end: int


@dataclass
class RedactionResult:
    """Result of a single redaction pass."""
    redacted_text: str
    pii_types_found: list[str] = field(default_factory=list)
    pii_redacted: bool = False
    spans: list[RedactionSpan] = field(default_factory=list)


# ---------------------------------------------------------------------------
# PIIRedactor
# ---------------------------------------------------------------------------

class PIIRedactor:
    """
    Local PII redactor using presidio-analyzer + presidio-anonymizer.

    Loads spaCy en_core_web_lg on construction (~750 MB, one-time cost).
    Call init_redactor() at server startup so the model load happens at a
    known, controlled moment rather than on the first request.
    """

    def __init__(self) -> None:
        nlp_config = {
            "nlp_engine_name": "spacy",
            "models": [{"lang_code": "en", "model_name": "en_core_web_lg"}],
        }
        nlp_engine = NlpEngineProvider(nlp_configuration=nlp_config).create_engine()

        self._analyzer = AnalyzerEngine(
            nlp_engine=nlp_engine,
            supported_languages=["en"],
        )
        # Register custom recognizers: Indian-specific PII + complaint ID deny-list
        for recognizer in (
            _aadhaar_recognizer(),
            _pan_recognizer(),
            _vehicle_registration_recognizer(),
            _self_introduced_name_recognizer(),
            _complaint_id_deny_recognizer(),
        ):
            self._analyzer.registry.add_recognizer(recognizer)

        self._anonymizer = AnonymizerEngine()

        # Build a static operator dict so every entity type maps to its own
        # <ENTITY_TYPE> placeholder string.  COMPLAINT_REF_ID uses "keep" so
        # its spans are preserved verbatim while still blocking US_BANK_NUMBER
        # from claiming the same character positions.
        self._operators: dict[str, OperatorConfig] = {
            entity: OperatorConfig("replace", {"new_value": f"<{entity}>"})
            for entity in _ENTITY_TYPES
        }
        self._operators["COMPLAINT_REF_ID"] = OperatorConfig("keep")

    def redact(self, text: str) -> RedactionResult:
        """
        Redact PII from *text* and return a RedactionResult.

        On any exception the original text is returned unchanged (fail-open).
        """
        if not text or not text.strip():
            return RedactionResult(redacted_text=text)

        try:
            analyzer_results = self._analyzer.analyze(
                text=text,
                entities=_ENTITY_TYPES + ["COMPLAINT_REF_ID"],
                language="en",
            )

            if not analyzer_results:
                return RedactionResult(redacted_text=text)

            anonymized = self._anonymizer.anonymize(
                text=text,
                analyzer_results=analyzer_results,
                operators=self._operators,
            )

            # COMPLAINT_REF_ID spans use the "keep" operator: they are preserved
            # verbatim AND, being higher-confidence (0.99), they win Presidio's
            # conflict resolution over any lower-score entity that overlaps them
            # (e.g. an 8-digit run inside "REF-20260501-001" mis-detected as
            # US_BANK_NUMBER). Such overlapping entities are therefore NEVER
            # replaced in the output. We must drop them from the reported spans —
            # otherwise the UI claims a redaction that did not happen and the
            # audit leak-check sees the "redacted" value still present in the
            # transmitted text and falsely flags a leak.
            ref_spans = [
                (r.start, r.end)
                for r in analyzer_results
                if r.entity_type == "COMPLAINT_REF_ID"
            ]

            def _overlaps_kept_ref(r) -> bool:
                return any(r.start < end and start < r.end for start, end in ref_spans)

            # Real, actually-redacted results: not the internal sentinel, and not
            # swallowed by a kept COMPLAINT_REF_ID span.
            real_results = [
                r
                for r in analyzer_results
                if r.entity_type != "COMPLAINT_REF_ID" and not _overlaps_kept_ref(r)
            ]
            pii_types_found = sorted({r.entity_type for r in real_results})

            # Build the per-span reveal list (original substring + placeholder),
            # sorted by position so the UI can render them in reading order.
            spans = [
                RedactionSpan(
                    entity_type=r.entity_type,
                    original=text[r.start:r.end],
                    placeholder=f"<{r.entity_type}>",
                    start=r.start,
                    end=r.end,
                )
                for r in sorted(real_results, key=lambda r: r.start)
            ]

            return RedactionResult(
                redacted_text=anonymized.text,
                pii_types_found=pii_types_found,
                # True only when real PII was actually replaced — a message whose
                # only matches were kept ref-IDs (or ref-ID-overlapping noise) is
                # an unmodified passthrough, not a redaction.
                pii_redacted=bool(real_results),
                spans=spans,
            )

        except Exception:
            logger.warning(
                "PIIRedactor.redact failed — returning original text unchanged",
                exc_info=True,
            )
            return RedactionResult(redacted_text=text)


# ---------------------------------------------------------------------------
# Singleton accessors
# ---------------------------------------------------------------------------

_redactor: Optional[PIIRedactor] = None


def init_redactor() -> PIIRedactor:
    """
    Initialise (or reinitialise) the module-level PIIRedactor singleton.

    Call this once at server startup so the spaCy model is loaded eagerly
    rather than on the first request.
    """
    global _redactor
    logger.info("Loading PIIRedactor (spaCy en_core_web_lg)…")
    _redactor = PIIRedactor()
    logger.info("PIIRedactor ready.")
    return _redactor


def get_redactor() -> PIIRedactor:
    """
    Return the PIIRedactor singleton, initialising it lazily if needed.

    Prefer calling init_redactor() explicitly at startup for predictable
    load-time behaviour.
    """
    global _redactor
    if _redactor is None:
        init_redactor()
    return _redactor