src/core/trace_redaction.py

"""Deterministic, no-LLM redaction of the audit trace before it is persisted (ADR-0013).

The always-on audit trace (ADR-0008) captures the full prompt and the code the
agent generated and ran, verbatim, into the trace store — and that store *is* the
audit control. A secret, credential, or personal identifier landing in it is a
leak into the very artifact meant to be trusted. This module scrubs that payload
*once*, at a single seam right before it reaches any sink (``local`` / ``hf`` /
``s3``), so the sinks stay dumb and every path inherits redaction for free.

Design (mirrors the shared-helper pattern of :mod:`src.core.integrity`):

- :func:`redact_trace` is a **pure function**: it deep-copies the trace and walks
  every string value, replacing matches with a typed placeholder
  (``«REDACTED:anthropic_key»``, ``«REDACTED:hf_token»``, …). The in-memory trace
  the UI / eval harness reads is never mutated — only the persisted copy.
- :func:`redact_trace_safe` is the **fail-closed** wrapper the run path calls: if
  redaction itself raises, it returns a *minimal* trace (run_id + timestamp +
  ``redaction_error``) rather than the raw payload — never write an unredacted
  trace, but never crash the run either (the logging sink is fail-open around
  this in turn).

Redaction is **best-effort scrubbing, not a proof of secret-free logs**: the
patterns are high-precision regex, so a novel secret format can slip through.
It is defense-in-depth layered with not exposing secrets to the agent in the
first place (ADR-0012 service-token hygiene, ADR-0014 secret scanning) — it
reduces *accidental* capture, it does not license putting secrets in prompts.

Config (all env, all optional):

  TRACE_REDACTION            on (default) | off   — ``off`` is local-debug only and
                                                    is NEVER the prod posture.
  TRACE_MAX_FIELD_CHARS      per-value size cap    (default: 20000)
  TRACE_REDACTION_EXTRA      extra regexes, newline- or comma-separated, each
                             scrubbed to ``«REDACTED:custom»`` (extensible tuning).
"""

from __future__ import annotations

import copy
import math
import os
import re
import time
from typing import Any

# --------------------------------------------------------------------------- #
# Config
# --------------------------------------------------------------------------- #
_DEFAULT_MAX_FIELD_CHARS = 20_000


def _redaction_enabled() -> bool:
    """True unless TRACE_REDACTION is explicitly ``off`` (default on).

    ``off`` is for local debugging only and is never the production posture.
    """
    return os.environ.get("TRACE_REDACTION", "on").strip().lower() != "off"


def _max_field_chars() -> int:
    """Per-value truncation cap. Non-positive / unparsable disables truncation."""
    raw = os.environ.get("TRACE_MAX_FIELD_CHARS")
    if raw is None:
        return _DEFAULT_MAX_FIELD_CHARS
    try:
        val = int(raw)
    except (TypeError, ValueError):
        return _DEFAULT_MAX_FIELD_CHARS
    return val if val > 0 else 0


# --------------------------------------------------------------------------- #
# Pattern set — deterministic, high-precision. Order matters: the most specific
# key shapes are scrubbed before the generic `sk-…` / Bearer catch-alls.
# --------------------------------------------------------------------------- #
def _placeholder(kind: str) -> str:
    return f"«REDACTED:{kind}»"


# (compiled_regex, kind). Applied in order via re.sub.
_PATTERNS: list[tuple[re.Pattern, str]] = [
    # Credentialed URL (user:pass@host) — scrub the embedded creds, keep the host
    # so a trace line stays readable. Must run before the generic patterns.
    (re.compile(r"(https?://)[^\s:/@]+:[^\s:/@]+@"), "credential_url"),
    # Anthropic keys (sk-ant-…) — before the generic OpenAI-style sk-…
    (re.compile(r"sk-ant-[A-Za-z0-9_\-]{16,}"), "anthropic_key"),
    # HuggingFace tokens
    (re.compile(r"\bhf_[A-Za-z0-9]{20,}\b"), "hf_token"),
    # AWS access key id
    (re.compile(r"\bAKIA[0-9A-Z]{16}\b"), "aws_access_key"),
    # OpenAI-style secret keys (generic sk-…), after sk-ant-… above
    (re.compile(r"\bsk-[A-Za-z0-9]{20,}\b"), "openai_key"),
    # Bearer <token>
    (re.compile(r"[Bb]earer\s+[A-Za-z0-9._\-]{8,}"), "bearer_token"),
    # GitHub / generic ghp_/gho_/ghs_ tokens (cheap, high precision)
    (re.compile(r"\bgh[pousr]_[A-Za-z0-9]{20,}\b"), "github_token"),
    # Email addresses (PII; conservative, on by default). Last, so it never
    # clobbers a token that happens to contain an @.
    (re.compile(r"\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}\b"), "email"),
]

# Credentialed-URL sub keeps the scheme + host, drops the "user:pass@".
_CREDENTIAL_URL_RE = _PATTERNS[0][0]

# An env line: UPPER_SNAKE=value. Three or more => an env dump; drop the whole value.
_ENV_LINE_RE = re.compile(r"^\s*[A-Z][A-Z0-9_]{2,}=.+$", re.MULTILINE)
_ENV_DUMP_MIN_LINES = 3

