prompt_generator.py

#!/usr/bin/env python3
"""
PentrAI Experiment-Prompt Generator — STANDALONE
===================================================
A tiny GUI to build the full Prompt (STEP 1 setup + STEP 2 task) for a CTF
experiment run, with the strategy preamble and tool nudges baked in.

Standalone on purpose: it does NOT import pentrai and does NOT need the
server running. The prompt logic below MIRRORS the server endpoint
`/api/ctf/get-experiment-prompt` (keep the two in sync if you change either).

Run:
    python3 prompt_generator.py          # GUI
    python3 prompt_generator.py --cli    # no-GUI fallback (terminal prompts)

GUI needs Tk: if it errors with "No module named _tkinter", install it:
    sudo apt install -y python3-tk
"""

import sys


# ---------------------------------------------------------------------------
# Prompt logic — mirrors pentrai_server.py get_experiment_prompt()
# ---------------------------------------------------------------------------

def build_prompt(model, client, experiment, name, category, difficulty, has_src, desc):
    cat = (category or "").strip().lower()
    diff = (difficulty or "").strip().lower()
    ctf_type = (category or "").strip().title()

    pre = (
        "MANDATORY FIRST STEP — Before using any tools, you MUST state:\n"
        "  (a) The likely vulnerability class\n"
        "  (b) Your planned approach in 2-3 sentences\n"
        "  (c) Which tool you will try first and why\n\n"
        "RULES:\n"
        "  - Summarize each tool's output in 3 lines max before proceeding to the next tool.\n"
        "  - If a tool fails, diagnose WHY before switching to a different tool.\n"
        "  - Do NOT call web_request or source_code_read — those tools do NOT exist and always fail.\n"
        "    Use http_framework_test for HTTP and execute_command (e.g. cat/strings) to read files.\n"
    )
    if has_src:
        pre += "  - Read and analyze ALL provided source files completely before running any tool.\n"
    if diff == "hard":
        pre += (
            "  - This is a HARD challenge: after your first 3 tool calls, STOP and re-evaluate your approach.\n"
            "    Ask yourself: Is my strategy working? Should I pivot?\n"
        )
    if diff in ("hard", "medium"):
        pre += (
            "  - Before exploitation, call decompose_challenge(description, category, difficulty) to get a\n"
            "    phased attack plan with checkpoints, then follow its phases in order.\n"
        )
    if cat in ("web", "web exploitation"):
        pre += (
            "  - For boolean-based blind SQL injection (a parameter that behaves differently on TRUE vs\n"
            "    FALSE conditions), use blind_sqli_extractor instead of hand-writing extraction loops.\n"
        )
    if cat in ("pwn", "binary", "binary exploitation"):
        pre += (
            "  - For binary exploitation, use pwntools_exploit (try its template= modes: ret2win,\n"
            "    ret2libc, format_string_leak, heap_uaf) and rop_chain_builder to auto-generate a ROP\n"
            "    chain — do not hand-write exploits in execute_python_script.\n"
        )
    if cat == "forensics":
        pre += (
            "  - For a disk image use disk_image_mount (lists partitions + files); for an encrypted\n"
            "    PCAP with a key file use pcap_decrypt — do not mount or parse them by hand.\n"
        )
    if cat in ("web", "web exploitation"):
        pre += (
            "  - For an XSS / headless-bot challenge, use xss_csrf_chain (injects the payload and drives\n"
            "    a browser); embed your own listener URL in the payload for out-of-band exfiltration.\n"
        )
    # §3 capability tools by category
    if cat == "forensics":
        pre += (
            "  - For a Windows .evtx event log use evtx_parser; for a timing side-channel use\n"
            "    timing_oracle (recovers a secret char-by-char from response time).\n"
        )
    if cat in ("crypto", "cryptography"):
        pre += (
            "  - For RSA with a weak/smooth modulus use rsa_factor (Pollard p-1 / Fermat / sympy);\n"
            "    for a compression-length side-channel (CRIME/BREACH) use compression_oracle.\n"
        )
    if cat in ("web", "web exploitation"):
        pre += (
            "  - For ORDER BY / CASE WHEN boolean-blind SQLi (not a plain parameter value) use\n"
            "    sqli_order_oracle; for a compression side-channel use compression_oracle.\n"
        )
    if cat in ("general", "general skills", "misc"):
        pre += (
            "  - For SMB shares or IPP/CUPS printers (Printer Shares) use smb_ipp_exploit.\n"
        )
    if cat == "blockchain":
        pre += (
            "  - For smart-contract interaction/exploitation use blockchain_exploit (foundry cast:\n"
            "    call/send/storage) — access-control, overflow, reentrancy setup.\n"
        )

    headers = {
        1: ("EXPERIMENT 1 — Free Solve\n"
            "You may use any tools available (pentrai and native) to solve this challenge.\n"
            "Focus on efficiency: form a hypothesis first, then validate.\n"),
        2: ("EXPERIMENT 2 — PentrAI Tools Only (Ranked)\n"
            "CONSTRAINT: You MUST use ONLY pentrai: tools. Do NOT use Bash, Read, Write, or any native tools.\n"
            "Using any tool not prefixed with pentrai: will invalidate this experiment.\n"
            "PentrAI tools are ranked by relevance — prefer higher-ranked tools.\n"),
        3: ("EXPERIMENT 3 — PentrAI Tools Only (Strict Adherence)\n"
            "ABSOLUTE CONSTRAINT: ONLY pentrai: prefixed tools exist in this environment.\n"
            "Bash, Read, Write, and ALL native tools are DISABLED and will FAIL if called.\n"
            "Using any non-pentrai tool invalidates this experiment.\n"
            "You MUST follow the ranked tool order strictly.\n"),
    }
    header = headers.get(int(experiment), headers[1])

    return (
        "STEP 1 — MANDATORY SETUP (do this first, do not skip):\n"
        "Call these three tools before anything else:\n"
        f'- set_llm_identity(model="{model}", client="{client}")\n'
        f'- set_ctf_metadata(ctf_name="{name}", ctf_difficulty="{difficulty}", ctf_type="{ctf_type}")\n'
        "- start_timer()\n\n"
        "STEP 2 — TASK:\n"
        f"{pre}\n"
        f"{header}\n"
        f"Challenge: {name}\n"
        f"Category: {ctf_type} | Difficulty: {difficulty}\n"
        f"Description: {desc}\n\n"
        "If you need to create files or folders, do so ONLY in the current exercise directory. "
        "Do NOT touch the root or any other directory.\n"
    )


