#!/usr/bin/env python3
"""Generate an angr symbolic-execution harness (and report angr availability).

`check` reports whether the angr package is importable. `harness` writes a ready-to-edit angr
script that loads a binary, explores to a target address while avoiding another, and prints the
solving input. This tool only generates text — it does not load the binary or run the sample.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

try:
    import angr  # type: ignore  # noqa: F401
    HAVE_ANGR = True
except ImportError:
    HAVE_ANGR = False


def check() -> dict:
    return {"angr_available": HAVE_ANGR,
            "note": "pip install angr to run generated harnesses" if not HAVE_ANGR else "ready"}


def harness(binary: str, find: str, avoid: str | None, out: str) -> dict:
    avoid_line = (f"    avoid_addr = {avoid}\n" if avoid else "    avoid_addr = []\n")
    explore = ("simgr.explore(find=find_addr, avoid=avoid_addr)"
               if avoid else "simgr.explore(find=find_addr)")
    text = (
        "#!/usr/bin/env python3\n"
        "# angr symbolic-execution harness (generated skeleton). Run in an isolated VM.\n"
        "import angr\n"
        "import claripy\n\n"
        "def solve():\n"
        f"    proj = angr.Project({binary!r}, auto_load_libs=False)\n"
        f"    find_addr = {find}\n"
        f"{avoid_line}"
        "    # Symbolic stdin of bounded length; tighten the charset/length below.\n"
        "    flag_len = 32\n"
        "    flag = claripy.BVS('flag', 8 * flag_len)\n"
        "    state = proj.factory.full_init_state(stdin=flag)\n"
        "    for i in range(flag_len):\n"
        "        b = flag.get_byte(i)\n"
        "        state.solver.add(b >= 0x20, b <= 0x7e)  # printable\n"
        "    simgr = proj.factory.simulation_manager(state)\n"
        "    simgr.use_technique(angr.exploration_techniques.Veritesting())\n"
        f"    {explore}\n"
        "    if simgr.found:\n"
        "        found = simgr.found[0]\n"
        "        print(found.posix.dumps(0))  # the solving input\n"
        "    else:\n"
        "        print('no solution reached the target')\n\n"
        "if __name__ == '__main__':\n"
        "    solve()\n"
    )
    Path(out).write_text(text, encoding="utf-8")
    return {"output": out, "find": find, "avoid": avoid, "angr_available": HAVE_ANGR}


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Generate an angr harness")
    sub = parser.add_subparsers(dest="cmd", required=True)
    sub.add_parser("check")
    ph = sub.add_parser("harness")
    ph.add_argument("--binary", required=True)
    ph.add_argument("--find", required=True, help="target address, e.g. 0x401234")
    ph.add_argument("--avoid", help="address to avoid, e.g. 0x401260")
    ph.add_argument("--out", required=True)
    args = parser.parse_args(argv)
    if args.cmd == "check":
        print(json.dumps(check(), indent=2))
    else:
        print(json.dumps(harness(args.binary, args.find, args.avoid, args.out), indent=2))
    return 0


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