Audit Trails & Logging

from __future__ import annotations
import hashlib
import json
import logging
import os
import socket
from datetime import datetime, timezone
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Literal
import urllib.request

EventType = Literal[
    "DETECTION", "REDACTION", "ACCESS", "EXPORT",
    "CONFIG_CHANGE", "SCAN_START", "SCAN_COMPLETE"
]

DataClassification = Literal["PII", "CUI", "BOTH", "NONE"]

# ── Audit event dataclass ─────────────────────────────────────────────────────

@dataclass
class AuditEvent:
    event_type: EventType
    agent_id: str
    data_classification: DataClassification
    action_taken: str
    # Mandatory fields set automatically
    timestamp: str = field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )
    # Optional context fields
    document_id: str | None = None
    entity_types: list[str] = field(default_factory=list)
    entity_hashes: list[str] = field(default_factory=list)
    entity_count: int = 0
    confidence_scores: list[float] = field(default_factory=list)
    strategy: str | None = None
    operator_id: str | None = None
    scan_id: str | None = None
    extra: dict[str, Any] = field(default_factory=dict)

    def to_json(self) -> str:
        d = {k: v for k, v in asdict(self).items() if v is not None and v != [] and v != {}}
        return json.dumps(d, separators=(",", ":"))

# ── Audit logger ──────────────────────────────────────────────────────────────

class AuditLogger:
    """Append-only structured audit log with optional SIEM forwarding."""

    def __init__(
        self,
        log_path: Path,
        siem_url: str | None = None,
        siem_token: str | None = None,
    ) -> None:
        self.log_path = log_path
        self.siem_url = siem_url
        self.siem_token = siem_token or os.getenv("SIEM_TOKEN")
        log_path.parent.mkdir(parents=True, exist_ok=True)
        # Open in append mode — never truncate
        self._file = log_path.open("a", encoding="utf-8", buffering=1)

    def emit(self, event: AuditEvent) -> None:
        line = event.to_json() + "\n"
        self._file.write(line)
        self._file.flush()
        if self.siem_url:
            self._forward_to_siem(line)

    def _forward_to_siem(self, line: str) -> None:
        """HTTP forward to SIEM (Splunk HEC, Elastic, etc.)."""
        try:
            data = line.encode()
            req = urllib.request.Request(
                self.siem_url,
                data=data,
                method="POST",
                headers={
                    "Content-Type": "application/json",
                    "Authorization": f"Splunk {self.siem_token}",
                },
            )
            with urllib.request.urlopen(req, timeout=2):
                pass
        except Exception as exc:
            # SIEM forwarding must never block the main pipeline
            logging.warning("SIEM forward failed: %s", exc)

    def close(self) -> None:
        self._file.close()

# ── Convenience helpers ───────────────────────────────────────────────────────

def hash_pii_value(value: str) -> str:
    """SHA-256 of the PII value for audit traceability without storing the value."""
    return hashlib.sha256(value.encode()).hexdigest()

def make_detection_event(
    agent_id: str,
    document_id: str,
    entities: list,   # list of DetectedEntity from detector.py
) -> AuditEvent:
    classification: DataClassification = "NONE"
    cui_types = {"CUI_EXPORT", "CUI_LAW_ENFORCEMENT", "CUI_GOVERNMENT_CONTRACT"}
    has_pii = any(e.pii_type not in cui_types for e in entities)
    has_cui = any(e.pii_type in cui_types for e in entities)
    if has_pii and has_cui:
        classification = "BOTH"
    elif has_pii:
        classification = "PII"
    elif has_cui:
        classification = "CUI"

    return AuditEvent(
        event_type="DETECTION",
        agent_id=agent_id,
        data_classification=classification,
        action_taken="Detection recorded; awaiting redaction policy",
        document_id=document_id,
        entity_types=[e.pii_type for e in entities],
        entity_hashes=[hash_pii_value(e.text) for e in entities],
        entity_count=len(entities),
        confidence_scores=[round(e.confidence, 3) for e in entities],
    )