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

code formatting and newly created enums to clean up code

Browse source
This commit is contained in:
Aéna Aria 2026-04-02 14:02:07 +02:00
parent f46cc2438f
commit 6f4fccd350
5 changed files with 111 additions and 84 deletions
Download patch file Download diff file Expand all files Collapse all files

28
iat.py
View file

@ -1,8 +1,10 @@
import argparse
import json
import lief
import patch
import cfg_parser
import patch
import reginit
import utils
@ -12,8 +14,6 @@ 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()
@ -58,8 +58,13 @@ 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.", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser = argparse.ArgumentParser(
prog="iat.py",
description="Create a patched PE from a binary dump and a traceCFG file.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
# Input arguments
parser.add_argument("dump", type=str, help="The path to the wave dump file (usually ends with .dump)")
@ -68,7 +73,7 @@ def main():
# Additional arguments
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("-v", "--verbose", action="store_true", help="Output additional debug info")
args = parser.parse_args()
utils.set_verbose(args.verbose)
@ -85,24 +90,24 @@ def main():
utils.print_debug(f"Opened file {args.trace} as the TraceCFG JSON")
# determine target wave
if args.wave == None and args.dump[-5:] == ".dump":
if args.wave is 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)
calls = cfg_parser.parse_syscalls(cfg, wave)
wave_entry = cfg_parser.parse_wave_entrypoint(cfg, wave)
# create new section
iatpatch_section = lief.PE.Section(".iatpatch")
iatpatch_content = []
# registers initiation
iatpatch_content += reginit.generate_reg_init_code(cfg,pe,wave,wave_entry)
iatpatch_content += reginit.generate_reg_init_code(cfg, pe, wave, wave_entry)
# write patch section code
iatpatch_section.content = iatpatch_content # pyright: ignore[reportAttributeAccessIssue]
iatpatch_section.content = iatpatch_content # pyright: ignore[reportAttributeAccessIssue]
# add new section to PE
pe.add_section(iatpatch_section)
@ -160,7 +165,7 @@ def main():
# 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")
utils.print_debug("Done!\n")
# write result
config = lief.PE.Builder.config_t()
@ -170,5 +175,6 @@ def main():
pe.write(output_path, config)
print(f"Wrote the patched executable as {output_path}")
if __name__ == "__main__":
main()