Add `replace_expr` to MLIL/HLIL in Python API

3 files changed, 45 insertions, 1 deletions

diff --git a/python/highlevelil.py b/python/highlevelil.py
index fc9477cd..ca6a8aaf 100644
--- a/python/highlevelil.py
+++ b/python/highlevelil.py
@@ -50,6 +50,8 @@ TokenList = List['function.InstructionTextToken']
 LinesType = Generator['function.DisassemblyTextLine', None, None]
 ExpressionIndex = NewType('ExpressionIndex', int)
 InstructionIndex = NewType('InstructionIndex', int)
+Index = Union[ExpressionIndex, InstructionIndex]
+InstructionOrExpression = Union['HighLevelILInstruction', Index]
 HLILInstructionsType = Generator['HighLevelILInstruction', None, None]
 HLILBasicBlocksType = Generator['HighLevelILBasicBlock', None, None]
 OperandsType = Tuple[ExpressionIndex, ExpressionIndex, ExpressionIndex, ExpressionIndex, ExpressionIndex]
@@ -2396,6 +2398,26 @@ class HighLevelILFunction:
 			operation_value = operation.value
 		return ExpressionIndex(core.BNHighLevelILAddExpr(self.handle, operation_value, size, a, b, c, d, e))
 
+	def replace_expr(self, original: InstructionOrExpression, new: InstructionOrExpression) -> None:
+		"""
+		``replace_expr`` allows modification of HLIL expressions
+
+		:param ExpressionIndex original: the ExpressionIndex to replace (may also be an expression index)
+		:param ExpressionIndex new: the ExpressionIndex to add to the current LowLevelILFunction (may also be an expression index)
+		:rtype: None
+		"""
+		if isinstance(original, HighLevelILInstruction):
+			original = original.expr_index
+		elif isinstance(original, int):
+			original = ExpressionIndex(original)
+
+		if isinstance(new, HighLevelILInstruction):
+			new = new.expr_index
+		elif isinstance(new, int):
+			new = ExpressionIndex(new)
+
+		core.BNReplaceHighLevelILExpr(self.handle, original, new)
+
 	def add_operand_list(self, operands: List[int]) -> ExpressionIndex:
 		"""
 		``add_operand_list`` returns an operand list expression for the given list of integer operands.
diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index 39461500..4a45514a 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -3146,7 +3146,7 @@ class LowLevelILFunction:
 
 	def replace_expr(self, original: InstructionOrExpression, new: InstructionOrExpression) -> None:
 		"""
-		``replace_expr`` allows modification of ExpressionIndexessions but ONLY during lifting.
+		``replace_expr`` allows modification of expressions but ONLY during lifting.
 
 		.. warning:: This function should ONLY be called as a part of a lifter. It will otherwise not do anything useful as there's no way to trigger re-analysis of IL levels at this time.
 
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index 937fc07e..49c932f3 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -48,6 +48,8 @@ from .commonil import (
 TokenList = List['function.InstructionTextToken']
 ExpressionIndex = NewType('ExpressionIndex', int)
 InstructionIndex = NewType('InstructionIndex', int)
+Index = Union[ExpressionIndex, InstructionIndex]
+InstructionOrExpression = Union['MediumLevelILInstruction', Index]
 MLILInstructionsType = Generator['MediumLevelILInstruction', None, None]
 MLILBasicBlocksType = Generator['MediumLevelILBasicBlock', None, None]
 OperandsType = Tuple[ExpressionIndex, ExpressionIndex, ExpressionIndex, ExpressionIndex, ExpressionIndex]
@@ -2879,6 +2881,26 @@ class MediumLevelILFunction:
 			_operation = operation.value
 		return ExpressionIndex(core.BNMediumLevelILAddExpr(self.handle, _operation, size, a, b, c, d, e))
 
+	def replace_expr(self, original: InstructionOrExpression, new: InstructionOrExpression) -> None:
+		"""
+		``replace_expr`` allows modification of MLIL expressions
+
+		:param ExpressionIndex original: the ExpressionIndex to replace (may also be an expression index)
+		:param ExpressionIndex new: the ExpressionIndex to add to the current LowLevelILFunction (may also be an expression index)
+		:rtype: None
+		"""
+		if isinstance(original, MediumLevelILInstruction):
+			original = original.expr_index
+		elif isinstance(original, int):
+			original = ExpressionIndex(original)
+
+		if isinstance(new, MediumLevelILInstruction):
+			new = new.expr_index
+		elif isinstance(new, int):
+			new = ExpressionIndex(new)
+
+		core.BNReplaceMediumLevelILExpr(self.handle, original, new)
+
 	def append(self, expr: ExpressionIndex) -> int:
 		"""
 		``append`` adds the ExpressionIndex ``expr`` to the current MediumLevelILFunction.