Python: Add ability to construct expressions with locations

Co-Authored-By: ltlly <a1253213025@163.com>

4 files changed, 83 insertions, 9 deletions

diff --git a/python/commonil.py b/python/commonil.py
index b665165d..74399184 100644
--- a/python/commonil.py
+++ b/python/commonil.py
@@ -19,10 +19,14 @@
 # IN THE SOFTWARE.
 
 from dataclasses import dataclass
+from typing import Union
 from .flowgraph import FlowGraph, FlowGraphNode
 from .enums import BranchType
 from .interaction import show_graph_report
 from .log import log_warn
+from . import lowlevelil
+from . import mediumlevelil
+from . import highlevelil
 
 
 # This file contains a list of top level abstract classes for implementing BNIL instructions
@@ -204,3 +208,27 @@ class SSAVariableInstruction(SSA, VariableInstruction):
 @dataclass(frozen=True, repr=False, eq=False)
 class AliasedVariableInstruction(VariableInstruction):
 	pass
+
+
+@dataclass
+class ILSourceLocation:
+	"""
+	ILSourceLocation is used to indicate where expressions were defined during the lifting process
+	and gets propagated through the lifting process as an instruction's address/source_operand properties.
+	These are used for, for example, integer display types and expression addresses.
+	"""
+	address: int
+	source_operand: int
+
+	@classmethod
+	def from_instruction(
+			cls,
+			instr: Union['lowlevelil.LowLevelILInstruction', 'mediumlevelil.MediumLevelILInstruction', 'highlevelil.HighLevelILInstruction']
+	) -> '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)
+
diff --git a/python/highlevelil.py b/python/highlevelil.py
index 154af642..c73df7f0 100644
--- a/python/highlevelil.py
+++ b/python/highlevelil.py
@@ -46,7 +46,7 @@ from .interaction import show_graph_report
 from .commonil import (
     BaseILInstruction, Tailcall, Syscall, Localcall, Comparison, Signed, UnaryOperation, BinaryOperation, SSA, Phi,
     Loop, ControlFlow, Memory, Constant, Arithmetic, DoublePrecision, Terminal, FloatingPoint, Intrinsic, Return,
-    VariableInstruction, SSAVariableInstruction, SetVar
+    VariableInstruction, SSAVariableInstruction, SetVar, ILSourceLocation
 )
 from . import deprecation
 
@@ -2851,14 +2851,29 @@ class HighLevelILFunction:
 
 	def expr(
 	    self, operation: Union[str, HighLevelILOperation], a: int = 0, b: int = 0, c: int = 0, d: int = 0, e: int = 0,
-	    size: int = 0
+	    size: int = 0,
+	    source_location: Optional['ILSourceLocation'] = None
 	) -> ExpressionIndex:
 		if isinstance(operation, str):
 			operation_value = HighLevelILOperation[operation]
 		else:
 			assert isinstance(operation, HighLevelILOperation)
 			operation_value = operation.value
-		return ExpressionIndex(core.BNHighLevelILAddExpr(self.handle, operation_value, size, a, b, c, d, e))
+		if source_location is not None:
+			return ExpressionIndex(core.BNHighLevelILAddExprWithLocation(
+				self.handle,
+				operation_value,
+				source_location.address,
+				source_location.source_operand,
+				size,
+				a,
+				b,
+				c,
+				d,
+				e
+			))
+		else:
+			return ExpressionIndex(core.BNHighLevelILAddExpr(self.handle, operation_value, size, a, b, c, d, e))
 
 	def get_expr_count(self) -> int:
 		"""
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index 4c2f0de9..76ee18cf 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -40,7 +40,7 @@ from .interaction import show_graph_report
 from .commonil import (
     BaseILInstruction, Constant, BinaryOperation, Tailcall, UnaryOperation, Comparison, SSA, Phi, FloatingPoint,
     ControlFlow, Terminal, Syscall, Localcall, StackOperation, Return, Signed, Arithmetic, Carry, DoublePrecision,
-    Memory, Load, Store, RegisterStack, SetReg, Intrinsic
+    Memory, Load, Store, RegisterStack, SetReg, Intrinsic, ILSourceLocation
 )
 
 ExpressionIndex = NewType('ExpressionIndex', int)
@@ -3779,7 +3779,8 @@ class LowLevelILFunction:
 
 	def expr(
 	    self, operation, a: ExpressionIndex = 0, b: ExpressionIndex = 0, c: ExpressionIndex = 0, d: ExpressionIndex = 0, size: int = 0,
-	    flags: Optional[Union['architecture.FlagWriteTypeName', 'architecture.FlagType', 'architecture.FlagIndex']] = None
+	    flags: Optional[Union['architecture.FlagWriteTypeName', 'architecture.FlagType', 'architecture.FlagIndex']] = None,
+	    source_location: Optional['ILSourceLocation'] = None
 	) -> ExpressionIndex:
 		_flags = architecture.FlagIndex(0)
 		if isinstance(operation, str):
@@ -3796,7 +3797,21 @@ class LowLevelILFunction:
 			_flags = architecture.FlagIndex(0)
 		else:
 			assert False, "flags type unsupported"
-		return ExpressionIndex(core.BNLowLevelILAddExpr(self.handle, operation, size, _flags, a, b, c, d))
+		if source_location is not None:
+			return ExpressionIndex(core.BNLowLevelILAddExprWithLocation(
+				self.handle,
+				source_location.address,
+				source_location.source_operand,
+				operation,
+				size,
+				_flags,
+				a,
+				b,
+				c,
+				d
+			))
+		else:
+			return ExpressionIndex(core.BNLowLevelILAddExpr(self.handle, operation, size, _flags, a, b, c, d))
 
 	def get_expr_count(self) -> int:
 		"""
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index 4ec79d8d..9a73e95f 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -42,7 +42,8 @@ from .interaction import show_graph_report
 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
+    Store, RegisterStack, SetVar, Intrinsic, VariableInstruction, SSAVariableInstruction, AliasedVariableInstruction,
+    ILSourceLocation
 )
 
 TokenList = List['function.InstructionTextToken']
@@ -3494,14 +3495,29 @@ class MediumLevelILFunction:
 
 	def expr(
 	    self, operation: MediumLevelILOperation, a: int = 0, b: int = 0, c: int = 0, d: int = 0, e: int = 0,
-	    size: int = 0
+	    size: int = 0,
+	    source_location: Optional['ILSourceLocation'] = None
 	) -> ExpressionIndex:
 		_operation = operation
 		if isinstance(operation, str):
 			_operation = MediumLevelILOperation[operation]
 		elif isinstance(operation, MediumLevelILOperation):
 			_operation = operation.value
-		return ExpressionIndex(core.BNMediumLevelILAddExpr(self.handle, _operation, size, a, b, c, d, e))
+		if source_location is not None:
+			return ExpressionIndex(core.BNMediumLevelILAddExprWithLocation(
+				self.handle,
+				_operation,
+				source_location.address,
+				source_location.source_operand,
+				size,
+				a,
+				b,
+				c,
+				d,
+				e
+			))
+		else:
+			return ExpressionIndex(core.BNMediumLevelILAddExpr(self.handle, _operation, size, a, b, c, d, e))
 
 	def get_expr_count(self) -> int:
 		"""