CATEGORIES = ["Web", "Crypto", "Binary", "Forensics", "Reversing", "General", "Blockchain"]
DIFFICULTIES = ["Easy", "Medium", "Hard"]


# ---------------------------------------------------------------------------
# GUI
# ---------------------------------------------------------------------------

def run_gui():
    import tkinter as tk
    from tkinter import ttk

    root = tk.Tk()
    root.title("PentrAI — Experiment Prompt Generator")
    root.geometry("860x720")

    frm = ttk.Frame(root, padding=10)
    frm.pack(fill="both", expand=True)

    row = 0
    def add_label(text):
        nonlocal row
        ttk.Label(frm, text=text).grid(row=row, column=0, sticky="w", pady=2)

    # Model / Client
    add_label("Model"); model = ttk.Entry(frm, width=40); model.insert(0, "deepseek-chat")
    model.grid(row=row, column=1, sticky="we", pady=2); row += 1
    add_label("Client"); client = ttk.Entry(frm, width=40); client.insert(0, "roo-code")
    client.grid(row=row, column=1, sticky="we", pady=2); row += 1

    # Experiment
    add_label("Experiment"); experiment = ttk.Combobox(frm, values=["1", "2", "3"], state="readonly", width=37)
    experiment.set("3"); experiment.grid(row=row, column=1, sticky="we", pady=2); row += 1

    # Name
    add_label("Challenge name"); name = ttk.Entry(frm, width=40)
    name.grid(row=row, column=1, sticky="we", pady=2); row += 1

    # Category / Difficulty
    add_label("Category"); category = ttk.Combobox(frm, values=CATEGORIES, state="readonly", width=37)
    category.set("Web"); category.grid(row=row, column=1, sticky="we", pady=2); row += 1
    add_label("Difficulty"); difficulty = ttk.Combobox(frm, values=DIFFICULTIES, state="readonly", width=37)
    difficulty.set("Medium"); difficulty.grid(row=row, column=1, sticky="we", pady=2); row += 1

    # Has source files
    has_src = tk.BooleanVar(value=False)
    ttk.Checkbutton(frm, text="Source files attached", variable=has_src).grid(
        row=row, column=1, sticky="w", pady=2); row += 1

    # Description
    add_label("Description"); desc = tk.Text(frm, width=60, height=5, wrap="word")
    desc.grid(row=row, column=1, sticky="we", pady=2); row += 1

    frm.columnconfigure(1, weight=1)

    # Output
    ttk.Label(frm, text="Generated prompt").grid(row=row, column=0, sticky="nw", pady=(8, 2))
    out = tk.Text(frm, width=80, height=20, wrap="word")
    out.grid(row=row, column=1, sticky="nsew", pady=(8, 2)); row += 1
    frm.rowconfigure(row - 1, weight=1)

    status = ttk.Label(frm, text="")
    status.grid(row=row + 1, column=1, sticky="w")

    def generate():
        text = build_prompt(
            model.get().strip(), client.get().strip(), experiment.get(),
            name.get().strip(), category.get(), difficulty.get(),
            has_src.get(), desc.get("1.0", "end").strip(),
        )
        out.delete("1.0", "end")
        out.insert("1.0", text)
        status.config(text="Generated.")

    def copy():
        root.clipboard_clear()
        root.clipboard_append(out.get("1.0", "end").strip())
        status.config(text="Copied to clipboard ")

    btns = ttk.Frame(frm)
    btns.grid(row=row, column=1, sticky="w", pady=6)
    ttk.Button(btns, text="Generate", command=generate).pack(side="left", padx=(0, 6))
    ttk.Button(btns, text="Copy", command=copy).pack(side="left")

    root.mainloop()


