Added a new argument to disable the registers initialization if wanted

iat.py

@@ -71,9 +71,10 @@ def main():
     parser.add_argument("trace", type=str, help="The path to the traceCFG file (.json)")
 
     # Additional arguments
+    parser.add_argument("-v", "--verbose", action="store_true", help="Output additional debug info")
     parser.add_argument("-o", "--output", type=str, default="patched.exe", 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")
+    parser.add_argument("--disable-reginit", action="store_true", help="Disable initialization of the registry before jumping to the wave start")
 
     args = parser.parse_args()
     utils.set_verbose(args.verbose)
@@ -103,16 +104,21 @@ def main():
     iatpatch_section = lief.PE.Section(".iatpatch")
     iatpatch_content = []
 
-    # registers initiation
+    # registers initialization
+    if not args.disable_reginit:
         iatpatch_content += reginit.generate_reg_init_code(cfg, pe, wave, wave_entry)
 
     # write patch section code
+    if iatpatch_content != []:
         iatpatch_section.content = iatpatch_content  # pyright: ignore[reportAttributeAccessIssue]
 
         # add new section to PE
         pe.add_section(iatpatch_section)
 
     # patch entrypoint
+    if args.disable_reginit:
+        entrypoint_format = int(hex(cfg_parser.parse_wave_entrypoint(cfg, wave))[-4:], 16)
+    else:
         entrypoint_format = int(hex(pe.get_section(".iatpatch").virtual_address)[-4:], 16)
     pe.optional_header.addressof_entrypoint = entrypoint_format