Python: Collect and pass mappings when copying MLIL functions

Author: CouleeApps

binaryninjacore.h

@@ -3711,6 +3711,16 @@ extern "C"
 		AlwaysEnabledRenderLayerDefaultEnableState,
 	} BNRenderLayerDefaultEnableState;
 
+	typedef struct BNExprMapInfo
+	{
+		size_t lowerIndex;
+		size_t higherIndex;
+		bool mapLowerToHigher;
+		bool mapHigherToLower;
+		bool lowerToHigherDirect;
+		bool higherToLowerDirect;
+	} BNExprMapInfo;
+
 	BINARYNINJACOREAPI char* BNAllocString(const char* contents);
 	BINARYNINJACOREAPI char* BNAllocStringWithLength(const char* contents, size_t len);
 	BINARYNINJACOREAPI void BNFreeString(char* str);
@@ -5647,7 +5657,9 @@ extern "C"
 	BINARYNINJACOREAPI void BNSetLowLevelILFunction(
 	    BNAnalysisContext* analysisContext, BNLowLevelILFunction* lowLevelIL);
 	BINARYNINJACOREAPI void BNSetMediumLevelILFunction(
-	    BNAnalysisContext* analysisContext, BNMediumLevelILFunction* mediumLevelIL);
+	    BNAnalysisContext* analysisContext, BNMediumLevelILFunction* mediumLevelIL,
+	    size_t* llilSSAToMLILSSAInstrMap, size_t instrCount,
+	    BNExprMapInfo* llilSSAToMLILSSAExprMap, size_t exprCount);
 	BINARYNINJACOREAPI void BNSetHighLevelILFunction(
 	    BNAnalysisContext* analysisContext, BNHighLevelILFunction* highLevelIL);
 	BINARYNINJACOREAPI bool BNAnalysisContextInform(BNAnalysisContext* analysisContext, const char* request);

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 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"<ILSourceLocation: {self.address:x}, {self.source_operand}{instr}>"
+
+	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

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 translate_instr(
         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,31 +266,52 @@ def translate_instr(
         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
 ))
 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
 ))
 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
 ))