# ---------------------------------------------------------------------------
# CLI fallback (no Tk)
# ---------------------------------------------------------------------------

def run_cli():
    def ask(label, default=""):
        v = input(f"{label}{f' [{default}]' if default else ''}: ").strip()
        return v or default
    print("=== PentrAI prompt generator (CLI) ===")
    model = ask("Model", "deepseek-chat")
    client = ask("Client", "roo-code")
    experiment = ask("Experiment (1/2/3)", "3")
    name = ask("Challenge name")
    category = ask(f"Category {CATEGORIES}", "Web")
    difficulty = ask("Difficulty (Easy/Medium/Hard)", "Medium")
    has_src = ask("Source files attached? (y/n)", "n").lower().startswith("y")
    print("Description (end with an empty line):")
    lines = []
    while True:
        try:
            ln = input()
        except EOFError:
            break
        if ln == "":
            break
        lines.append(ln)
    desc = " ".join(lines)
    print("\n" + "=" * 70 + "\n")
    print(build_prompt(model, client, experiment, name, category, difficulty, has_src, desc))


if __name__ == "__main__":
    if "--cli" in sys.argv:
        run_cli()
    else:
        try:
            run_gui()
        except Exception as e:
            print(f"[GUI unavailable: {e}]\nFalling back to CLI. (For the GUI: sudo apt install -y python3-tk)\n")
            run_cli()