cfg_parser.py

@@ -1,3 +1,5 @@
+import utils
+
 def parse_wave_nodes(cfg,wave: int) -> list:
     return list(filter(lambda node: node["wave"] == wave,cfg["nodes"]))
 

iat.py

@@ -4,12 +4,16 @@ import lief
 import patch
 import cfg_parser
 import reginit
+import utils
 
 lief.disable_leak_warning()  # warnings to disable for the callback
 
 with open("lib/WindowsDllsExport/win10-19043-exports.json", "rb") as f:
     api_info = json.load(f)
 
+
+
+
 # Retrives all unique DLL names being imported
 def get_used_dlls(calls: list[dict[str, str]]) -> set[str]:
     res = set()
@@ -54,7 +58,6 @@ def link_func_to_dll(func_list):
             res.append(res_new)
     return res
 
-
 def main():
     parser = argparse.ArgumentParser(prog="iat.py", description="Create a patched PE from a binary dump and a traceCFG file.")
 
@@ -65,22 +68,27 @@ def main():
     # Additional arguments
     parser.add_argument("-o", "--output", type=str, help="Specify an output filepath for the patched PE.")
     parser.add_argument("-w", "--wave", type=int, help="Specify the wave number for the binary dump (if it can't be inferred from the filename)")
+    parser.add_argument("-v", '--verbose', action='store_true', help="Output additional debug info")
 
     args = parser.parse_args()
+    utils.set_verbose(args.verbose)
 
     # open wave dump file
     with open(args.dump, "rb") as f:
         pe = lief.parse(f)
         assert isinstance(pe, lief.PE.Binary)
+        utils.print_debug(f"Opened file {args.dump} as the binary dump")
 
     # open traceCFG json
     with open(args.trace, "r") as f:
         cfg = json.load(f)
+        utils.print_debug(f"Opened file {args.trace} as the TraceCFG JSON")
 
     if args.wave == None and args.dump[-5:] == ".dump":
         wave = int(args.dump[-9:-5])
     else:
         wave = args.wave
+    utils.print_debug(f"Determined wave to be {wave}")
 
     calls = cfg_parser.parse_syscalls(cfg,wave)
     wave_entry = cfg_parser.parse_wave_entrypoint(cfg,wave)
@@ -149,11 +157,13 @@ def main():
     config.imports = True  # allows the config of the writer to write a new IAT
 
     def patching_callback(pe: lief.PE.Binary, imp: lief.PE.Import, entry: lief.PE.ImportEntry, rva: int):
+        utils.print_debug(f"Now trying to patch {entry.name}!{imp.name}...")
         for call in filter(lambda x: x["name"] == f"{imp.name.upper()}!{entry.name}", calls):
             patch.patch_instr_to_new_IAT_entry(pe, call, rva)
         # patch additional non-call related info
         for func in filter(lambda x: x["name"] == entry.name and x["dll"] == imp.name, func_dll_list):
             patch.patch_addr_found_in_mem(pe, rva, func["addr"])
+        utils.print_debug(f"Done!\n")
     config.resolved_iat_cbk = patching_callback  # callback after the IAT has been written
     pe.write("patched.exe" if args.output == None else args.output, config)
     print("Wrote the patched executable as patched.exe")

utils.py

@@ -8,3 +8,11 @@ def hex_address_to_memory_representation(hex_addr: str, is_32b: bool, is_little_
     if not is_little_endian:
         mem_value = mem_value[::-1]  # reverse byte order for big endian
     return mem_value
+
+verbose = False
+def print_debug(msg:str):
+    if(verbose): print(msg)
+
+def set_verbose(value:bool):
+    global verbose
+    verbose = value