WIP: Python: Collect and pass mappings when copying MLIL functions

Author: CouleeApps

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
@@ -210,7 +210,6 @@ class AliasedVariableInstruction(VariableInstruction):
 	pass
 
 
-@dataclass
 class ILSourceLocation:
 	"""
 	ILSourceLocation is used to indicate where expressions were defined during the lifting process
@@ -220,15 +219,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

@@ -67,7 +67,7 @@ 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):
+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 +78,12 @@ 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
+    assert old_insn.expr_type == new_insn.expr_type, 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)):
@@ -241,30 +243,45 @@ def translate_instr(
             assert_mlil_eq(old_insn, new_insn)
 
 
-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
 ))

python/mediumlevelil.py

@@ -3302,6 +3302,11 @@ def __init__(
 		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)
@@ -3548,7 +3553,7 @@ def expr(
 		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,
@@ -3560,6 +3565,17 @@ def expr(
 				d,
 				e
 			))
+			# 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))
+			return index
 		else:
 			return ExpressionIndex(core.BNMediumLevelILAddExpr(self.handle, _operation, size, a, b, c, d, e))
 
@@ -3946,7 +3962,7 @@ def translate(
 			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
 
@@ -3970,15 +3986,85 @@ def set_expr_attributes(self, expr: InstructionOrExpression, value: ILInstructio
 			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.
 
 		:param ExpressionIndex expr: the ExpressionIndex to add to the current 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))
+
+		# 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))
+		return index
+
+	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)
+						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:
+					if old_index not in llil_ssa_to_mlil_expr_map:
+						llil_ssa_to_mlil_expr_map[old_index] = []
+					for (new_index, new_direct) in new_indices:
+						llil_ssa_to_mlil_expr_map[old_index].append(new_index)
+		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:
+						if llil_ssa_index not in llil_ssa_to_mlil_expr_map:
+							llil_ssa_to_mlil_expr_map[llil_ssa_index] = []
+						if llil_ssa_index == llil_ssa_direct:
+							llil_ssa_to_mlil_expr_map[llil_ssa_index].insert(0, expr.expr_index)
+						else:
+							llil_ssa_to_mlil_expr_map[llil_ssa_index].append(expr.expr_index)
+
+		return llil_ssa_to_mlil_expr_map
 
 	def nop(self, loc: Optional['ILSourceLocation'] = None) -> ExpressionIndex:
 		"""

python/workflow.py

@@ -130,10 +130,10 @@ def set_mlil_function(
 		                                  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:
-			llil_ssa_to_mlil_instr_map = {}
-			llil_ssa_to_mlil_expr_map = {}
+			# 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