From 803f0e1cd11e9042fb10e2e72b6d161921733c88 Mon Sep 17 00:00:00 2001
From: Seliaste <aena.aria@posteo.net>
Date: Mon, 23 Mar 2026 15:20:29 +0100
Subject: [PATCH] it works!!!

---
 .gitignore |  1 +
 iat.py     | 66 ++++++++++++++++++++++++++++++++++++++++++++----------
 2 files changed, 55 insertions(+), 12 deletions(-)

diff --git a/.gitignore b/.gitignore
index ab3e8ce..9f0ebbc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -162,3 +162,4 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
+patched.exe
\ No newline at end of file
diff --git a/iat.py b/iat.py
index 97d38aa..de99567 100644
--- a/iat.py
+++ b/iat.py
@@ -1,30 +1,72 @@
 import json
-
 import lief
+lief.disable_leak_warning() # warnings to disable for the callback
 
-# wave to parse
+# 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)
 
-# 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)
+# 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
-pe.write("patched.exe")
+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)