From 7871953325024282191116771022201477045749 Mon Sep 17 00:00:00 2001 From: Glenn Smith Date: Mon, 30 Jun 2025 13:46:52 -0400 Subject: Python: Collect and pass mappings when copying MLIL functions --- python/commonil.py | 49 +++++++++-- python/examples/wf_test_copy_expr.py | 59 +++++++++++-- python/mediumlevelil.py | 160 +++++++++++++++++++++++++++++++++-- python/workflow.py | 48 ++++++++++- 4 files changed, 295 insertions(+), 21 deletions(-) (limited to 'python') diff --git a/python/commonil.py b/python/commonil.py index e7822518..ae520627 100644 --- a/python/commonil.py +++ b/python/commonil.py @@ -19,7 +19,7 @@ # IN THE SOFTWARE. from dataclasses import dataclass -from typing import Union +from typing import Union, Optional from .flowgraph import FlowGraph, FlowGraphNode from .enums import BranchType from .interaction import show_graph_report @@ -29,6 +29,9 @@ from . import mediumlevelil from . import highlevelil +invalid_il_index = 0xffffffffffffffff + + # This file contains a list of top level abstract classes for implementing BNIL instructions @dataclass(frozen=True, repr=False, eq=False) class BaseILInstruction: @@ -210,7 +213,6 @@ class AliasedVariableInstruction(VariableInstruction): pass -@dataclass class ILSourceLocation: """ ILSourceLocation is used to indicate where expressions were defined during the lifting process @@ -220,15 +222,52 @@ class ILSourceLocation: address: int source_operand: int + source_llil_instruction: Optional['lowlevelil.LowLevelILInstruction'] = None + source_mlil_instruction: Optional['mediumlevelil.MediumLevelILInstruction'] = None + source_hlil_instruction: Optional['highlevelil.HighLevelILInstruction'] = None + il_direct: bool = True + + def __init__(self, address: int, source_operand: int): + self.address = address + self.source_operand = source_operand + + def __repr__(self): + instr = "" + if self.source_llil_instruction is not None: + instr = f" (from LLIL {self.source_llil_instruction})" + if self.source_mlil_instruction is not None: + instr = f" (from MLIL {self.source_mlil_instruction})" + if self.source_hlil_instruction is not None: + instr = f" (from HLIL {self.source_hlil_instruction})" + return f"" + + def __hash__(self): + return hash((self.address, self.source_operand)) + + def __eq__(self, other): + if not isinstance(other, ILSourceLocation): + return False + return self.address == other.address and self.source_operand == other.source_operand + @classmethod def from_instruction( cls, - instr: Union['lowlevelil.LowLevelILInstruction', 'mediumlevelil.MediumLevelILInstruction', 'highlevelil.HighLevelILInstruction'] + instr: Union['lowlevelil.LowLevelILInstruction', 'mediumlevelil.MediumLevelILInstruction', 'highlevelil.HighLevelILInstruction'], + il_direct: bool = True ) -> 'ILSourceLocation': """ Get the source location of a given instruction :param instr: Instruction, Low, Medium, or High level :return: Its location """ - return cls(instr.address, instr.source_operand) - + loc = cls(instr.address, instr.source_operand) + if isinstance(instr, lowlevelil.LowLevelILInstruction): + loc.source_llil_instruction = instr + elif isinstance(instr, mediumlevelil.MediumLevelILInstruction): + loc.source_mlil_instruction = instr + elif isinstance(instr, highlevelil.HighLevelILInstruction): + loc.source_hlil_instruction = instr + else: + log_warn(f"Unknown instruction type {type(instr)}") + loc.il_direct = il_direct + return loc diff --git a/python/examples/wf_test_copy_expr.py b/python/examples/wf_test_copy_expr.py index 7a3da06c..77e6e159 100644 --- a/python/examples/wf_test_copy_expr.py +++ b/python/examples/wf_test_copy_expr.py @@ -1,3 +1,4 @@ +import functools import json import math @@ -67,7 +68,14 @@ def assert_llil_eq(old_insn: LowLevelILInstruction, new_insn: LowLevelILInstruct assert old_op[1] == new_op[1], err_msg -def assert_mlil_eq(old_insn: LowLevelILInstruction, new_insn: LowLevelILInstruction): +@functools.lru_cache(maxsize=8) +def get_mlil_maps(mlil: MediumLevelILFunction, builders: bool) -> Tuple[LLILSSAToMLILInstructionMapping, LLILSSAToMLILExpressionMapping]: + instr_map = mlil._get_llil_ssa_to_mlil_instr_map(builders) + expr_map = mlil._get_llil_ssa_to_mlil_expr_map(builders) + return instr_map, expr_map + + +def assert_mlil_eq(old_insn: MediumLevelILInstruction, new_insn: MediumLevelILInstruction): """ Make sure that these two instructions are the same (probably correct). Asserts otherwise. @@ -78,10 +86,23 @@ def assert_mlil_eq(old_insn: LowLevelILInstruction, new_insn: LowLevelILInstruct """ err_msg = (hex(old_insn.address), old_insn, new_insn) assert old_insn.operation == new_insn.operation, err_msg - # assert old_insn.attributes == new_insn.attributes, err_msg + assert old_insn.attributes == new_insn.attributes, err_msg assert old_insn.size == new_insn.size, err_msg assert old_insn.source_location == new_insn.source_location, err_msg assert len(old_insn.operands) == len(new_insn.operands), err_msg + # Type only applies once we've generated SSA form (probably not consistent) + # assert old_insn.expr_type == new_insn.expr_type, f"{err_msg} {old_insn.expr_type} {new_insn.expr_type}" + + instr_map, expr_map = get_mlil_maps(new_insn.function, True) + + # Compare that the instruction's LLIL SSA map is the same as the old function + if old_insn.instr_index is not None and old_insn.function.get_expr_index_for_instruction(old_insn.instr_index) == old_insn.expr_index: + old_llil_ssa = old_insn.function.get_low_level_il_instruction_index(old_insn.instr_index) + if old_llil_ssa is not None: + assert [mlil for (llil, mlil) in instr_map.items() if llil == old_llil_ssa] == [new_insn.instr_index], err_msg + else: + assert [mlil for (llil, mlil) in instr_map.items() if llil == old_llil_ssa] == [], err_msg + # Can't compare operands directly since IL expression indices might change when # copying an instruction to another function for i, (old_op, new_op) in enumerate(zip(old_insn.detailed_operands, new_insn.detailed_operands)): @@ -231,6 +252,11 @@ def mlil_action(context: AnalysisContext): report.append(FlowGraphReport("new graph", new_mlil.create_graph_immediate(settings))) show_report_collection("copy expr test", report) + # Check expr mappings are the same + new_map = list(sorted(new_mlil._get_llil_ssa_to_mlil_expr_map(True), key=lambda o: (o.lower_index, o.higher_index))) + old_map = list(sorted(old_mlil._get_llil_ssa_to_mlil_expr_map(False), key=lambda o: (o.lower_index, o.higher_index))) + assert old_map == new_map + # Check all BBs have all the same instructions # Technically, this misses any instructions outside a BB, but those are not # picked up by analysis anyway, and therefore don't matter. @@ -240,15 +266,26 @@ def mlil_action(context: AnalysisContext): for old_insn, new_insn in zip(old_bb, new_bb): assert_mlil_eq(old_insn, new_insn) + # Make sure mappings update correctly following set + new_map = list(sorted(new_mlil._get_llil_ssa_to_mlil_expr_map(True), key=lambda o: (o.lower_index, o.higher_index))) + context.mlil = new_mlil + newer_map = list(sorted(context.mlil._get_llil_ssa_to_mlil_expr_map(False), key=lambda o: (o.lower_index, o.higher_index))) + assert new_map == newer_map + -wf = Workflow("core.function.metaAnalysis").clone("TestCopyExpr") +wf = Workflow("core.function.metaAnalysis").clone("core.function.metaAnalysis") # Define the custom activity configuration wf.register_activity(Activity( configuration=json.dumps({ "name": "extension.test_copy_expr.lil_action", "title": "Lifted IL copy_expr Test", - "description": "Makes sure copy_expr works on Lifted IL functions." + "description": "Makes sure copy_expr works on Lifted IL functions.", + "eligibility": { + "auto": { + "default": False + } + } }), action=lil_action )) @@ -256,7 +293,12 @@ wf.register_activity(Activity( configuration=json.dumps({ "name": "extension.test_copy_expr.llil_action", "title": "Low Level IL copy_expr Test", - "description": "Makes sure copy_expr works on Low Level IL functions." + "description": "Makes sure copy_expr works on Low Level IL functions.", + "eligibility": { + "auto": { + "default": False + } + } }), action=llil_action )) @@ -264,7 +306,12 @@ wf.register_activity(Activity( configuration=json.dumps({ "name": "extension.test_copy_expr.mlil_action", "title": "Medium Level IL copy_expr Test", - "description": "Makes sure copy_expr works on Medium Level IL functions." + "description": "Makes sure copy_expr works on Medium Level IL functions.", + "eligibility": { + "auto": { + "default": False + } + } }), action=mlil_action )) diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py index cc99db8f..5d386b3f 100644 --- a/python/mediumlevelil.py +++ b/python/mediumlevelil.py @@ -20,7 +20,7 @@ import ctypes import struct -from typing import (Optional, List, Union, Mapping, +from typing import (Optional, List, Union, Mapping, MutableMapping, Generator, NewType, Tuple, ClassVar, Dict, Set, Callable, Any, Iterator, overload) from dataclasses import dataclass from . import deprecation @@ -43,7 +43,7 @@ from .commonil import ( BaseILInstruction, Constant, BinaryOperation, UnaryOperation, Comparison, SSA, Phi, FloatingPoint, ControlFlow, Terminal, Call, Localcall, Syscall, Tailcall, Return, Signed, Arithmetic, Carry, DoublePrecision, Memory, Load, Store, RegisterStack, SetVar, Intrinsic, VariableInstruction, SSAVariableInstruction, AliasedVariableInstruction, - ILSourceLocation + ILSourceLocation, invalid_il_index ) TokenList = List['function.InstructionTextToken'] @@ -61,6 +61,30 @@ MediumLevelILOperandType = Union[int, float, 'MediumLevelILOperationAndSize', 'M MediumLevelILVisitorCallback = Callable[[str, MediumLevelILOperandType, str, Optional['MediumLevelILInstruction']], bool] StringOrType = Union[str, '_types.Type', '_types.TypeBuilder'] ILInstructionAttributeSet = Union[Set[ILInstructionAttribute], List[ILInstructionAttribute]] +LLILSSAToMLILInstructionMapping = MutableMapping['lowlevelil.InstructionIndex', InstructionIndex] + + +@dataclass(frozen=True) +class LLILSSAToMLILExpressionMap: + lower_index: 'lowlevelil.ExpressionIndex' + higher_index: ExpressionIndex + map_lower_to_higher: bool + map_higher_to_lower: bool + lower_to_higher_direct: bool + higher_to_lower_direct: bool + + def _to_core_struct(self): + result = core.BNExprMapInfo() + result.lowerIndex = self.lower_index + result.higherIndex = self.higher_index + result.mapLowerToHigher = self.map_lower_to_higher + result.mapHigherToLower = self.map_higher_to_lower + result.lowerToHigherDirect = self.lower_to_higher_direct + result.higherToLowerDirect = self.higher_to_lower_direct + return result + + +LLILSSAToMLILExpressionMapping = List['LLILSSAToMLILExpressionMap'] @dataclass(frozen=True, repr=False, order=True) @@ -3300,6 +3324,11 @@ class MediumLevelILFunction: self._arch = _arch self._source_function = _source_function + self._mlil_to_mlil_expr_map: dict['MediumLevelILInstruction', List[Tuple[ExpressionIndex, bool]]] = {} + self._mlil_to_mlil_instr_map: dict['MediumLevelILInstruction', List[Tuple[InstructionIndex, bool]]] = {} + self._llil_ssa_to_mlil_expr_map: dict['lowlevelil.LowLevelILInstruction', List[Tuple[ExpressionIndex, bool]]] = {} + self._llil_ssa_to_mlil_instr_map: dict['lowlevelil.LowLevelILInstruction', List[Tuple[InstructionIndex, bool]]] = {} + def __del__(self): if core is not None: core.BNFreeMediumLevelILFunction(self.handle) @@ -3546,7 +3575,7 @@ class MediumLevelILFunction: elif isinstance(operation, MediumLevelILOperation): _operation = operation.value if source_location is not None: - return ExpressionIndex(core.BNMediumLevelILAddExprWithLocation( + index = ExpressionIndex(core.BNMediumLevelILAddExprWithLocation( self.handle, _operation, source_location.address, @@ -3558,6 +3587,8 @@ class MediumLevelILFunction: d, e )) + self._record_mlil_to_mlil_expr_map(index, source_location) + return index else: return ExpressionIndex(core.BNMediumLevelILAddExpr(self.handle, _operation, size, a, b, c, d, e)) @@ -3914,7 +3945,6 @@ class MediumLevelILFunction: new_index = do_copy(expr, dest, sub_expr_handler) # Copy expression metadata as well dest.set_expr_attributes(new_index, expr.attributes) - dest.set_expr_type(new_index, self.get_expr_type(expr.expr_index)) return new_index def translate( @@ -3944,7 +3974,7 @@ class MediumLevelILFunction: for instr_index in range(block.start, block.end): instr: MediumLevelILInstruction = self[InstructionIndex(instr_index)] propagated_func.set_current_address(instr.address, block.arch) - propagated_func.append(expr_handler(propagated_func, block, instr)) + propagated_func.append(expr_handler(propagated_func, block, instr), ILSourceLocation.from_instruction(instr)) return propagated_func @@ -3968,7 +3998,7 @@ class MediumLevelILFunction: result |= flag.value core.BNSetMediumLevelILExprAttributes(self.handle, expr, result) - def append(self, expr: ExpressionIndex) -> InstructionIndex: + def append(self, expr: ExpressionIndex, source_location: Optional['ILSourceLocation'] = None) -> InstructionIndex: """ ``append`` adds the ExpressionIndex ``expr`` to the current MediumLevelILFunction. @@ -3976,7 +4006,115 @@ class MediumLevelILFunction: :return: Index of added instruction in the current function :rtype: int """ - return InstructionIndex(core.BNMediumLevelILAddInstruction(self.handle, expr)) + index = InstructionIndex(core.BNMediumLevelILAddInstruction(self.handle, expr)) + self._record_mlil_to_mlil_instr_map(index, source_location) + return index + + def _record_mlil_to_mlil_instr_map(self, index, source_location: 'ILSourceLocation'): + # Update internal mappings to remember this + if source_location is not None: + if source_location.source_mlil_instruction is not None: + if source_location.source_mlil_instruction not in self._mlil_to_mlil_instr_map: + self._mlil_to_mlil_instr_map[source_location.source_mlil_instruction] = [] + self._mlil_to_mlil_instr_map[source_location.source_mlil_instruction].append((index, source_location.il_direct)) + if source_location.source_llil_instruction is not None \ + and source_location.source_llil_instruction.function.il_form == FunctionGraphType.LowLevelILSSAFormFunctionGraph: + if source_location.source_llil_instruction not in self._llil_ssa_to_mlil_instr_map: + self._llil_ssa_to_mlil_instr_map[source_location.source_llil_instruction] = [] + self._llil_ssa_to_mlil_instr_map[source_location.source_llil_instruction].append((index, source_location.il_direct)) + + def _record_mlil_to_mlil_expr_map(self, index, source_location: 'ILSourceLocation'): + # Update internal mappings to remember this + if source_location.source_mlil_instruction is not None: + if source_location.source_mlil_instruction not in self._mlil_to_mlil_expr_map: + self._mlil_to_mlil_expr_map[source_location.source_mlil_instruction] = [] + self._mlil_to_mlil_expr_map[source_location.source_mlil_instruction].append((index, source_location.il_direct)) + if source_location.source_llil_instruction is not None \ + and source_location.source_llil_instruction.function.il_form == FunctionGraphType.LowLevelILSSAFormFunctionGraph: + if source_location.source_llil_instruction not in self._llil_ssa_to_mlil_expr_map: + self._llil_ssa_to_mlil_expr_map[source_location.source_llil_instruction] = [] + self._llil_ssa_to_mlil_expr_map[source_location.source_llil_instruction].append((index, source_location.il_direct)) + + def _get_llil_ssa_to_mlil_instr_map(self, from_builders: bool) -> LLILSSAToMLILInstructionMapping: + llil_ssa_to_mlil_instr_map = {} + + if from_builders: + for (old_instr, new_indices) in self._mlil_to_mlil_instr_map.items(): + old_instr: MediumLevelILInstruction + new_indices: List[InstructionIndex] + + # Look up the LLIL SSA instruction for the old instr in its function + # And then store that mapping for the new function + + for (new_index, new_direct) in new_indices: + # Instructions are always mapped 1 to 1. If the map is marked indirect + # then just ignore it + if new_direct: + old_llil_ssa_index = old_instr.function.get_low_level_il_instruction_index(old_instr.instr_index) + if old_llil_ssa_index is not None: + llil_ssa_to_mlil_instr_map[old_llil_ssa_index] = new_index + else: + for instr in self.instructions: + llil_ssa_index = self.get_low_level_il_instruction_index(instr.instr_index) + llil_ssa_to_mlil_instr_map[llil_ssa_index] = instr.instr_index + + return llil_ssa_to_mlil_instr_map + + def _get_llil_ssa_to_mlil_expr_map(self, from_builders: bool) -> LLILSSAToMLILExpressionMapping: + llil_ssa_to_mlil_expr_map = [] + + if from_builders: + for (old_expr, new_indices) in self._mlil_to_mlil_expr_map.items(): + old_expr: MediumLevelILInstruction + new_indices: List[ExpressionIndex] + + # Look up the LLIL SSA expression for the old expr in its function + # And then store that mapping for the new function + + old_llil_ssa_direct = old_expr.function.get_low_level_il_expr_index(old_expr.expr_index) + old_llil_ssa_indices = old_expr.function.get_low_level_il_expr_indexes(old_expr.expr_index) + for old_index in old_llil_ssa_indices: + old_reverse_direct = old_expr.function.low_level_il.ssa_form.get_medium_level_il_expr_index(old_index) + old_reverse_all = old_expr.function.low_level_il.ssa_form.get_medium_level_il_expr_indexes(old_index) + + for (new_index, new_direct) in new_indices: + lower_to_higher_direct = new_direct and old_reverse_direct == old_expr.expr_index + higher_to_lower_direct = new_direct and old_index == old_llil_ssa_direct + map_lower_to_higher = old_expr.expr_index in old_reverse_all + map_higher_to_lower = True + + llil_ssa_to_mlil_expr_map.append(LLILSSAToMLILExpressionMap( + old_index, + new_index, + map_lower_to_higher, + map_higher_to_lower, + lower_to_higher_direct, + higher_to_lower_direct + )) + else: + for instr in self.instructions: + for expr in instr.traverse(lambda e: e): + llil_ssa_direct = self.get_low_level_il_expr_index(expr.expr_index) + llil_ssa_indices = self.get_low_level_il_expr_indexes(expr.expr_index) + for llil_ssa_index in llil_ssa_indices: + reverse_direct = self.low_level_il.ssa_form.get_medium_level_il_expr_index(llil_ssa_index) + reverse_all = self.low_level_il.ssa_form.get_medium_level_il_expr_indexes(llil_ssa_index) + + lower_to_higher_direct = reverse_direct == expr.expr_index + higher_to_lower_direct = llil_ssa_index == llil_ssa_direct + map_lower_to_higher = expr.expr_index in reverse_all + map_higher_to_lower = True + + llil_ssa_to_mlil_expr_map.append(LLILSSAToMLILExpressionMap( + llil_ssa_index, + expr.expr_index, + map_lower_to_higher, + map_higher_to_lower, + lower_to_higher_direct, + higher_to_lower_direct + )) + + return llil_ssa_to_mlil_expr_map def nop(self, loc: Optional['ILSourceLocation'] = None) -> ExpressionIndex: """ @@ -5619,7 +5757,9 @@ class MediumLevelILFunction: :rtype: ExpressionIndex """ if loc is not None: - return ExpressionIndex(core.BNMediumLevelILGotoWithLocation(self.handle, label.handle, loc.address, loc.source_operand)) + index = ExpressionIndex(core.BNMediumLevelILGotoWithLocation(self.handle, label.handle, loc.address, loc.source_operand)) + self._record_mlil_to_mlil_expr_map(index, loc) + return index else: return ExpressionIndex(core.BNMediumLevelILGoto(self.handle, label.handle)) @@ -5639,7 +5779,9 @@ class MediumLevelILFunction: :rtype: ExpressionIndex """ if loc is not None: - return ExpressionIndex(core.BNMediumLevelILIfWithLocation(self.handle, operand, t.handle, f.handle, loc.address, loc.source_operand)) + index = ExpressionIndex(core.BNMediumLevelILIfWithLocation(self.handle, operand, t.handle, f.handle, loc.address, loc.source_operand)) + self._record_mlil_to_mlil_expr_map(index, loc) + return index else: return ExpressionIndex(core.BNMediumLevelILIf(self.handle, operand, t.handle, f.handle)) diff --git a/python/workflow.py b/python/workflow.py index c63c20b9..bf48e5d8 100644 --- a/python/workflow.py +++ b/python/workflow.py @@ -112,7 +112,53 @@ class AnalysisContext: @mlil.setter def mlil(self, value: mediumlevelil.MediumLevelILFunction) -> None: - core.BNSetMediumLevelILFunction(self.handle, value.handle) + self.set_mlil_function(value) + + def set_mlil_function( + self, + new_func: mediumlevelil.MediumLevelILFunction, + llil_ssa_to_mlil_instr_map: Optional['mediumlevelil.LLILSSAToMLILInstructionMapping'] = None, + llil_ssa_to_mlil_expr_map: Optional['mediumlevelil.LLILSSAToMLILExpressionMapping'] = None, + ) -> None: + """ + Set the Medium Level IL function in the current analysis, giving updated + Low Level IL (SSA) to Medium Level IL instruction and expression mappings. + :param new_func: New MLIL function + :param llil_ssa_to_mlil_instr_map: Mapping from every LLIL SSA instruction to + every MLIL instruction + :param llil_ssa_to_mlil_expr_map: Mapping from every LLIL SSA expression to + one or more MLIL expressions (first expression + will be the primary) + """ + if llil_ssa_to_mlil_instr_map is None or llil_ssa_to_mlil_expr_map is None: + # Build up maps from existing data in the function + llil_ssa_to_mlil_instr_map = new_func._get_llil_ssa_to_mlil_instr_map(True) + llil_ssa_to_mlil_expr_map = new_func._get_llil_ssa_to_mlil_expr_map(True) + + # Number of instructions + instr_count = 0 + if len(llil_ssa_to_mlil_instr_map) > 0: + instr_count = max(llil_ssa_to_mlil_instr_map.keys()) + 1 + ffi_instr_map = (ctypes.c_size_t * instr_count)() + for i in range(instr_count): + ffi_instr_map[i] = 0xffffffffffffffff + for (key, value) in llil_ssa_to_mlil_instr_map.items(): + ffi_instr_map[key] = value + + # Number of map entries, not highest index + expr_count = len(llil_ssa_to_mlil_expr_map) + ffi_expr_map = (core.BNExprMapInfo * expr_count)() + for i, map in enumerate(llil_ssa_to_mlil_expr_map): + ffi_expr_map[i] = map._to_core_struct() + + core.BNSetMediumLevelILFunction( + self.handle, + new_func.handle, + ffi_instr_map, + instr_count, + ffi_expr_map, + expr_count + ) @property def hlil(self) -> Optional[highlevelil.HighLevelILFunction]: -- cgit v1.3.1