utils.py

"""
utils.py - PII Masking utilities using Regex and spaCy NER.

Detects and masks the following PII/PCI entities:
    - full_name
    - email
    - phone_number
    - dob
    - aadhar_num
    - credit_debit_no
    - cvv_no
    - expiry_no
"""

import re
from typing import Any, Dict, List, Tuple

# ---------------------------------------------------------------------------
# Compiled regex patterns (ordered: most-specific / longest match first)
# ---------------------------------------------------------------------------

# Credit / debit card: 16 digits with optional space/hyphen separators
_CREDIT_DEBIT_PAT = re.compile(
    r"\b(\d{4}[\s\-]?\d{4}[\s\-]?\d{4}[\s\-]?\d{4})\b"
)

# Aadhar card: 12 digits (4-4-4 with optional space/hyphen)
# Must come AFTER credit-card so 16-digit hits are already consumed.
_AADHAR_PAT = re.compile(
    r"\b(\d{4}[\s\-]?\d{4}[\s\-]?\d{4})\b"
)

# CVV: only the digit value (3–4 digits) after the CVV/CVC keyword
_CVV_PAT = re.compile(
    r"\b(?:CVV|CVC)[:\s#]+(?P<cvv>\d{3,4})\b",
    re.IGNORECASE,
)

# Card expiry: MM/YY or MM/YYYY
_EXPIRY_PAT = re.compile(
    r"\b(0[1-9]|1[0-2])\s*/\s*(\d{2}|\d{4})\b"
)

# Date of birth: DD/MM/YYYY or DD-MM-YYYY
_DOB_PAT = re.compile(
    r"\b(0?[1-9]|[12]\d|3[01])[\/\-](0?[1-9]|1[0-2])[\/\-](19|20)\d{2}\b"
)

# Email address
_EMAIL_PAT = re.compile(
    r"\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}\b"
)

# Phone number: international or local formats
_PHONE_PAT = re.compile(
    r"(?<!\d)(\+\d{1,3}[\s\-]?)?"
    r"(\(?\d{3}\)?[\s\-]?\d{3}[\s\-]?\d{4})(?!\d)"
)

# Full name: explicit trigger phrases followed by a capitalised name.
# Named group 'name' captures only the name, not the trigger phrase.
_NAME_EXPLICIT_PAT = re.compile(
    r"(?:My name is|I am|name\s*:|naam\s*:|ich bin"
    r"|je suis|mi nombre es)\s+"
    r"(?P<name>[A-Z][a-z]+(?:[\s][A-Z][a-z]+){1,3})",
    re.IGNORECASE,
)

# Words that must never be tagged as a person name (spaCy false-positive guard)
_NAME_BLOCKLIST = frozenset(
    [
        "subject", "dear", "hello", "hi", "regards", "sincerely",
        "thanks", "thank", "please", "note", "expiry", "card",
        "number", "account", "billing", "payment", "date", "issue",
        "support", "customer", "request", "change", "problem",
        "incident", "update", "new", "old", "service",
    ]
)


def _load_spacy():
    """Lazily load spaCy model; returns None if unavailable."""
    try:
        import spacy  # noqa: PLC0415
        try:
            return spacy.load("en_core_web_sm")
        except OSError:
            try:
                return spacy.load("xx_ent_wiki_sm")
            except OSError:
                return None
    except ImportError:
        return None


_NLP = None  # populated on first call


def _get_nlp():
    global _NLP  # noqa: PLW0603
    if _NLP is None:
        _NLP = _load_spacy()
    return _NLP


# ---------------------------------------------------------------------------
# Core masking logic
# ---------------------------------------------------------------------------

def _find_entities(text: str) -> List[Dict[str, Any]]:
    """
    Find all PII entities in *text*.

    Returns a list of dicts:
        position       : [start_char, end_char] in the original text
        classification : entity label string
        entity         : original entity value
    Overlapping spans are resolved by keeping the first (longest) match.
    """
    # (start, end, label, value)
    raw_hits: List[Tuple[int, int, str, str]] = []

    # 1. Regex pipeline -------------------------------------------------------
    regex_patterns = [
        ("credit_debit_no", _CREDIT_DEBIT_PAT, "full"),
        ("aadhar_num", _AADHAR_PAT, "full"),
        ("cvv_no", _CVV_PAT, "cvv"),       # named group
        ("expiry_no", _EXPIRY_PAT, "full"),
        ("dob", _DOB_PAT, "full"),
        ("email", _EMAIL_PAT, "full"),
        ("phone_number", _PHONE_PAT, "full"),
        ("full_name", _NAME_EXPLICIT_PAT, "name"),  # named group
    ]

    for label, pattern, capture in regex_patterns:
        for m in pattern.finditer(text):
            if capture == "full":
                raw_hits.append((m.start(), m.end(), label, m.group()))
            else:
                # Use named capture group; fall back to full match
                try:
                    val = m.group(capture)
                    start = m.start(capture)
                    end = m.end(capture)
                    raw_hits.append((start, end, label, val))
                except (IndexError, error := Exception):  # noqa: F841
                    raw_hits.append((m.start(), m.end(), label, m.group()))

    # 2. spaCy PERSON entities (supplement regex name detection) --------------
    nlp = _get_nlp()
    if nlp is not None:
        doc = nlp(text)
        for ent in doc.ents:
            if ent.label_ == "PERSON":
                name = ent.text.strip()
                # Filter single-token or blocklisted false positives
                if (
                    len(name.split()) >= 2
                    and name.lower() not in _NAME_BLOCKLIST
                ):
                    raw_hits.append(
                        (ent.start_char, ent.end_char, "full_name", name)
                    )

    # 3. Sort: start ascending, length descending (longest span wins) ---------
    raw_hits.sort(key=lambda x: (x[0], -(x[1] - x[0])))

    # 4. Remove overlaps (greedy) --------------------------------------------
    merged: List[Tuple[int, int, str, str]] = []
    last_end = -1
    for start, end, label, value in raw_hits:
        if start >= last_end:
            merged.append((start, end, label, value))
            last_end = end

    return [
        {"position": [start, end], "classification": label, "entity": value}
        for start, end, label, value in merged
    ]


def mask_pii(text: str) -> Tuple[str, List[Dict[str, Any]]]:
    """
    Mask PII in *text*.

    Returns:
        masked_text : email body with PII replaced by [entity_type] tokens
        entities    : list of entity dicts (positions refer to original text)
    """
    entities = _find_entities(text)

    # Build masked string left-to-right using a running offset
    masked = text
    offset = 0
    for ent in entities:
        start = ent["position"][0] + offset
        end = ent["position"][1] + offset
        token = f"[{ent['classification']}]"
        masked = masked[:start] + token + masked[end:]
        offset += len(token) - (end - start)

    return masked, entities


def demask_email(masked_text: str, entities: List[Dict[str, Any]]) -> str:
    """
    Restore original PII values from *entities* back into *masked_text*.
    Each [entity_type] token is replaced exactly once in order.
    """
    result = masked_text
    for ent in entities:
        token = f"[{ent['classification']}]"
        result = result.replace(token, ent["entity"], 1)
    return result