Topic 5: Security and Governance

# module12/governance.py
"""Access control at query time, tenant isolation, personal data, audit trails, and residency."""
from __future__ import annotations

import hashlib
import json
import re
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime, timedelta, timezone
from pathlib import Path

# ---------------------------------------------------------------- 1. personal data
PATTERNS = {
    "email": re.compile(r"\b[\w.+-]+@[\w-]+\.[\w.]+\b"),
    "phone": re.compile(r"\b(?:\+?\d{1,3}[ -]?)?\d{10}\b"),
    "card": re.compile(r"\b(?:\d[ -]?){13,16}\b"),
    "order_id": re.compile(r"\bORD-\d{3,}\b", re.I),
    "aadhaar": re.compile(r"\b\d{4}\s?\d{4}\s?\d{4}\b"),
    "ip": re.compile(r"\b\d{1,3}(?:\.\d{1,3}){3}\b"),
}


def detect_pii(text: str) -> dict[str, list[str]]:
    return {name: pattern.findall(text) for name, pattern in PATTERNS.items() if pattern.search(text)}


def redact(text: str, keep: set[str] = frozenset({"order_id"})) -> str:
    """Replace personal data with labels. Order ids are usually kept: they are needed to help."""
    for name, pattern in PATTERNS.items():
        if name not in keep:
            text = pattern.sub(f"[{name} removed]", text)
    return text


def pseudonymize(value: str, salt: str = "shopsphere") -> str:
    """Stable identifier for analytics without storing who it was."""
    return "u_" + hashlib.sha256(f"{salt}|{value}".encode()).hexdigest()[:16]


# ---------------------------------------------------------------- 2. access control
@dataclass(frozen=True)
class Principal:
    user_id: str
    tenant_id: str
    groups: frozenset = frozenset()
    region: str = "IN"
    scopes: frozenset = frozenset()


FIELD_POLICIES = {
    # field name -> the groups allowed to see it; everyone else gets it masked
    "internal_notes": frozenset({"support-team", "admin"}),
    "cost_price": frozenset({"admin"}),
    "customer_email": frozenset({"support-team", "admin"}),
}


def can_read_document(principal: Principal, doc: dict) -> bool:
    if doc.get("tenant_id", principal.tenant_id) != principal.tenant_id:
        return False                                        # tenant isolation comes first
    visibility = doc.get("visibility", "public")
    if visibility == "public":
        return True
    if principal.groups & frozenset(doc.get("allowed_groups", [])):
        return True
    return principal.user_id in doc.get("allowed_users", [])


def apply_field_policy(principal: Principal, doc: dict) -> dict:
    """Document-level access is not enough: some fields are restricted inside a readable document."""
    filtered = {}
    for field_name, value in doc.items():
        allowed = FIELD_POLICIES.get(field_name)
        filtered[field_name] = value if allowed is None or (principal.groups & allowed) else "[restricted]"
    return filtered


def authorize_documents(principal: Principal, docs: list[dict]) -> tuple[list[dict], list[str]]:
    allowed, denied = [], []
    for doc in docs:
        if can_read_document(principal, doc):
            allowed.append(apply_field_policy(principal, doc))
        else:
            denied.append(doc["doc_id"])
    return allowed, denied


class TenantIsolationError(Exception):
    """Raised when anything in the answer path touches another tenant's data."""


def assert_no_cross_tenant(principal: Principal, docs: list[dict]) -> None:
    leaked = [d["doc_id"] for d in docs if d.get("tenant_id", principal.tenant_id) != principal.tenant_id]
    if leaked:
        raise TenantIsolationError(f"documents from another tenant reached the answer path: {leaked}")


# ---------------------------------------------------------------- 3. residency
RESIDENCY = {"IN": {"regions": {"ap-south-1"}, "note": "India: keep customer data in-region"},
             "EU": {"regions": {"eu-west-1", "eu-central-1"}, "note": "EU: GDPR transfer rules"},
             "US": {"regions": {"us-east-1", "us-west-2"}, "note": "US"}}


def residency_ok(principal: Principal, index_region: str, model_region: str) -> tuple[bool, str]:
    allowed = RESIDENCY.get(principal.region, {}).get("regions", set())
    problems = [name for name, region in (("index", index_region), ("model", model_region))
                if region not in allowed]
    return (not problems, "ok" if not problems else
            f"{principal.region} data would be processed outside its region by: {problems}")


# ---------------------------------------------------------------- 4. audit trail
@dataclass
class AuditRecord:
    correlation_id: str
    at: str
    tenant_id: str
    user_ref: str                     # pseudonymized, not the raw id
    question: str                     # redacted
    retrieved_ids: list[str]
    cited_ids: list[str]
    answered: bool
    refused: bool
    denied_ids: list[str] = field(default_factory=list)
    policy_flags: list[str] = field(default_factory=list)


class AuditLog:
    """Who asked what, what was retrieved, what was answered. Append-only, retained on a schedule."""

    def __init__(self, path: str = "data/audit.jsonl", retention_days: int = 180):
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.retention_days = retention_days

    def write(self, record: AuditRecord) -> None:
        with self.path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(asdict(record)) + "\n")

    def read(self) -> list[dict]:
        if not self.path.exists():
            return []
        return [json.loads(line) for line in self.path.read_text(encoding="utf-8").splitlines()]

    def for_user(self, user_ref: str) -> list[dict]:
        return [r for r in self.read() if r["user_ref"] == user_ref]

    def enforce_retention(self, now: datetime | None = None) -> int:
        """Delete records past the retention window. Run this on a schedule, and prove it ran."""
        now = now or datetime.now(timezone.utc)
        cutoff = now - timedelta(days=self.retention_days)
        rows = self.read()
        kept = [r for r in rows if datetime.fromisoformat(r["at"]) >= cutoff]
        self.path.write_text("\n".join(json.dumps(r) for r in kept) + ("\n" if kept else ""),
                             encoding="utf-8")
        return len(rows) - len(kept)

    def erase_user(self, user_ref: str) -> int:
        """Right-to-erasure request: remove this user's records."""
        rows = self.read()
        kept = [r for r in rows if r["user_ref"] != user_ref]
        self.path.write_text("\n".join(json.dumps(r) for r in kept) + ("\n" if kept else ""),
                             encoding="utf-8")
        return len(rows) - len(kept)


def audit_from_trace(trace, principal: Principal, retrieved_ids, cited_ids, denied_ids,
                     refused: bool, flags=()) -> AuditRecord:
    return AuditRecord(
        correlation_id=trace.correlation_id,
        at=datetime.now(timezone.utc).isoformat(timespec="seconds"),
        tenant_id=principal.tenant_id,
        user_ref=pseudonymize(principal.user_id),
        question=redact(trace.query),
        retrieved_ids=list(retrieved_ids), cited_ids=list(cited_ids), denied_ids=list(denied_ids),
        answered=not refused, refused=refused, policy_flags=list(flags))