Python: Collect and pass mappings when copying MLIL functions

1 files changed, 44 insertions, 5 deletions

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"<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