#!/usr/bin/env python3
"""
scaffold.py — generate a new CTF / task skeleton (p4-team style).

Usage:
  python3 scaffold.py new hitcon "HITCON CTF 2024" 5 700
  python3 scaffold.py add hitcon "Safe Bank" pwn
"""
from pathlib import Path
from datetime import datetime
import argparse
import re

TEMPLATES = {
    "pwn": "pwn_skeleton.py",
    "crypto": "crypto_skeleton.py",
    "web": "web_skeleton.py",
    "re": "re_skeleton.py",
    "reverse": "re_skeleton.py",
    "forensics": "forensics_skeleton.py",
    "misc": "misc_skeleton.py",
}

ROOT = Path(__file__).parent


def new_ctf(args):
    stamp = datetime.now().isoformat()[:10]
    dirname = f"{stamp}-{args.slug}"
    newdir = ROOT / dirname
    newdir.mkdir(exist_ok=True)
    (newdir / "README.md").write_text(f"# {args.name}\n\n### Table of contents\n")
    print(f"[+] created {dirname}/")


def add_task(args):
    matches = sorted(ROOT.glob(f"*-{args.slug}"))
    if not matches:
        print("no such CTF slug"); return
    d = matches[-1]
    slug = re.sub(r"[^a-z0-9]+", "-", args.task.lower())
    taskdir = d / slug
    taskdir.mkdir(exist_ok=True)
    tmpl = TEMPLATES.get(args.category, "pwn_skeleton.py")
    src = (ROOT / "templates" / tmpl).read_text()
    (taskdir / "solve.py").write_text(src)
    rd = d / "README.md"
    rd.write_text(rd.read_text() + f"* [{args.task} ({args.category})]({slug})\n")
    print(f"[+] added {taskdir}/solve.py")


def main():
    p = argparse.ArgumentParser()
    sub = p.add_subparsers(dest="cmd", required=True)
    n = sub.add_parser("new"); n.add_argument("slug"); n.add_argument("name"); n.add_argument("place"); n.add_argument("teams"); n.set_defaults(func=new_ctf)
    a = sub.add_parser("add"); a.add_argument("slug"); a.add_argument("task"); a.add_argument("category"); a.set_defaults(func=add_task)
    args = p.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()