Redaction Strategies

from __future__ import annotations
import hashlib
import hmac
import json
import os
import re
from dataclasses import dataclass, field
from typing import Literal

# ── Shared types ─────────────────────────────────────────────────────────────

Strategy = Literal["mask", "replace", "pseudonymize", "fpe"]

@dataclass
class RedactionRecord:
    original_hash: str   # SHA-256 of original value — never store the value itself
    pii_type: str
    strategy: Strategy
    replacement: str
    start: int
    end: int

# ── Fake value generators for type-aware replacement ─────────────────────────

_FAKE_BY_TYPE: dict[str, str] = {
    "SSN": "999-00-0000",
    "EMAIL": "[email protected]",
    "PHONE": "555-000-0000",
    "CREDIT_CARD": "4000-0000-0000-0000",
    "PASSPORT": "X00000000",
    "IP_ADDRESS": "192.0.2.0",   # TEST-NET-1, RFC 5737
    "NAME": "Jane Doe",
    "ADDRESS": "1 Placeholder Ln, Anytown, ST 00000",
    "DATE_OF_BIRTH": "1900-01-01",
    "MEDICAL_RECORD": "MRN-REDACTED",
}

def _fake_value(pii_type: str, original: str) -> str:
    return _FAKE_BY_TYPE.get(pii_type, f"[{pii_type}]")

# ── Pseudonymization: deterministic HMAC-based mapping ───────────────────────

_PSEUDO_SECRET = os.getenv("PII_PSEUDO_SECRET", "change-me-in-production")

def _pseudonymize(value: str, pii_type: str) -> str:
    """Produce a deterministic, reversible-with-secret pseudonym."""
    digest = hmac.new(
        _PSEUDO_SECRET.encode(),
        value.encode(),
        hashlib.sha256
    ).hexdigest()[:12]
    return f"{pii_type}-{digest}"

# ── Format-Preserving Encryption stub (FF3-1) ─────────────────────────────────
# Production: use the 'ff3' PyPI package (AES-FF3-1, NIST SP 800-38G Rev. 1)
# Stub shown here for illustration — not cryptographically secure.

def _fpe_encrypt(value: str, pii_type: str) -> str:
    """
    Stub FPE: in production replace with ff3.FF3Cipher().encrypt().
    Preserves digit-only format for numeric types.
    """
    digits = re.sub(r'\D', '', value)
    if not digits:
        return f"[FPE:{pii_type}]"
    # Real FPE would encrypt these digits while preserving length
    encrypted_digits = str(int(digits) ^ 0xDEADBEEF)[:len(digits)].zfill(len(digits))
    return re.sub(r'\d+', encrypted_digits, value, count=1)

# ── Core redaction engine ─────────────────────────────────────────────────────

@dataclass
class DetectedEntity:
    text: str
    start: int
    end: int
    pii_type: str
    confidence: float
    reason: str

def redact(
    text: str,
    entities: list[DetectedEntity],
    strategy: Strategy = "mask",
) -> tuple[str, list[RedactionRecord]]:
    """
    Apply the chosen strategy to all detected entities.
    Returns (redacted_text, list_of_redaction_records).
    Entities must be sorted by start position; overlaps are skipped.
    """
    entities_sorted = sorted(entities, key=lambda e: e.start)
    records: list[RedactionRecord] = []
    result_parts: list[str] = []
    cursor = 0

    for entity in entities_sorted:
        if entity.start < cursor:
            continue  # skip overlapping entity

        result_parts.append(text[cursor : entity.start])

        original_hash = hashlib.sha256(entity.text.encode()).hexdigest()

        if strategy == "mask":
            replacement = f"[{entity.pii_type}]"
        elif strategy == "replace":
            replacement = _fake_value(entity.pii_type, entity.text)
        elif strategy == "pseudonymize":
            replacement = _pseudonymize(entity.text, entity.pii_type)
        elif strategy == "fpe":
            replacement = _fpe_encrypt(entity.text, entity.pii_type)
        else:
            replacement = f"[{entity.pii_type}]"

        result_parts.append(replacement)
        records.append(RedactionRecord(
            original_hash=original_hash,
            pii_type=entity.pii_type,
            strategy=strategy,
            replacement=replacement,
            start=entity.start,
            end=entity.end,
        ))
        cursor = entity.end

    result_parts.append(text[cursor:])
    return "".join(result_parts), records