chemenu/tools/chemenu/telemetry/scrub.py

"""Redaction applied to every trace event, in one place: the emitter.

The repo records prompts and assistant replies in cleartext, because a failure
taxonomy cannot be read out of hashes - that is the whole point of the
comprehension phase. Cleartext is only defensible with three guards, all of them
here rather than at each call site:

1. **Secret scrubbing.** Pattern-based, best effort, never a substitute for
   discipline - `reports/` is gitignored and a secret scan runs in verification.
2. **A content cap**, so one 5 MB tool result cannot dominate a trace.
3. **A kill switch.** `WIKI_TRACE_CONTENT=0` drops model-facing text and keeps
   only its length and SHA-256.

`raw/` file *contents* never reach a trace at all, whatever these settings say:
that text is data, not instruction (AGENTS.md invariant 4), and a trace is read
back later. Callers record a path plus a digest instead.
"""
from __future__ import annotations

import hashlib
import os
import re

# Sized like Claude Code's own OTel content limit (60 KiB), which is in turn
# sized for backends that cap an attribute at 64 KiB.
DEFAULT_MAX_CONTENT = 61440

# Attribute names holding model-facing text. Only these obey the kill switch;
# paths, tool names and exit codes stay readable either way.
CONTENT_KEYS = frozenset(
    {"prompt", "response", "message", "text", "tool_input", "tool_output", "error"}
)


def _mask(name: str) -> str:
    return f"[REDACTED:{name}]"


def _keep_key(name: str):
    """Replace the value of a `key: value` pair, keep the key readable."""

    def repl(match: re.Match) -> str:
        return f"{match.group(1)}{match.group(2)}{_mask(name)}"

    return repl


# Order matters: the most specific pattern must match before a generic one can
# swallow part of it.
SECRET_PATTERNS: list[tuple[str, re.Pattern, object]] = [
    (
        "private-key",
        re.compile(r"-----BEGIN [A-Z ]*PRIVATE KEY-----[\s\S]*?-----END [A-Z ]*PRIVATE KEY-----"),
        None,
    ),
    ("github-token", re.compile(r"\bgh[pousr]_[A-Za-z0-9]{16,}\b"), None),
    ("github-pat", re.compile(r"\bgithub_pat_[A-Za-z0-9_]{20,}\b"), None),
    ("aws-access-key", re.compile(r"\b(?:AKIA|ASIA)[0-9A-Z]{16}\b"), None),
    ("slack-token", re.compile(r"\bxox[baprs]-[A-Za-z0-9-]{10,}\b"), None),
    ("anthropic-key", re.compile(r"\bsk-ant-[A-Za-z0-9_-]{20,}\b"), None),
    ("openai-key", re.compile(r"\bsk-[A-Za-z0-9]{20,}\b"), None),
    ("google-api-key", re.compile(r"\bAIza[0-9A-Za-z_-]{35}\b"), None),
    ("onepassword-token", re.compile(r"\bops_[A-Za-z0-9+/=_-]{40,}\b"), None),
    (
        "jwt",
        re.compile(r"\beyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}\b"),
        None,
    ),
    (
        # Stops at a quote, comma or whitespace after the credential: a tool
        # input is usually one line of JSON, and eating to end-of-line would
        # take the rest of the payload with the secret.
        "auth-header",
        re.compile(
            r"""(?ix)
            \b(authorization|x-api-key|proxy-authorization)
            (\s*[:=]\s*)
            (?:bearer|basic|token)?\s*
            [^\s"',;\\]+
            """
        ),
        "keep-key",
    ),
    (
        "secret-assignment",
        re.compile(
            r"""(?ix)
            \b([A-Za-z0-9_.-]*
               (?:api[_-]?key|secret|token|password|passwd|credential)
               [A-Za-z0-9_.-]*)
            (\s*[:=]\s*)
            (["']?[^\s"',;]{8,}["']?)
            """
        ),
        "keep-key",
    ),
]


def max_content() -> int:
    raw = os.environ.get("WIKI_TRACE_MAX_CONTENT")
    if raw and raw.isdigit() and int(raw) > 0:
        return int(raw)
    return DEFAULT_MAX_CONTENT


def content_enabled() -> bool:
    """Cleartext prompts/responses on by default; `WIKI_TRACE_CONTENT=0` opts out."""
    return os.environ.get("WIKI_TRACE_CONTENT", "1") != "0"


def sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8", "replace")).hexdigest()


def scrub_text(text: str) -> tuple[str, list[str]]:
    """Replace known secret shapes. Returns the text and the pattern names hit."""
    hits: list[str] = []
    for name, pattern, mode in SECRET_PATTERNS:
        repl = _keep_key(name) if mode == "keep-key" else _mask(name)
        text, count = pattern.subn(repl, text)
        if count:
            hits.append(name)
    return text, hits


def cap_text(text: str, limit: int | None = None) -> str:
    limit = limit or max_content()
    if len(text) <= limit:
        return text
    dropped = len(text) - limit
    return text[:limit] + f"... [TRUNCATED {dropped} chars]"


def _clean_string(value: str, hits: list[str]) -> str:
    scrubbed, found = scrub_text(value)
    hits.extend(h for h in found if h not in hits)
    return cap_text(scrubbed)


def scrub_attrs(attrs: dict) -> tuple[dict, list[str]]:
    """Walk an attribute tree: scrub every string, cap every string, and apply
    the content kill switch to the keys that hold model-facing text.

    Content keys always gain `<key>_length` and `<key>_sha256` siblings, so two
    traces recorded under different settings stay comparable.
    """
    hits: list[str] = []
    keep_content = content_enabled()

    def walk(value):
        if isinstance(value, dict):
            out = {}
            for key, item in value.items():
                if key in CONTENT_KEYS and isinstance(item, str):
                    out[f"{key}_length"] = len(item)
                    out[f"{key}_sha256"] = sha256_text(item)
                    if keep_content:
                        out[key] = _clean_string(item, hits)
                else:
                    out[key] = walk(item)
            return out
        if isinstance(value, list):
            return [walk(item) for item in value]
        if isinstance(value, str):
            return _clean_string(value, hits)
        return value

    return walk(attrs), hits