diff options
Diffstat (limited to 'python/examples')
| -rw-r--r-- | python/examples/wf_unflatten.py | 293 |
1 files changed, 293 insertions, 0 deletions
diff --git a/python/examples/wf_unflatten.py b/python/examples/wf_unflatten.py new file mode 100644 index 00000000..ccda2be2 --- /dev/null +++ b/python/examples/wf_unflatten.py @@ -0,0 +1,293 @@ +import json +from binaryninja import Workflow, Activity, AnalysisContext, MediumLevelILBasicBlock, \ + MediumLevelILInstruction, ReportCollection, FlowGraphReport, show_report_collection, \ + FlowGraph, DisassemblySettings +from binaryninja.mediumlevelil import * +from binaryninja.enums import * + + +""" +This workflow reverses the control-flow flattening algorithm of Limoncello[1], +at least some of the time. It does this in a relatively simple manner: + +1. Find "dispatcher blocks" which contain a MLIL_JUMP_TO whose targets all have the + dispatcher as a post-dominator (ie all targets must flow back to the dispatcher) +2. Walk backwards from all unconditional branches into the dispatcher until finding blocks + with conditional branches and build a list of all unconditional continuation blocks + for each block that flows into the dispatcher. +3. Copy those continuation blocks and the dispatcher into each block that calls them. + This leaves you with a copy of the dispatcher block (and intermediate bookkeeping) in + every path that would have normally flowed into the dispatcher. +4. Since each path now has its own copy of the dispatcher, use MLIL dataflow to solve for + which branch of the dispatcher is taken at each copy. When a dispatcher with a solved + target is encountered, rewrite it as an MLIL_GOTO with the target directly. + +[1] https://github.com/jonpalmisc/limoncello +""" + + +def is_dispatcher(block: MediumLevelILBasicBlock) -> bool: + """ + Determine if a block looks like a CFF dispatcher, i.e. all outgoing edges + post-dominate it (or return) + + :param block: The block to check + :return: True if it's a dispatcher + """ + for ins in block: + ins: MediumLevelILInstruction + if ins.operation == MediumLevelILOperation.MLIL_JUMP_TO: + for edge in block.outgoing_edges: + out_block = edge.target + # This is too trivial for fancier cases with multiple blocks that all lead + # to a return, but whatever this is just a sample plugin + if len(out_block.outgoing_edges) == 0: + continue + if block not in out_block.post_dominators: + return False + + return True + return False + + +def graph_pls(fn: MediumLevelILFunction) -> FlowGraph: + # For ReportCollection, create graph with the settings I want + settings = DisassemblySettings() + settings.set_option(DisassemblyOption.ShowAddress, True) + return fn.create_graph_immediate(settings=settings) + + +def rewrite_action(context: AnalysisContext, do_it: bool): + # Main workflow action + + # ===================================================================== + # Custom debug report + + report = None + if context.function.check_for_debug_report("unflatten"): + report = ReportCollection() + + try: + if report is not None: + init_graph = graph_pls(context.mlil) + report.append(FlowGraphReport("Initial", init_graph, context.view)) + + # ===================================================================== + # Finding flattened control flow and resolving continuations + + # Look for dispatcher block + dispatcher = None + for block in context.mlil.basic_blocks: + if is_dispatcher(block): + dispatcher = block + break + + if dispatcher is None: + return + + if report is not None: + init_graph = graph_pls(context.mlil) + for node_index, node in enumerate(init_graph.nodes): + if node.basic_block.start == dispatcher.start: + node.highlight = HighlightStandardColor.RedHighlightColor + init_graph.replace(node_index, node) + report.append(FlowGraphReport("Finding Dispatcher", init_graph, context.view)) + + # Find blocks that flow directly into the dispatcher + queue = [] + to_copy = {} + for incoming in dispatcher.incoming_edges: + if dispatcher in incoming.source.dominators: + to_copy[incoming.source] = [dispatcher] + queue.append(incoming.source) + + # For each of these, walk back along unconditional branch edges to find their + # list of continuation blocks which will be copied at their end + while len(queue) > 0: + # This is a sort of backwards BFS + top = queue.pop(0) + any_conditional = False + for incoming in top.incoming_edges: + if incoming.type != BranchType.UnconditionalBranch: + any_conditional = True + # Having any conditional branches ends the unconditional chain + if not any_conditional: + for incoming in top.incoming_edges: + if incoming.type == BranchType.UnconditionalBranch: + if incoming.source not in to_copy: + queue.append(incoming.source) + to_copy[incoming.source] = [top] + to_copy[top] + to_copy[top] = [] + + # Ignore empty continuation paths + for bb in list(to_copy.keys()): + if len(to_copy[bb]) == 0: + to_copy.pop(bb) + + if report is not None: + for block, path in to_copy.items(): + init_graph = graph_pls(context.mlil) + for node_index, node in enumerate(init_graph.nodes): + if node.basic_block.start in [bb.start for bb in path]: + node.highlight = HighlightStandardColor.RedHighlightColor + if node.basic_block.start == block.start: + node.highlight = HighlightStandardColor.GreenHighlightColor + init_graph.replace(node_index, node) + report.append(FlowGraphReport(" Blocks flowing into", init_graph, context.view)) + + # ===================================================================== + # Modify the IL to copy the continuations into all the blocks calling the dispatcher + + # Make a new IL function and append the modified instructions to it + old_mlil = context.mlil + new_mlil = MediumLevelILFunction(old_mlil.arch, low_level_il=context.llil) + new_mlil.prepare_to_copy_function(old_mlil) + block_map_starts = {} + + # Copy all instructions in all blocks of the old version of the function + for block in old_mlil.basic_blocks: + new_mlil.prepare_to_copy_block(block) + block_map_starts[block] = len(new_mlil) + + new_mlil.set_current_address(old_mlil[InstructionIndex(block.start)].address, block.arch) + + for instr_index in range(block.start, block.end): + instr: MediumLevelILInstruction = old_mlil[InstructionIndex(instr_index)] + + # Copy continuation blocks to end of block calling dispatcher + if block in to_copy: + path = to_copy[block] + + if instr_index == block.end - 1: + # For every block in the continuation, copy it at the end of this block + for copy_block in path: + new_mlil.prepare_to_copy_block(copy_block) + + # Skip the final instruction in the continuations because it is a MLIL_GOTO + end = copy_block.end - 1 + if copy_block == dispatcher: + end = copy_block.end + + for copy_block_instr_index in range(copy_block.start, end): + # Copy instructions as-is + copy_block_instr: MediumLevelILInstruction = old_mlil[InstructionIndex(copy_block_instr_index)] + new_mlil.set_current_address(copy_block_instr.address, copy_block.arch) + new_mlil.append(copy_block_instr.copy_to(new_mlil)) + continue + + # Otherwise, copy the instruction as-is + new_mlil.set_current_address(instr.address, block.arch) + new_mlil.append(instr.copy_to(new_mlil)) + + # Generate blocks and SSA (for dataflow) for the next part + new_mlil.finalize() + new_mlil.generate_ssa_form() + + # Since we're constructing a new function twice, we need to commit the mappings + # of the intermediate function before copying again so that mappings will resolve + # all the way to the end (gross) + # TODO: Construct from another function without needing this + context.mlil = new_mlil + + if report is not None: + newer_graph = graph_pls(new_mlil) + report.append(FlowGraphReport("Swapped dispatch with jump_to", newer_graph, context.view)) + + # ===================================================================== + # Now convert all MLIL_JUMP_TO with known dest to a jump + + # Maybe this should be a separate workflow action (so it can be composed) + old_mlil = new_mlil + new_mlil = MediumLevelILFunction(old_mlil.arch, low_level_il=context.llil) + new_mlil.prepare_to_copy_function(old_mlil) + + # Keep a running list of blocks and labels so we can line up the MLIL_GOTOs + block_labels = {} + + for old_block in old_mlil.basic_blocks: + new_mlil.prepare_to_copy_block(old_block) + new_mlil.set_current_address(old_mlil[InstructionIndex(old_block.start)].address, old_block.arch) + + # Update block label list for the MLIL_GOTOs + if old_block.start in block_labels: + label = block_labels[old_block.start] + else: + label = MediumLevelILLabel() + block_labels[old_block.start] = label + new_mlil.mark_label(label) + + for instr_index in range(old_block.start, old_block.end): + instr: MediumLevelILInstruction = old_mlil[InstructionIndex(instr_index)] + # If we find a MLIL_JUMP_TO with a known constant dest, then rewrite it + # to a MLIL_GOTO with the known dest filled in + if instr.operation == MediumLevelILOperation.MLIL_JUMP_TO: + if instr.dest.value.type == RegisterValueType.ConstantPointerValue: + dest_value = instr.dest.value.value + if dest_value in instr.targets: + old_target_index = instr.targets[dest_value] + if old_target_index in block_labels: + target_label = block_labels[old_target_index] + else: + target_label = MediumLevelILLabel() + block_labels[old_target_index] = target_label + new_mlil.append(new_mlil.goto(target_label, ILSourceLocation.from_instruction(instr))) + continue + + # Otherwise, copy the instruction as-is + new_mlil.set_current_address(instr.address, old_block.arch) + new_mlil.append(instr.copy_to(new_mlil)) + + new_mlil.finalize() + new_mlil.generate_ssa_form() + + if report is not None: + newer_graph = graph_pls(new_mlil) + report.append(FlowGraphReport("Resolved constant jump_to's", newer_graph, context.view)) + + # ===================================================================== + # And we're done + + if do_it: + context.mlil = new_mlil + finally: + # Show debug report if requested, even on exception thrown + if report is not None: + show_report_collection("Unflatten Debug Report", report) + + +# Create and register the workflow for translating these instructions +wf = Workflow("core.function.metaAnalysis").clone("core.function.metaAnalysis") + +# Define the custom activity configuration +wf.register_activity(Activity( + configuration=json.dumps({ + "name": "extension.unflatten_limoncello.unflatten.dry_run", + "title": "Unflatten (Limoncello) Dry Run", + "description": "Detect and reverse Limoncello's Control Flow Flattening scheme.", + "eligibility": { + "auto": { + "default": False + } + } + }), + action=lambda context: rewrite_action(context, False) +)) +wf.register_activity(Activity( + configuration=json.dumps({ + "name": "extension.unflatten_limoncello.unflatten", + "title": "Unflatten (Limoncello)", + "description": "Detect and reverse Limoncello's Control Flow Flattening scheme.", + "eligibility": { + "auto": { + "default": False + } + } + }), + action=lambda context: rewrite_action(context, True) +)) + +wf.insert_after("core.function.generateMediumLevelIL",[ + "extension.unflatten_limoncello.unflatten.dry_run", + "extension.unflatten_limoncello.unflatten" +]) +wf.register() |