# AWS secret-access-key shape: a 40-or-more-char base64-ish run, standalone (AWS
# secrets are exactly 40; ≥40 also catches other long secret blobs). Paired with a
# Shannon-entropy gate so ordinary long text or gene-id runs don't trip it.
# Anchored on non-base64 boundaries to avoid slicing a longer token.
_SECRET40_RE = re.compile(r"(?<![A-Za-z0-9/+])[A-Za-z0-9/+]{40,}(?![A-Za-z0-9/+])")
_SECRET40_MIN_ENTROPY = 4.0  # bits/char; random base64 ≈ 6, English prose ≈ 2–3.


def _shannon_entropy(s: str) -> float:
    """Bits-per-character Shannon entropy of ``s`` (0.0 for empty)."""
    if not s:
        return 0.0
    counts: dict[str, int] = {}
    for ch in s:
        counts[ch] = counts.get(ch, 0) + 1
    n = len(s)
    return -sum((c / n) * math.log2(c / n) for c in counts.values())


def _looks_like_env_dump(s: str) -> bool:
    """True if ``s`` contains several UPPER_SNAKE=value lines (a leaked env)."""
    return len(_ENV_LINE_RE.findall(s)) >= _ENV_DUMP_MIN_LINES


def _extra_patterns() -> list[re.Pattern]:
    """Compile TRACE_REDACTION_EXTRA (newline/comma-separated regexes)."""
    raw = os.environ.get("TRACE_REDACTION_EXTRA")
    if not raw:
        return []
    out: list[re.Pattern] = []
    for piece in re.split(r"[\n,]", raw):
        piece = piece.strip()
        if not piece:
            continue
        try:
            out.append(re.compile(piece))
        except re.error:
            # A bad custom pattern must never break redaction of everything else.
            continue
    return out


def _scrub_string(s: str, max_chars: int, extra: list[re.Pattern]) -> str:
    """Scrub one string value: env-dump drop, pattern subs, high-entropy secrets, cap.

    Scrubbing runs *before* truncation so a secret straddling the cap boundary is
    still removed rather than half-exposed.
    """
    if not s:
        return s

    # Whole-value drop for an obvious environment dump.
    if _looks_like_env_dump(s):
        return _placeholder("env_dump")

    # Credentialed URL: keep scheme+host, drop the user:pass@ segment.
    s = _CREDENTIAL_URL_RE.sub(lambda m: m.group(1) + _placeholder("credential_url") + "@", s)

    # Remaining fixed-shape secrets / PII.
    for pattern, kind in _PATTERNS[1:]:
        s = pattern.sub(_placeholder(kind), s)

    # High-entropy 40-char base64 runs (AWS secret-key shape), entropy-gated.
    def _maybe_secret(m: re.Match) -> str:
        tok = m.group(0)
        return (
            _placeholder("high_entropy_secret")
            if _shannon_entropy(tok) >= _SECRET40_MIN_ENTROPY
            else tok
        )

    s = _SECRET40_RE.sub(_maybe_secret, s)

    # Operator-supplied extra patterns.
    for pattern in extra:
        s = pattern.sub(_placeholder("custom"), s)

    # Size cap last (defends against a pathological paste bloating the store).
    if max_chars and len(s) > max_chars:
        s = s[:max_chars] + _placeholder("truncated")

    return s


def _walk(node: Any, max_chars: int, extra: list[re.Pattern]) -> Any:
    """Recursively scrub every string in a nested dict/list structure in place."""
    if isinstance(node, str):
        return _scrub_string(node, max_chars, extra)
    if isinstance(node, dict):
        for k in list(node.keys()):
            node[k] = _walk(node[k], max_chars, extra)
        return node
    if isinstance(node, list):
        for i, v in enumerate(node):
            node[i] = _walk(v, max_chars, extra)
        return node
    if isinstance(node, tuple):
        return tuple(_walk(v, max_chars, extra) for v in node)
    return node


# --------------------------------------------------------------------------- #
# Public API
# --------------------------------------------------------------------------- #
def redact_trace(trace: dict) -> dict:
    """Return a deep-copied, redacted copy of ``trace`` (pure; input unmutated).

    Walks every string value under ``messages`` and ``trace_logs`` (and any other
    key present) and replaces secrets / credentials / PII with typed placeholders.
    When ``TRACE_REDACTION=off`` the trace is returned unchanged (local debug only).

    Raises only on a genuinely malformed input; callers on the persist path use
    :func:`redact_trace_safe`, which fails closed to a minimal trace instead.
    """
    if not _redaction_enabled():
        return trace
    redacted = copy.deepcopy(trace)
    return _walk(redacted, _max_field_chars(), _extra_patterns())


def minimal_trace(run_id: str | None, reason: str = "redaction_error") -> dict:
    """A stripped trace written when redaction fails — carries no raw payload."""
    return {
        "run_id": run_id,
        "execution_time": time.strftime("%Y-%m-%d %H:%M:%S"),
        "redaction_error": reason,
        "messages": [],
        "trace_logs": [],
    }


def redact_trace_safe(trace: dict, run_id: str | None = None) -> dict:
    """Redact ``trace``, failing **closed** to a minimal trace on any error.

    This is the wrapper the always-on persist path and the file dump call: it
    never returns the raw payload if scrubbing did not complete, and it never
    raises (so the surrounding fail-open logging wrapper still governs crashes).
    """
    try:
        return redact_trace(trace)
    except Exception as e:  # noqa: BLE001 — never persist a raw trace on redaction failure
        print(f"[trace_redaction] redaction failed, writing minimal trace: {e}")
        return minimal_trace(run_id, reason=f"redaction_error: {type(e).__name__}")