seliaste/lief-iat-reconstruction
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Added a new argument to disable the registers initialization if wanted

Browse source
This commit is contained in:
Aéna Aria 2026-04-02 14:21:12 +02:00
parent fc0b7ac3d3
commit bc2753f98a
1 changed files with 13 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

20
iat.py
View file

@ -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,17 +104,22 @@ def main():
iatpatch_section = lief.PE.Section(".iatpatch")
iatpatch_content = []
# registers initiation
iatpatch_content += reginit.generate_reg_init_code(cfg, pe, wave, wave_entry)
# registers initialization
if not args.disable_reginit:
iatpatch_content += reginit.generate_reg_init_code(cfg, pe, wave, wave_entry)
# write patch section code
iatpatch_section.content = iatpatch_content # pyright: ignore[reportAttributeAccessIssue]
if iatpatch_content != []:
iatpatch_section.content = iatpatch_content # pyright: ignore[reportAttributeAccessIssue]
# add new section to PE
pe.add_section(iatpatch_section)
# add new section to PE
pe.add_section(iatpatch_section)
# patch entrypoint
entrypoint_format = int(hex(pe.get_section(".iatpatch").virtual_address)[-4:], 16)
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
# remove all current imports