diff --git a/cfg_parser.py b/cfg_parser.py index 113fb8e..2ea23a9 100644 --- a/cfg_parser.py +++ b/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"])) diff --git a/iat.py b/iat.py index 1ee858d..72dc957 100644 --- a/iat.py +++ b/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") diff --git a/utils.py b/utils.py index 9fe3f2b..3c0e069 100644 --- a/utils.py +++ b/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