#!/usr/bin/env python3

from pwn import *
import sys
import stat
from pathlib import Path
import yaml

nop_length = 256
e = ELF("challenge", checksec=False)
context.binary = e
nops_offset = list(e.search(b"\x90"*nop_length))[0]


def main():
    calls = yaml.safe_load(open("patches.yaml"))

    assembly = ""
    asm_len = 0

    for call in calls:
        args = call[1]
        addr = call[0]
        tmp = ""

        if len(args) >= 1:
            tmp += f"mov rdi, {hex(args[0])}\n"
        if len(args) >= 2:
            tmp += f"mov rsi, {hex(args[1])}\n"
        if len(args) >= 3:
            tmp += f"mov rdx, {hex(args[2])}\n"
        if len(args) >= 4:
            tmp += f"mov rcx, {hex(args[3])}\n"
        if len(args) >= 5:
            tmp += f"mov r8, {hex(args[4])}\n"
        if len(args) >= 6:
            tmp += f"mov r9, {hex(args[5])}\n"
        args = args[6:]
        args = args[::-1]
        while len(args) > 0:
            tmp += f"push {hex(args[0])}\n"
            args = args[1:]

        asm_len += len(asm(tmp)) + 5
        assembly += tmp + f"call {addr}\n"

    if asm_len > nop_length:
        print("Too much calls/arguments, not enough space")
        exit(1)

    e.asm(nops_offset, assembly)
    e.save("patched")

    p = Path("patched")
    p.chmod(p.stat().st_mode | stat.S_IEXEC)


if __name__ == "__main__":
    main()
    print("Patched version created. Run './patched' to see the results.")