scrub.py

"""Redact credentials from the agent's trace as it streams to the bucket.

The trace records every command the agent ran and everything those commands printed, and
the Job's environment holds four live tokens — so the trace is one `env` away from being a
credential dump. Filtering in the pipe means raw content never reaches durable storage,
and a partial trace from a killed run is scrubbed too.

Two passes, because they fail differently:
  1. The exact values this Job was handed. Deterministic, and catches a token however it
     was mangled on the way out — split across a header, quoted, URL-embedded.
  2. A shape sweep, for credentials the Job was never given and cannot match by value:
     something the agent read from a file, or a token belonging to someone else entirely.
"""

import os
import re
import sys

SECRET_ENV = ("SLACK_BOT_TOKEN", "GH_TOKEN", "CLAUDE_CODE_OAUTH_TOKEN", "HF_TOKEN")

# Longest first: if one secret is a substring of another, redacting the short one first
# would leave the tail of the long one exposed.
values = sorted(
    (v for k in SECRET_ENV if len(v := os.environ.get(k, "")) > 8), key=len, reverse=True
)

SHAPES = re.compile(
    r"xox[abposr]-[A-Za-z0-9-]{10,}"  # slack
    r"|gh[pousr]_[A-Za-z0-9]{16,}"  # github classic / oauth
    r"|github_pat_[A-Za-z0-9_]{20,}"  # github fine-grained
    r"|sk-ant-[A-Za-z0-9_-]{16,}"  # anthropic api key / oauth token
    r"|hf_[A-Za-z0-9]{16,}"  # hugging face
    r"|AKIA[A-Z0-9]{16}"  # aws access key id
)

REDACTED = "[REDACTED]"

count = 0
for line in sys.stdin:
    for value in values:
        if value in line:
            count += line.count(value)
            line = line.replace(value, REDACTED)
    line, n = SHAPES.subn(REDACTED, line)
    count += n
    sys.stdout.write(line)
    sys.stdout.flush()

# Goes to the run log, where a non-zero count is worth noticing: it means something in the
# pass was handling a credential in a way that reached the transcript.
print(f"[scrub] redacted {count} credential occurrence(s)", file=sys.stderr)