Better comments

iat.py

@@ -84,6 +84,7 @@ def main():
         cfg = json.load(f)
         utils.print_debug(f"Opened file {args.trace} as the TraceCFG JSON")
 
+    # determine target wave
     if args.wave == None and args.dump[-5:] == ".dump":
         wave = int(args.dump[-9:-5])
     else:
@@ -141,9 +142,8 @@ def main():
             func_calls_list.append(func["name"])
             imported_dll.add_entry(func["name"])
 
-    # 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.
+    # Define all sections as writeable, to prevent permission issues.
-
+    # Ideally, we would like to have the actual permitions from Goatracer at some point in the future
-    # Define all sections as writeable, to help with some weird stuff we're seeing
     for section in pe.sections:
         section.characteristics = (
             lief.PE.Section.CHARACTERISTICS.MEM_WRITE.value
@@ -152,10 +152,7 @@ def main():
             + lief.PE.Section.CHARACTERISTICS.CNT_INITIALIZED_DATA.value
         )
 
-    # write result
+    # 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.
-    config = lief.PE.Builder.config_t()
-    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):
@@ -164,9 +161,14 @@ def main():
         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)
+    # write result
-    print("Wrote the patched executable as 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 = patching_callback  # Define the callback
+    output_path = args.output
+    pe.write(output_path, config)
+    print(f"Wrote the patched executable as {output_path}")
 
 if __name__ == "__main__":
     main()