Add function to generate traces

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()