import json
import lief
lief.disable_leak_warning() # warnings to disable for the callback

# Retrives all unique DLL names being imported
def get_used_dlls(calls: list[dict[str,str]]) -> set[str]:
    res = set()
    for call in calls:
        res.add(call["name"].split('!')[0])
    return res

# Retrieves all unique function names used for a single DLL name
def get_used_functions_from_dll(dllname,calls):
    res = set()
    for [dll,func] in map(lambda x: x["name"].split('!'),calls):
        if dll == dllname:
            res.add(func)
    return res
    
def patch_call_to_new_IAT_entry(pe: lief.PE.Binary, call: dict[str,str], rva: int):
    print(call)
    instruction_offset = int(call["adress"],16)-0x1000000
    # We can manually patch the instruction here: FF 15 08 10 00 01 represents `call [0x01001080]`
    # print(hex(rva))
    pe.patch_address(instruction_offset, [0xFF,0x15,int(hex(rva)[4:6],16),int(hex(rva)[2:4],16),0x00,0x01], lief.Binary.VA_TYPES.RVA)
    print([0xFF,0x15,hex(rva)[4:6],hex(rva)[2:4],0x00,0x01])
    # section.content[instruction_offset_from_section+0] = 0xFF
    # section.content[instruction_offset_from_section+1] = 0x15
    # section.content[instruction_offset_from_section+2] = int(hex(rva)[:2],16)
    # section.content[instruction_offset_from_section+3] = int(hex(rva)[:-2],16)
    # section.content[instruction_offset_from_section+4] = 0x00 # TODO: Un-hardcode this!
    # section.content[instruction_offset_from_section+5] = 0x01

def patch_calls_to_new_IAT(pe: lief.PE.Binary, imp: lief.PE.Import, entry:lief.PE.ImportEntry, rva: int):
    print(f"{imp.name}!{entry.name}: 0x{rva:010x}")
    for call in filter(lambda x : x["name"] == f"{imp.name.upper()}!{entry.name}" ,calls):
        patch_call_to_new_IAT_entry(pe,call,rva)
    
    

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

# JSON generated with the python reader files
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)

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

# remove all current imports
pe.remove_all_imports()

# recreate all DLL imports
for dll in get_used_dlls(calls):
    imported_dll = pe.add_import(dll.lower())
    # recreate all function calls related to that dll import
    for func in get_used_functions_from_dll(dll,calls):
        imported_dll.add_entry(func)
        
# At this point, the new IAT will only be constructed when the PE is written. We therefore need to make a callback function to patch calls afterwards.

# write result
config = lief.PE.Builder.config_t()
config.imports = True # allows the config of the writer to write a new IAT
config.resolved_iat_cbk = patch_calls_to_new_IAT # callback after the IAT has been written
pe.write("patched.exe", config)