import json

import lief

# wave to parse
with open("rsc/wave-0001.dump", "rb") as f:
    pe = lief.parse(f)
    assert isinstance(pe, lief.PE.Binary)

with open("rsc/upx-hostname.exe.bin_iat_wave1.json", "r") as iat_json_input:
    iat_data = json.load(iat_json_input)
calls:list[dict[str,str]] = iat_data["calls"]
wave_entry = int(iat_data["entry"],16)

# print(pe.rich_header)

# for section in pe.sections:
#     print(section.name, len(section.content))

# patch entrypoint
entrypoint_format = int(hex(wave_entry)[-4:],16)
pe.optional_header.addressof_entrypoint = entrypoint_format

# create new iat section
section = lief.PE.Section(".patchiat")
section.content = [0xCC] * 0x100
pe.add_section(section)

# write result
pe.write("patched.exe")