From 786decd8ccc2db8ac13ba46bc9bc0fedb87ab6bb Mon Sep 17 00:00:00 2001 From: Seliaste Date: Wed, 25 Mar 2026 15:05:21 +0100 Subject: [PATCH] Small factoring and formatting issues --- iat.py | 85 +++++++++++++++++++++++++++++++++++----------------------- 1 file changed, 52 insertions(+), 33 deletions(-) diff --git a/iat.py b/iat.py index 570855d..36f47c5 100644 --- a/iat.py +++ b/iat.py @@ -1,56 +1,73 @@ import json + import lief -lief.disable_leak_warning() # warnings to disable for the callback + +lief.disable_leak_warning() # warnings to disable for the callback + + +dump_path = "rsc/wave-0001.dump" +# dump_path = "rsc/wave-0002.dump" +iat_json_path = "rsc/upx-hostname.exe.bin_iat_wave1.json" +# iat_json_path = "rsc/000155f2e0360f6ff6cd.exe_iat_wave2.json" + # Retrives all unique DLL names being imported -def get_used_dlls(calls: list[dict[str,str]]) -> set[str]: +def get_used_dlls(calls: list[dict[str, str]]) -> set[str]: res = set() for call in calls: - res.add(call["name"].split('!')[0]) + 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): +def get_used_functions_from_dll(dllname, calls): res = set() - for [dll,func] in map(lambda x: x["name"].split('!'),calls): + 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): + + +def patch_call_to_new_IAT_entry(pe: lief.PE.Binary, call: dict[str, str], rva: int): base = pe.imagebase - instruction_offset = int(call["adress"],16)-base + instruction_offset = int(call["adress"], 16) - base # We can manually patch the instruction here: FF 15 08 10 00 01 represents `call [0x01001080]` - adress_size = 4 if pe.abstract.header.is_32 else 8 is_little_endian = pe.abstract.header.endianness == lief.Header.ENDIANNESS.LITTLE - new_value = [0x00]*adress_size - hex_adress = hex(rva + pe.imagebase)[::-1][:-2] # reversing order and stripping zero - for i in range(0,adress_size): - byte_str = hex_adress[i*2:(i+1)*2][::-1] - new_value[i] += int(byte_str,16) - if(not is_little_endian): - new_value = new_value[::-1] # reverse byte order for big endian - pe.patch_address(instruction_offset, [0xFF,0x15] + new_value, lief.Binary.VA_TYPES.RVA) - + new_value = [0x00] * adress_size + hex_adress = hex(rva + pe.imagebase)[::-1][ + :-2 + ] # reversing order and stripping zero + for i in range(0, adress_size): + byte_str = hex_adress[i * 2 : (i + 1) * 2][::-1] + new_value[i] += int(byte_str, 16) + if not is_little_endian: + new_value = new_value[::-1] # reverse byte order for big endian + pe.patch_address( + instruction_offset, [0xFF, 0x15] + new_value, lief.Binary.VA_TYPES.RVA + ) -def patch_calls_to_new_IAT(pe: lief.PE.Binary, imp: lief.PE.Import, entry:lief.PE.ImportEntry, rva: int): + +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) - - + 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: +with open(dump_path, "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: +with open(iat_json_path, "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) +calls: list[dict[str, str]] = iat_data["calls"] +wave_entry = int(iat_data["entry"], 16) # Define all sections as writeable, to help with some weird stuff we're seeing for section in pe.sections: @@ -58,7 +75,7 @@ for section in pe.sections: section.characteristics_lists.append(lief.PE.Section.CHARACTERISTICS.MEM_EXECUTE) # patch entrypoint -entrypoint_format = int(hex(wave_entry)[-4:],16) +entrypoint_format = int(hex(wave_entry)[-4:], 16) pe.optional_header.addressof_entrypoint = entrypoint_format # remove all current imports @@ -68,14 +85,16 @@ pe.optional_header.addressof_entrypoint = entrypoint_format 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): + 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 +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) -print("Wrote the patched executable as patched.exe") \ No newline at end of file +print("Wrote the patched executable as patched.exe")