lief-iat-reconstruction/cfg_parser.py

36 lines
1.5 KiB
Python

def parse_wave_nodes(cfg,wave: int) -> list:
return list(filter(lambda node: node["wave"] == wave,cfg["nodes"]))
def parse_procaddr_calls(cfg,wave:int):
res = []
wave_nodes:list[dict] = parse_wave_nodes(cfg,wave)
for node in wave_nodes:
if "syscalls" in node.keys():
for syscall in node["syscalls"]:
if syscall["name"] == "KERNEL32.DLL!GetProcAddress":
funcname = syscall["arguments"][-1].split("\"")[1]
func_addr = syscall["return"]
res.append({"name": funcname, "addr": func_addr})
return res
def parse_syscalls(cfg,wave: int) -> list[dict[str, str]]:
res: list[dict[str,str]] = []
wave_nodes:list[dict] = parse_wave_nodes(cfg,wave)
no_repeat = []
for node in wave_nodes:
if "syscalls" in node.keys():
for syscall in node["syscalls"]:
if node["last_instr"] in no_repeat:
continue
adress = node["last_instr"] # call is at the end of the basic block
name = syscall["name"]
current_instruction = node["instructions"][-1]["mnemonic"]
no_repeat.append(adress)
res.append({"adress":adress,"name":name})
return res
def parse_wave_entrypoint(cfg,wave: int) -> int:
return int(parse_wave_nodes(cfg,wave)[0]["start"],16)
def parse_bb_registers(cfg,wave:int,n_bb:int) -> dict[str,str]:
return parse_wave_nodes(cfg,wave)[n_bb]["registers"]