diff --git a/generate_deep_traces.py b/generate_deep_traces.py new file mode 100644 index 0000000..1b003e9 --- /dev/null +++ b/generate_deep_traces.py @@ -0,0 +1,229 @@ +import click +import json +import pathlib + +from minor_patch_diff import get_opcode_offset, get_longest_switch, ZONE_PROTO_DOWN_SIG + + +class RefNode: + def __init__(self, ea): + self.ea = ea + self.calls = [] + self.branch0 = None + self.branch1 = None + + def add_call(self, ea): + node = RefNode(ea) + self.calls.append(node) + return node + + def add_b0(self, ea): + if self.branch0: + raise Exception(f"Node {repr(self)} already has a branch0!") + self.branch0 = RefNode(ea) + return self.branch0 + + def add_b1(self, ea): + if self.branch1: + raise Exception(f"Node {repr(self)} already has a branch1!") + self.branch1 = RefNode(ea) + return self.branch1 + + def __repr__(self): + calls = f", calls: {self.calls}" if self.calls else "" + branch0 = f", b0: {self.branch0}" if self.branch0 else "" + branch1 = f", b1: {self.branch1}" if self.branch1 else "" + return f"{{ ea: {hex(self.ea)}{calls}{branch0}{branch1} }}" + + +class OpcodeCase: + def __init__(self, ea, opcodes): + self.ref = RefNode(ea) + self.opcodes = opcodes + + def __repr__(self): + return f"{[hex(opcode) for opcode in self.opcodes]}: { repr(self.ref) }" + + +class BlockDict(dict): + def missing_set(self, eas): + missing = set() + for ea in eas: + if ea not in self: + missing.add(ea) + return missing + + def update_with_missing(self, r2, eas): + missing = self.missing_set(eas) + if len(missing) == 0: + return + + block_json_list = r2.cmd( + f"pdbj @@={' '.join((hex(ea) for ea in missing))}" + ).splitlines() + missing_blocks = [json.loads(block_json) for block_json in block_json_list] + additional_blocks = { + missing_block[0]["offset"]: missing_block + for missing_block in missing_blocks + } + self.update(additional_blocks) + + if len(additional_blocks) < len(missing): + return self.missing_set(eas) + + def mark_missing_as_unknown(self, eas): + missing = self.missing_set(eas) + self.update({ea: [{"unknown": True}] for ea in missing}) + + +def populate_child_refs(blocks, refs): + """Runs a single step of a BFS for traversing calls/jumps""" + child_refs = [] + for ref in refs: + for insn in blocks[ref.ea]: + if "jump" in insn: + jump_ea = insn["jump"] + if insn["type"] == "call": + child_refs.append(ref.add_call(jump_ea)) + else: + child_refs.append(ref.add_b1(jump_ea)) + if "fail" in insn: + child_refs.append(ref.add_b0(insn["fail"])) + return child_refs + + +def generate_opcodes_db(r2, switch, opcode_offset, fn_graph): + opcodes_db = dict() + + for case_ea, data in switch.items(): + opcodes = data["opcodes"] + resolved_opcodes = [int(opcode) + opcode_offset for opcode in opcodes] + opcodes_db[resolved_opcodes[0]] = OpcodeCase(case_ea, resolved_opcodes) + + blocks = BlockDict() + for bb in fn_graph["bbs"]: + blocks[bb["addr"]] = bb["ops"] + + # This should be a no-op since it is assumed the fn_graph would have + # every block in the function. + still_missing = blocks.update_with_missing(r2, switch.keys()) + if still_missing: + raise Exception("There's no way there should be any missing blocks here") + + ref_nodes = [] + for opcase in opcodes_db.values(): + ref_nodes.append(opcase.ref) + + for i in range(10): + child_refs = populate_child_refs(blocks, ref_nodes) + eas = [ref.ea for ref in child_refs] + still_missing = blocks.update_with_missing(r2, eas) + if still_missing: + r2.cmd(f"af @@={' '.join((hex(ea) for ea in still_missing))}") + yet_still_missing = blocks.update_with_missing(r2, still_missing) + if yet_still_missing: + blocks.mark_missing_as_unknown(yet_still_missing) + ref_nodes = child_refs + + return opcodes_db, blocks + + +def extract_opcode_data(exe_file): + from utils import eprint, create_r2_byte_pattern, sync_r2_output + + import r2pipe + + r2 = r2pipe.open(exe_file, ["-2"]) + eprint(f"Radare loaded {exe_file}") + + sync_r2_output(r2) + + p = create_r2_byte_pattern(ZONE_PROTO_DOWN_SIG) + target = r2.cmd(f"/x {p}").split()[0] # Find byte pattern + packet_handler_ea = int(target, 16) + + r2.cmd(f"s {target}") # Seek to target + + ## STEP 1: Grab switch cases + r2.cmd("f--") # Delete existing flags + r2.cmd("afr") # Analyze function recursively + switch_cases = r2.cmdj(f"fj") + + eprint(f" Loaded switch cases") + + ## STEP 2: Grab opcode offset + opcode_offset = get_opcode_offset(r2) + eprint(f" Found opcode offset: {opcode_offset}") + + ## STEP 3: Grab function graph + fn_graph = r2.cmdj(f"pdrj") + + ## STEP 4: Process data + packet_handler_switch = get_longest_switch(switch_cases) + + opcodes_db, blocks = generate_opcodes_db( + r2, packet_handler_switch, opcode_offset, fn_graph + ) + + eprint(f" Loaded {len(opcodes_db)} cases from packet handler") + + r2.quit() + + return opcodes_db, blocks + + +def bb_lines(blocks, ref): + lines = [] + block = blocks[ref.ea] + for insn in block: + if "unknown" in insn: + lines.append("UNKNOWN_BLOCK") + else: + lines.append(insn["opcode"]) + return lines + + +def trace_lines(blocks, ref): + # Generate traces in a BFS fashion + lines = [] + refs = [(None, ref)] + while refs: + (name, ref) = refs.pop(0) + if name: + lines.append(name) + lines.extend(bb_lines(blocks, ref)) + if name: + lines.append(f"{name}_END") + for i, call_ref in enumerate(ref.calls): + refs.append((f"CALL{i}", call_ref)) + if ref.branch0: + refs.append(("BRANCH0", ref.branch0)) + if ref.branch1: + refs.append(("BRANCH1", ref.branch1)) + return lines + + +@click.command() +@click.argument( + "exe_file", type=click.Path(exists=True, dir_okay=False, resolve_path=True) +) +@click.argument("output_dir", type=click.Path(file_okay=False)) +def generate_deep_traces(exe_file, output_dir): + opcodes_db, blocks = extract_opcode_data(exe_file) + + pathlib.Path(output_dir).mkdir(parents=True, exist_ok=True) + + for opcode, opcase in opcodes_db.items(): + with open(f"{output_dir}/{hex(opcode)}.asm", "w+") as f: + trace = trace_lines(blocks, opcase.ref) + f.writelines(s + "\n" for s in trace) + + with open(f"{output_dir}/opcode_sets.json", "w+") as f: + opcode_sets = { + opcode: opcase.opcodes for (opcode, opcase) in opcodes_db.items() + } + json.dump(opcode_sets, f) + + +if __name__ == "__main__": + generate_deep_traces()