import click import json import pathlib from minor_patch_diff import ( get_opcode_offset, get_longest_switch, get_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(get_zone_proto_down_sig(exe_file)) 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): """ Generates deep traces for every packet handler in the target EXE_FILE. This outputs an ASM trace as an .asm file for each pointer opcode. It also outputs an `opcode_sets.json` that maps each pointer opcode to the full set of opcodes handled by that trace. These aren't normal traces by any measure; they are simply basic blocks printed in BFS order, where children blocks are added to the BFS tree by traversing calls and jumps. Example: python generate_deep_traces.py ffxiv_dx11.6.28h.exe 6.28h-traces """ 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()