#!/usr/bin/env python3
"""Binary exploitation analysis agent.

# For authorized security testing and CTF challenges only

Analyzes ELF binaries for security mitigations, discovers ROP gadgets,
and assists exploit development using pwntools and checksec.
"""

import argparse
import json
import subprocess
import sys
import datetime

try:
    from pwn import ELF, ROP
    HAS_PWNTOOLS = True
except ImportError:
    HAS_PWNTOOLS = False


def run_checksec(binary_path):
    """Analyze binary security mitigations using checksec."""
    if HAS_PWNTOOLS:
        try:
            elf = ELF(binary_path, checksec=False)
            return {
                "arch": elf.arch,
                "bits": elf.bits,
                "endian": elf.endian,
                "nx": elf.nx,
                "pie": elf.pie,
                "canary": elf.canary,
                "relro": "Full" if elf.relro == "Full" else ("Partial" if elf.relro else "None"),
                "stripped": not elf.sym,
                "static": elf.statically_linked,
            }
        except Exception as e:
            return {"error": str(e)}
    try:
        result = subprocess.run(["checksec", "--file", binary_path, "--output", "json"],
                                capture_output=True, text=True, timeout=10)
        if result.stdout:
            return json.loads(result.stdout)
    except (FileNotFoundError, subprocess.TimeoutExpired, json.JSONDecodeError):
        pass
    return {"error": "Neither pwntools nor checksec available"}


def find_rop_gadgets(binary_path, max_gadgets=20):
    """Find ROP gadgets using pwntools or ROPgadget."""
    if HAS_PWNTOOLS:
        try:
            elf = ELF(binary_path, checksec=False)
            rop = ROP(elf)
            gadgets = []
            for gadget in rop.gadgets.values():
                if len(gadgets) >= max_gadgets:
                    break
                gadgets.append({
                    "address": hex(gadget.address),
                    "insns": "; ".join(gadget.insns),
                })
            return gadgets
        except Exception as e:
            return [{"error": str(e)}]
    try:
        result = subprocess.run(
            ["ROPgadget", "--binary", binary_path, "--count", str(max_gadgets)],
            capture_output=True, text=True, timeout=30
        )
        gadgets = []
        for line in result.stdout.splitlines():
            if " : " in line:
                parts = line.split(" : ", 1)
                gadgets.append({"address": parts[0].strip(), "insns": parts[1].strip()})
        return gadgets[:max_gadgets]
    except (FileNotFoundError, subprocess.TimeoutExpired):
        return [{"error": "Neither pwntools ROP nor ROPgadget available"}]


def find_useful_functions(binary_path):
    """Find useful functions for exploitation (system, exec, write, etc.)."""
    if not HAS_PWNTOOLS:
        return {"error": "pwntools not available"}
    try:
        elf = ELF(binary_path, checksec=False)
        interesting = ["system", "execve", "exec", "popen", "gets", "strcpy",
                       "sprintf", "read", "write", "puts", "printf", "mprotect"]
        found = {}
        for func in interesting:
            addr = elf.sym.get(func) or elf.plt.get(func)
            if addr:
                found[func] = hex(addr)
        got_entries = {}
        for name in ["system", "printf", "puts", "__libc_start_main"]:
            if name in elf.got:
                got_entries[name] = hex(elf.got[name])
        return {"functions": found, "got_entries": got_entries}
    except Exception as e:
        return {"error": str(e)}


def find_vulnerable_functions(binary_path):
    """Identify potentially vulnerable functions in the binary."""
    dangerous = {"gets": "Unbounded read - guaranteed buffer overflow",
                 "strcpy": "No length check - possible overflow",
                 "strcat": "No length check - possible overflow",
                 "sprintf": "No length check - possible overflow",
                 "scanf": "Possible format string / overflow",
                 "vsprintf": "No length check - possible overflow"}
    if not HAS_PWNTOOLS:
        return {"error": "pwntools not available"}
    try:
        elf = ELF(binary_path, checksec=False)
        found = []
        for func, reason in dangerous.items():
            if func in elf.plt or func in elf.sym:
                found.append({"function": func, "reason": reason,
                              "address": hex(elf.plt.get(func, elf.sym.get(func, 0)))})
        return found
    except Exception as e:
        return [{"error": str(e)}]


def analyze_binary(binary_path):
    """Full binary exploitation analysis."""
    report = {
        "binary": binary_path,
        "timestamp": datetime.datetime.utcnow().isoformat() + "Z",
        "checksec": run_checksec(binary_path),
        "dangerous_functions": find_vulnerable_functions(binary_path),
        "useful_functions": find_useful_functions(binary_path),
        "rop_gadgets": find_rop_gadgets(binary_path, max_gadgets=15),
    }
    mitigations = report["checksec"]
    if isinstance(mitigations, dict) and "error" not in mitigations:
        report["exploit_difficulty"] = "HARD" if all([
            mitigations.get("nx"), mitigations.get("pie"),
            mitigations.get("canary"), mitigations.get("relro") == "Full"
        ]) else "MEDIUM" if mitigations.get("nx") else "EASY"
    return report


def main():
    parser = argparse.ArgumentParser(
        description="Binary exploitation analysis agent (authorized testing only)"
    )
    parser.add_argument("binary", nargs="?", help="Path to ELF binary")
    parser.add_argument("--checksec-only", action="store_true", help="Only run checksec")
    parser.add_argument("--gadgets", type=int, default=15, help="Max ROP gadgets to find")
    parser.add_argument("--output", "-o", help="Output JSON report path")
    args = parser.parse_args()

    print("[*] Binary Exploitation Analysis Agent")
    print("[*] For authorized security testing and CTF challenges only")
    print(f"    pwntools available: {HAS_PWNTOOLS}")

    if not args.binary:
        print("\nUsage: python agent.py /path/to/binary [--checksec-only] [--gadgets 20]")
        print("  Analyzes: mitigations, dangerous functions, ROP gadgets, GOT entries")
        print(json.dumps({"demo": True, "pwntools": HAS_PWNTOOLS}, indent=2))
        sys.exit(0)

    if args.checksec_only:
        result = run_checksec(args.binary)
        print(json.dumps(result, indent=2))
        sys.exit(0)

    report = analyze_binary(args.binary)
    checksec = report.get("checksec", {})
    if isinstance(checksec, dict) and "error" not in checksec:
        print(f"\n[*] Architecture: {checksec.get('arch')} ({checksec.get('bits')}-bit)")
        print(f"    NX: {checksec.get('nx')} | PIE: {checksec.get('pie')} | "
              f"Canary: {checksec.get('canary')} | RELRO: {checksec.get('relro')}")
        print(f"    Exploit difficulty: {report.get('exploit_difficulty', '?')}")

    dangerous = report.get("dangerous_functions", [])
    if isinstance(dangerous, list) and dangerous:
        print(f"\n[!] Dangerous functions found: {len(dangerous)}")
        for d in dangerous:
            if "error" not in d:
                print(f"    {d['function']} @ {d['address']}: {d['reason']}")

    gadgets = report.get("rop_gadgets", [])
    if gadgets and "error" not in gadgets[0]:
        print(f"\n[*] ROP gadgets found: {len(gadgets)}")
        for g in gadgets[:5]:
            print(f"    {g['address']}: {g['insns']}")

    if args.output:
        with open(args.output, "w") as f:
            json.dump(report, f, indent=2)
    print(json.dumps({"difficulty": report.get("exploit_difficulty", "unknown"),
                       "gadgets": len(gadgets)}, indent=2))


if __name__ == "__main__":
    main()
