Small factoring and formatting issues

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")
+print("Wrote the patched executable as patched.exe")