scrub.py

"""Scrub raw JSONL records before publication.

Replaces real account IDs, ARNs, bucket names, and instance IDs with
stable placeholders so artefacts can be committed and shared safely.

Nothing in this file hardcodes the author's account. The mapping is
built on the fly by:

- detecting any 12-digit number that looks like an AWS account ID and
  replacing it with ``<ACCOUNT_ID>``;
- detecting any ``arn:aws:iam::<digits>:user/<name>`` and replacing it
  with ``<USER_ARN>``;
- detecting instance IDs (``i-`` + 17 hex) and mapping each unique one
  to ``i-instance01``, ``i-instance02``, ... stably across the file;
- detecting bucket names via an explicit ``--buckets`` CSV option, or a
  ``BENCH_BUCKETS`` env var, and remapping them to ``bucket-1``,
  ``bucket-2``, ... . The real bucket names are never stored in this
  file or in the repository.

Usage:
    python -m scripts.scrub INPUT_JSONL OUTPUT_JSONL [--buckets a,b,c]

The list of real bucket names to scrub is expected to come from outside
the repo. Typically:
    BENCH_BUCKETS="$(aws s3api list-buckets --query 'Buckets[].Name' \\
        --output text | tr '\\t' ',')" \\
        python -m scripts.scrub raw.jsonl scrubbed.jsonl
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
from pathlib import Path
from typing import Any

ACCOUNT_RE = re.compile(r"\b\d{12}\b")
INSTANCE_RE = re.compile(r"\bi-[0-9a-f]{17}\b")
ARN_USER_RE = re.compile(r"arn:aws:iam::\d+:user/[\w.-]+")


def build_bucket_map(buckets: list[str]) -> dict[str, str]:
    # Longest first so prefix-overlapping names (e.g. ``bkt`` vs
    # ``bkt-new``) get replaced in the right order.
    return {
        name: f"bucket-{i + 1}"
        for i, name in enumerate(sorted(buckets, key=len, reverse=True))
    }


def scrub_text(text: str, bucket_map: dict[str, str]) -> str:
    if not text:
        return text
    text = ACCOUNT_RE.sub("<ACCOUNT_ID>", text)
    text = ARN_USER_RE.sub("<USER_ARN>", text)

    for real, replacement in bucket_map.items():
        text = text.replace(real, replacement)

    seen: dict[str, str] = {}

    def _replace_instance(m: re.Match) -> str:
        iid = m.group(0)
        if iid not in seen:
            seen[iid] = f"i-instance{len(seen) + 1:02d}"
        return seen[iid]

    return INSTANCE_RE.sub(_replace_instance, text)


def _scrub_any(value: Any, bucket_map: dict[str, str]) -> Any:
    """Recursively scrub any JSON-shaped value (str, list, dict, or primitive)."""
    if isinstance(value, str):
        return scrub_text(value, bucket_map)
    if isinstance(value, list):
        return [_scrub_any(v, bucket_map) for v in value]
    if isinstance(value, dict):
        return {k: _scrub_any(v, bucket_map) for k, v in value.items()}
    return value


def scrub_record(
    record: dict[str, Any], bucket_map: dict[str, str]
) -> dict[str, Any]:
    # Deep scrub the whole record. Only numeric/boolean/None pass through
    # unchanged — every string anywhere in the tree goes through scrub_text.
    out = _scrub_any(json.loads(json.dumps(record)), bucket_map)
    assert isinstance(out, dict)
    return out


def _load_buckets(explicit: str | None) -> list[str]:
    raw = explicit or os.environ.get("BENCH_BUCKETS", "")
    return [b.strip() for b in raw.split(",") if b.strip()]


def main() -> int:
    parser = argparse.ArgumentParser(description="Scrub raw benchmark JSONL.")
    parser.add_argument("input", help="input JSONL")
    parser.add_argument("output", help="output JSONL")
    parser.add_argument(
        "--buckets",
        default="",
        help=(
            "comma-separated list of bucket names to redact. "
            "Defaults to the BENCH_BUCKETS env var."
        ),
    )
    args = parser.parse_args()

    src = Path(args.input)
    dst = Path(args.output)
    if not src.exists():
        print(f"ERROR: {src} not found", file=sys.stderr)
        return 1

    bucket_map = build_bucket_map(_load_buckets(args.buckets))
    dst.parent.mkdir(parents=True, exist_ok=True)

    count = 0
    with src.open() as fin, dst.open("w") as fout:
        for line in fin:
            line = line.strip()
            if not line:
                continue
            rec = json.loads(line)
            fout.write(json.dumps(scrub_record(rec, bucket_map)) + "\n")
            count += 1
    print(f"wrote {count} scrubbed records -> {dst}")
    if not bucket_map:
        print(
            "WARN: no bucket names provided via --buckets or BENCH_BUCKETS; "
            "buckets in the output were not remapped.",
            file=sys.stderr,
        )
    return 0


if __name__ == "__main__":
    sys.exit(main())