repro-bundle-v0/scripts/prepare_cybench_tasks.py

"""Extract a subset of self-contained (no docker/network) CyBench tasks into
data/cybench_tasks/<slug>/{challenge/*, task.json}.

Each task's `challenge/` directory replicates what the CyBench `init_script.sh`
would drop into the agent's working directory (we replicate that copy logic
per-task below, since we don't run docker). `task.json` holds the prompt and
gold flag used only by the grading harness -- never exposed inside
`challenge/`.

Usage: python3 scripts/prepare_cybench_tasks.py --cybench-repo /tmp/cybench_probe
"""

from __future__ import annotations

import argparse
import json
import shutil
import zipfile
from pathlib import Path

# (slug, relative task dir under benchmark/, extraction kind, extra)
TASKS = [
    ("crypto_dynastic", "hackthebox/cyber-apocalypse-2024/crypto/[Very Easy] Dynastic", "copy_challenge", None),
    ("crypto_makeshift", "hackthebox/cyber-apocalypse-2024/crypto/[Very Easy] Makeshift", "copy_challenge", None),
    (
        "crypto_primary_knowledge",
        "hackthebox/cyber-apocalypse-2024/crypto/[Very Easy] Primary Knowledge",
        "copy_challenge",
        None,
    ),
    (
        "crypto_partial_tenacity",
        "hackthebox/cyber-apocalypse-2024/crypto/[Medium] Partial Tenacity",
        "copy_challenge",
        None,
    ),
    ("crypto_diffecient", "project-sekai-ctf/sekaictf-2022/crypto/diffecient", "copy_app", None),
    # NOTE: "forensics_it_has_begun" was dropped -- its challenge file
    # (recovered_log.sh) mimics ransomware/attacker persistence commands as
    # part of the CTF's forensics design, which Windows Defender intermittently
    # quarantines on the WSL /mnt/d mount (PermissionError on read/write
    # unrelated to our code). Using 9 tasks instead of 10.
    (
        "forensics_urgent",
        "hackthebox/cyber-apocalypse-2024/forensics/[Very Easy] Urgent",
        "copy_release_zip",
        "forensics_urgent.zip",
    ),
    (
        "reversing_boxcutter",
        "hackthebox/cyber-apocalypse-2024/reversing/[Very Easy] BoxCutter",
        "unzip_release_flatten",
        ("rev_boxcutter.zip", "rev_boxcutter", "cutter"),
    ),
    (
        "reversing_lootstash",
        "hackthebox/cyber-apocalypse-2024/reversing/[Very Easy] LootStash",
        "unzip_release_flatten",
        ("rev_lootstash.zip", "rev_lootstash", "stash"),
    ),
    (
        "reversing_crushing",
        "hackthebox/cyber-apocalypse-2024/reversing/[Easy] Crushing",
        "unzip_release_flatten_multi",
        ("rev_crushing.zip", "rev_crushing", ["crush", "message.txt.cz"]),
    ),
]


def build_task(cybench_repo: Path, out_root: Path, slug: str, rel_dir: str, kind: str, extra):
    src = cybench_repo / "benchmark" / rel_dir
    meta = json.loads((src / "metadata" / "metadata.json").read_text())
    prompt = meta["easy_prompt"].strip()
    if slug == "forensics_it_has_begun":
        prompt = prompt.replace("script.sh", "recovered_log.sh")
    flag_subtask = meta["subtasks"][-1]
    assert "flag" in flag_subtask["question"].lower(), flag_subtask["question"]
    flag = flag_subtask["answer"].strip()

    out_dir = out_root / slug
    if out_dir.exists():
        shutil.rmtree(out_dir)
    challenge_dir = out_dir / "challenge"
    challenge_dir.mkdir(parents=True)

    if kind == "copy_challenge":
        for f in (src / "challenge").iterdir():
            shutil.copy2(f, challenge_dir / f.name)
    elif kind == "copy_app":
        shutil.copy2(src / "challenge" / "app" / "diffecient.py", challenge_dir / "diffecient.py")
    elif kind == "copy_release_zip":
        zip_name = extra
        shutil.copy2(src / "release" / zip_name, challenge_dir / zip_name)
    elif kind == "unzip_release":
        zip_name = extra
        with zipfile.ZipFile(src / "release" / zip_name) as zf:
            zf.extractall(challenge_dir)
        # Windows/DrvFs (WSL /mnt/d) can wedge on the literal name "script.sh" at
        # this path (persistent PermissionError unrelated to our code); rename to
        # dodge it. Prompt text below is adjusted to match.
        script_sh = challenge_dir / "script.sh"
        if script_sh.exists():
            script_sh.rename(challenge_dir / "recovered_log.sh")
    elif kind == "unzip_release_flatten":
        zip_name, subdir, keep_file = extra
        tmp = out_dir / "_tmp_unzip"
        with zipfile.ZipFile(src / "release" / zip_name) as zf:
            zf.extractall(tmp)
        shutil.copy2(tmp / subdir / keep_file, challenge_dir / keep_file)
        shutil.rmtree(tmp)
    elif kind == "unzip_release_flatten_multi":
        zip_name, subdir, keep_files = extra
        tmp = out_dir / "_tmp_unzip"
        with zipfile.ZipFile(src / "release" / zip_name) as zf:
            zf.extractall(tmp)
        for kf in keep_files:
            shutil.copy2(tmp / subdir / kf, challenge_dir / kf)
        shutil.rmtree(tmp)
    else:
        raise ValueError(kind)

    task_json = {
        "name": slug,
        "source": f"cybench/{rel_dir}",
        "category": meta.get("categories", ["?"])[0],
        "difficulty": meta.get("difficulty", "?"),
        "prompt": prompt,
        "flag": flag,
    }
    (out_dir / "task.json").write_text(json.dumps(task_json, indent=2))
    n_files = sum(1 for _ in challenge_dir.rglob("*") if _.is_file())
    print(f"  built {slug}: {n_files} challenge file(s), flag={flag!r}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cybench-repo", required=True, type=str)
    ap.add_argument("--out", type=str, default=None)
    args = ap.parse_args()
    cybench_repo = Path(args.cybench_repo)
    out_root = Path(args.out) if args.out else Path(__file__).resolve().parent.parent / "data" / "cybench_tasks"
    out_root.mkdir(parents=True, exist_ok=True)
    for slug, rel_dir, kind, extra in TASKS:
        print(f"building {slug} ...")
        build_task(cybench_repo, out_root, slug, rel_dir, kind, extra)
    print(f"\nDone. {len(TASKS)} tasks written to {out_root}")


if __name__ == "__main__":
    main()