Add SetExprType API to MLIL/HLIL

Add get/set_expr_type to MLIL/HLIL Python APi

6 files changed, 130 insertions, 0 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 46e3fedf..fa66a799 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -8866,6 +8866,19 @@ namespace BinaryNinja {
 		Confidence<Ref<Type>> GetExprType(size_t expr);
 		Confidence<Ref<Type>> GetExprType(const MediumLevelILInstruction& expr);
 
+		/*! SetExprType sets the type of a given expression.
+
+			\warning This method is only meant for workflows or for debugging purposes, since the changes they make
+			are not persistent and get lost after a database save and reload. To make persistent changes to the analysis,
+			one should use other APIs to, for example, change the type of variables. The analysis will then propagate the
+			type of the variable and update the type of related expressions.
+
+		    \param expr index of the expression to set
+		    \param type new type of the expression
+		*/
+		void SetExprType(size_t expr, const Confidence<Ref<Type>>& type);
+		void SetExprType(const MediumLevelILInstruction& expr, const Confidence<Ref<Type>>& type);
+
 		static bool IsConstantType(BNMediumLevelILOperation op)
 		{
 			return op == MLIL_CONST || op == MLIL_CONST_PTR || op == MLIL_EXTERN_PTR;
@@ -9133,6 +9146,19 @@ namespace BinaryNinja {
 		Confidence<Ref<Type>> GetExprType(size_t expr);
 		Confidence<Ref<Type>> GetExprType(const HighLevelILInstruction& expr);
 
+		/*! SetExprType sets the type of a given expression.
+
+			\warning This method is only meant for workflows or for debugging purposes, since the changes they make
+			are not persistent and get lost after a database save and reload. To make persistent changes to the analysis,
+			one should use other APIs to, for example, change the type of variables. The analysis will then propagate the
+			type of the variable and update the type of related expressions.
+
+		    \param expr index of the expression to set
+		    \param type new type of the expression
+		*/
+		void SetExprType(size_t expr, const Confidence<Ref<Type>>& type);
+		void SetExprType(const HighLevelILInstruction& expr, const Confidence<Ref<Type>>& type);
+
 		void VisitAllExprs(const std::function<bool(const HighLevelILInstruction& expr)>& func);
 
 		Ref<FlowGraph> CreateFunctionGraph(DisassemblySettings* settings = nullptr);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index f7d1a75f..ac0d90d6 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -5006,6 +5006,7 @@ extern "C"
 	BINARYNINJACOREAPI size_t* BNGetHighLevelILExprIndexes(BNMediumLevelILFunction* func, size_t expr, size_t* count);
 
 	BINARYNINJACOREAPI BNTypeWithConfidence BNGetMediumLevelILExprType(BNMediumLevelILFunction* func, size_t expr);
+	BINARYNINJACOREAPI void BNSetMediumLevelILExprType(BNMediumLevelILFunction* func, size_t expr, BNTypeWithConfidence* type);
 
 	// High-level IL
 	BINARYNINJACOREAPI BNHighLevelILFunction* BNCreateHighLevelILFunction(BNArchitecture* arch, BNFunction* func);
@@ -5052,6 +5053,7 @@ extern "C"
 	    BNHighLevelILFunction* func, size_t expr, bool asFullAst, size_t* count, BNDisassemblySettings* settings);
 
 	BINARYNINJACOREAPI BNTypeWithConfidence BNGetHighLevelILExprType(BNHighLevelILFunction* func, size_t expr);
+	BINARYNINJACOREAPI void BNSetHighLevelILExprType(BNHighLevelILFunction* func, size_t expr, BNTypeWithConfidence* type);
 
 	BINARYNINJACOREAPI BNBasicBlock** BNGetHighLevelILBasicBlockList(BNHighLevelILFunction* func, size_t* count);
 	BINARYNINJACOREAPI BNBasicBlock* BNGetHighLevelILBasicBlockForInstruction(BNHighLevelILFunction* func, size_t i);
diff --git a/highlevelil.cpp b/highlevelil.cpp
index d0b3b4e2..7d3789d9 100644
--- a/highlevelil.cpp
+++ b/highlevelil.cpp
@@ -501,6 +501,22 @@ Confidence<Ref<Type>> HighLevelILFunction::GetExprType(const HighLevelILInstruct
 }
 
 
+void HighLevelILFunction::SetExprType(size_t expr, const Confidence<Ref<Type>>& type)
+{
+	BNTypeWithConfidence tc;
+	tc.type = type->GetObject();
+	tc.confidence = type.GetConfidence();
+	BNSetHighLevelILExprType(m_object, expr, &tc);
+}
+
+
+void HighLevelILFunction::SetExprType(const BinaryNinja::HighLevelILInstruction& expr,
+										const Confidence<Ref<BinaryNinja::Type>>& type)
+{
+	SetExprType(expr.exprIndex, type);
+}
+
+
 void HighLevelILFunction::VisitAllExprs(const function<bool(const HighLevelILInstruction& expr)>& func)
 {
 	GetRootExpr().VisitExprs([&](const HighLevelILInstruction& expr) { return func(expr); });
diff --git a/mediumlevelil.cpp b/mediumlevelil.cpp
index cc619708..ba2c94a3 100644
--- a/mediumlevelil.cpp
+++ b/mediumlevelil.cpp
@@ -842,6 +842,22 @@ Confidence<Ref<Type>> MediumLevelILFunction::GetExprType(const MediumLevelILInst
 }
 
 
+void MediumLevelILFunction::SetExprType(size_t expr, const Confidence<Ref<Type>>& type)
+{
+	BNTypeWithConfidence tc;
+	tc.type = type->GetObject();
+	tc.confidence = type.GetConfidence();
+	BNSetMediumLevelILExprType(m_object, expr, &tc);
+}
+
+
+void MediumLevelILFunction::SetExprType(const BinaryNinja::MediumLevelILInstruction& expr,
+										const Confidence<Ref<BinaryNinja::Type>>& type)
+{
+	SetExprType(expr.exprIndex, type);
+}
+
+
 Ref<FlowGraph> MediumLevelILFunction::CreateFunctionGraph(DisassemblySettings* settings)
 {
 	BNFlowGraph* graph = BNCreateMediumLevelILFunctionGraph(m_object, settings ? settings->GetObject() : nullptr);
diff --git a/python/highlevelil.py b/python/highlevelil.py
index 451942fc..bb62840b 100644
--- a/python/highlevelil.py
+++ b/python/highlevelil.py
@@ -56,6 +56,7 @@ HighLevelILOperandType = Union['HighLevelILInstruction', 'lowlevelil.ILIntrinsic
                                List['mediumlevelil.SSAVariable'], List['HighLevelILInstruction'], Optional[int], float,
                                'GotoLabel']
 VariablesList = List[Union['mediumlevelil.SSAVariable', 'variable.Variable']]
+StringOrType = Union[str, '_types.Type', '_types.TypeBuilder']
 
 
 class VariableReferenceType(Enum):
@@ -2556,6 +2557,40 @@ class HighLevelILFunction:
 		core.BNFreeILInstructionList(uses)
 		return result
 
+	def get_expr_type(self, expr_index: int) -> Optional['types.Type']:
+		"""
+		Get type of expression
+
+		:param int expr_index: index of the expression to retrieve
+		:rtype: Optional['types.Type']
+		"""
+		result = core.BNGetHighLevelILExprType(self.handle, expr_index)
+		if result.type:
+			platform = None
+			if self.source_function:
+				platform = self.source_function.platform
+			return types.Type.create(
+				core.BNNewTypeReference(result.type), platform=platform, confidence=result.confidence
+			)
+		return None
+
+	def set_expr_type(self, expr_index: int, expr_type: StringOrType) -> None:
+		"""
+		Set type of expression
+
+		This API is only meant for workflows or for debugging purposes, since the changes they make are not persistent
+		and get lost after a database save and reload. To make persistent changes to the analysis, one should use other
+		APIs to, for example, change the type of variables. The analysis will then propagate the type of the variable
+		and update the type of related expressions.
+
+		:param int expr_index: index of the expression to set
+		:param StringOrType: new type of the expression
+		"""
+		if isinstance(expr_type, str):
+			(expr_type, _) = self.view.parse_type_string(expr_type)
+		tc = expr_type._to_core_struct()
+		core.BNSetHighLevelILExprType(self.handle, expr_index, tc)
+
 
 class HighLevelILBasicBlock(basicblock.BasicBlock):
 	"""
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index 47e7e2e1..b665fef8 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -53,6 +53,7 @@ MediumLevelILOperandType = Union[int, float, 'MediumLevelILOperationAndSize', 'M
                                  'lowlevelil.ILIntrinsic', 'variable.Variable', 'SSAVariable', List[int],
                                  List['variable.Variable'], List['SSAVariable'], List['MediumLevelILInstruction'],
                                  Mapping[int, int]]
+StringOrType = Union[str, '_types.Type', '_types.TypeBuilder']
 
 
 @dataclass(frozen=True, repr=False, order=True)
@@ -3260,6 +3261,40 @@ class MediumLevelILFunction:
 
 		return []
 
+	def get_expr_type(self, expr_index: int) -> Optional['types.Type']:
+		"""
+		Get type of expression
+
+		:param int expr_index: index of the expression to retrieve
+		:rtype: Optional['types.Type']
+		"""
+		result = core.BNGetMediumLevelILExprType(self.handle, expr_index)
+		if result.type:
+			platform = None
+			if self.source_function:
+				platform = self.source_function.platform
+			return types.Type.create(
+				core.BNNewTypeReference(result.type), platform=platform, confidence=result.confidence
+			)
+		return None
+
+	def set_expr_type(self, expr_index: int, expr_type: StringOrType) -> None:
+		"""
+		Set type of expression
+
+		This API is only meant for workflows or for debugging purposes, since the changes they make are not persistent
+		and get lost after a database save and reload. To make persistent changes to the analysis, one should use other
+		APIs to, for example, change the type of variables. The analysis will then propagate the type of the variable
+		and update the type of related expressions.
+
+		:param int expr_index: index of the expression to set
+		:param StringOrType: new type of the expression
+		"""
+		if isinstance(expr_type, str):
+			(expr_type, _) = self.view.parse_type_string(expr_type)
+		tc = expr_type._to_core_struct()
+		core.BNSetMediumLevelILExprType(self.handle, expr_index, tc)
+
 
 class MediumLevelILBasicBlock(basicblock.BasicBlock):
 	"""