Changes to implement User-informed dataflow

4 files changed, 329 insertions, 1 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 195f9089..90bbc4c2 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -2938,13 +2938,16 @@ __attribute__ ((format (printf, 1, 2)))
 		std::vector<BNValueRange> ranges;
 		std::set<int64_t> valueSet;
 		std::vector<LookupTableEntry> table;
+		size_t count;
 
 		static PossibleValueSet FromAPIObject(BNPossibleValueSet& value);
+		BNPossibleValueSet ToAPIObject();
 	};
 
 	class FlowGraph;
 	class MediumLevelILFunction;
 	class HighLevelILFunction;
+	struct SSAVariable;
 
 	class Function: public CoreRefCountObject<BNFunction, BNNewFunctionReference, BNFreeFunction>
 	{
@@ -3160,6 +3163,10 @@ __attribute__ ((format (printf, 1, 2)))
 
 		Ref<FlowGraph> GetUnresolvedStackAdjustmentGraph();
 
+		void SetVariableValue(const Variable& var, uint64_t defAddr, PossibleValueSet& value);
+		void ClearInformedVariableValue(const Variable& var, uint64_t defAddr);
+		void ClearInformedVariableValues();
+
 		void RequestDebugReport(const std::string& name);
 
 		std::string GetGotoLabelName(uint64_t labelId);
@@ -3702,7 +3709,6 @@ __attribute__ ((format (printf, 1, 2)))
 	};
 
 	struct MediumLevelILInstruction;
-	struct SSAVariable;
 
 	class MediumLevelILFunction: public CoreRefCountObject<BNMediumLevelILFunction,
 		BNNewMediumLevelILFunctionReference, BNFreeMediumLevelILFunction>
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 4128372d..6c32545a 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -257,6 +257,8 @@ extern "C"
 		StructureHexDumpTextToken = 27,
 		GotoLabelToken = 28,
 		CommentToken = 29,
+		PossibleValueToken = 30,
+		PossibleValueTypeToken = 31,
 		// The following are output by the analysis system automatically, these should
 		// not be used directly by the architecture plugins
 		CodeSymbolToken = 64,
@@ -3368,6 +3370,10 @@ __attribute__ ((format (printf, 1, 2)))
 
 	BINARYNINJACOREAPI BNFlowGraph* BNGetUnresolvedStackAdjustmentGraph(BNFunction* func);
 
+	BINARYNINJACOREAPI void BNSetVariableValue(BNFunction* func, const BNVariable* var, const BNArchitectureAndAddress* defSite, const BNPossibleValueSet* value);
+	BINARYNINJACOREAPI void BNClearInformedVariableValue(BNFunction* func, const BNVariable* var, const BNArchitectureAndAddress* defSite);
+	BINARYNINJACOREAPI void BNClearInformedVariableValues(BNFunction* func);
+
 	BINARYNINJACOREAPI void BNRequestFunctionDebugReport(BNFunction* func, const char* name);
 
 	// Disassembly settings
diff --git a/function.cpp b/function.cpp
index fd3760b4..9dc9ea21 100644
--- a/function.cpp
+++ b/function.cpp
@@ -19,6 +19,8 @@
 // IN THE SOFTWARE.
 
 #include "binaryninjaapi.h"
+#include "mediumlevelilinstruction.h"
+#include <cstring>
 
 using namespace BinaryNinja;
 using namespace std;
@@ -354,11 +356,68 @@ PossibleValueSet PossibleValueSet::FromAPIObject(BNPossibleValueSet& value)
 		for (size_t i = 0; i < value.count; i++)
 			result.valueSet.insert(value.valueSet[i]);
 	}
+
+	result.count = value.count;
 	BNFreePossibleValueSet(&value);
 	return result;
 }
 
 
+BNPossibleValueSet PossibleValueSet::ToAPIObject ()
+{
+	BNPossibleValueSet result;
+	result.state = state;
+	result.value = value;
+	result.offset = offset;
+	result.count = 0;
+
+	if ((state == SignedRangeValue) || (state == UnsignedRangeValue))
+	{
+		result.ranges = new BNValueRange[ranges.size()];
+		result.count = ranges.size();
+		for (size_t i = 0; i < ranges.size(); i++)
+			result.ranges[i] = ranges[i];
+	}
+	else
+	{
+		result.ranges = nullptr;
+	}
+
+	if (state == LookupTableValue)
+	{
+		result.table = new BNLookupTableEntry[table.size()];
+		result.count = table.size();
+		for (size_t i = 0; i < table.size(); i++)
+		{
+			result.table[i].fromValues = new int64_t[table[i].fromValues.size()];
+			memcpy(result.table[i].fromValues, &table[i].fromValues[0], sizeof(int64_t) *
+				table[i].fromValues.size());
+			result.table[i].fromCount = table[i].fromValues.size();
+			result.table[i].toValue = table[i].toValue;
+		}
+	}
+	else
+	{
+		result.table = nullptr;
+	}
+
+	if ((state == InSetOfValues) || (state == NotInSetOfValues))
+	{
+		result.valueSet = new int64_t[valueSet.size()];
+		result.count = valueSet.size();
+		size_t i = 0;
+		for (auto j : valueSet)
+			result.valueSet[i++] = j;
+	}
+	else
+	{
+		result.valueSet = nullptr;
+	}
+
+	return result;
+}
+
+
 RegisterValue Function::GetRegisterValueAtInstruction(Architecture* arch, uint64_t addr, uint32_t reg)
 {
 	BNRegisterValue value = BNGetRegisterValueAtInstruction(m_object, arch->GetObject(), addr, reg);
@@ -1712,6 +1771,83 @@ Ref<FlowGraph> Function::GetUnresolvedStackAdjustmentGraph()
 }
 
 
+void Function::SetVariableValue(const Variable& var, uint64_t defAddr, PossibleValueSet& value)
+{
+	auto mlil = GetMediumLevelIL();
+	auto varDefs = mlil->GetVariableDefinitions(var);
+	if (varDefs.size() == 0)
+	{
+		LogError("Could not get definition for Variable");
+		return;
+	} 
+	bool found = false;
+	for (const auto& site : varDefs)
+	{
+		if (site == defAddr)
+		{
+			found = true;
+			break;
+		}
+	}
+	if (!found)
+	{
+		LogError("Could not find definition for variable at given address");
+	}
+	auto defSite = BNArchitectureAndAddress();
+	defSite.arch = GetArchitecture()->m_object;
+	defSite.address = defAddr;
+
+	auto var_data = BNVariable();
+	var_data.type = var.type;
+	var_data.index = var.index;
+	var_data.storage = var.storage;
+
+	auto valueObj = value.ToAPIObject();
+
+	BNSetVariableValue(m_object, &var_data, &defSite, &valueObj);
+}
+
+
+void Function::ClearInformedVariableValue(const Variable& var, uint64_t defAddr)
+{
+	auto mlil = GetMediumLevelIL();
+	auto varDefs = mlil->GetVariableDefinitions(var);
+	if (varDefs.size() == 0)
+	{
+		LogError("Could not get definition for Variable");
+	} 
+	bool found = false;
+	for (auto site : varDefs)
+	{
+		if (site == defAddr)
+		{
+			found = true;
+			break;
+		}
+	}
+	if (!found)
+	{
+		LogError("Could not find definition for variable at given address");
+	}
+	auto defSite = BNArchitectureAndAddress();
+	defSite.arch = GetArchitecture()->m_object;
+	defSite.address = defAddr;
+
+	auto var_data = BNVariable();
+	var_data.type = var.type;
+	var_data.index = var.index;
+	var_data.storage = var.storage;
+
+	BNClearInformedVariableValue(m_object, &var_data, &defSite);
+}
+
+
+void Function::ClearInformedVariableValues() 
+{
+	BNClearInformedVariableValues(m_object);
+}
+
+
 void Function::RequestDebugReport(const string& name)
 {
 	BNRequestFunctionDebugReport(m_object, name.c_str());
diff --git a/python/function.py b/python/function.py
index fd8f7674..42bc06f0 100644
--- a/python/function.py
+++ b/python/function.py
@@ -346,6 +346,7 @@ class PossibleValueSet(object):
 			self._values = set()
 			for i in range(0, value.count):
 				self._values.add(value.valueSet[i])
+		self._count = value.count
 
 	def __repr__(self):
 		if self._type == RegisterValueType.EntryValue:
@@ -383,6 +384,63 @@ class PossibleValueSet(object):
 			return NotImplemented
 		return not (self == other)
 
+	def _to_api_object(self):
+		result = core.BNPossibleValueSet()
+		result.state = RegisterValueType(self.type)
+		if self.type == RegisterValueType.UndeterminedValue:
+			return result
+		elif self.type == RegisterValueType.ConstantValue:
+			result.value = self.value
+		elif self.type == RegisterValueType.ConstantPointerValue:
+			result.value = self.value
+		elif self.type == RegisterValueType.StackFrameOffset:
+			result.offset = self.value
+		elif self.type == RegisterValueType.SignedRangeValue:
+			result.offst = self.value
+			result.ranges = (core.BNValueRange * self.count)()
+			for i in range(0, self.count):
+				start = self.ranges[i].start
+				end = self.ranges[i].end
+				if start & (1 << 63):
+					start |= ~((1 << 63) - 1)
+				if end & (1 << 63):
+					end |= ~((1 << 63) - 1)
+				value_range = core.BNValueRange()
+				value_range.start = start
+				value_range.end = end
+				value_range.step = self.ranges[i].step
+				result.ranges[i] = value_range
+			result.count = self.count
+		elif self.type == RegisterValueType.UnsignedRangeValue:
+			result.offset = self.value
+			result.ranges = (core.BNValueRange * self.count)()
+			for i in range(0, self.count):
+				value_range = core.BNValueRange()
+				value_range.start = self.ranges[i].start
+				value_range.end = self.ranges[i].end
+				value_range.step = self.ranges[i].step
+				result.ranges[i] = value_range
+			result.count = self.count
+		elif self.type == RegisterValueType.LookupTableValue:
+			result.table = []
+			result.mapping = {}
+			for i in range(self.count):
+				from_list = []
+				for j in range(0, self.table[i].fromCount):
+					from_list.append(self.table[i].fromValues[j])
+					result.mapping[self.table[i].fromValues[j]] = result.table[i].toValue
+				result.table.append(LookupTableEntry(from_list, result.table[i].toValue))
+			result.count = self.count
+		elif (self.type == RegisterValueType.InSetOfValues) or (self.type == RegisterValueType.NotInSetOfValues):
+			values = (ctypes.c_long * self.count)()
+			i = 0
+			for value in self.values:
+				values[i] = value
+				i += 1
+			result.valueSet = ctypes.cast(values, ctypes.POINTER(ctypes.c_long))
+			result.count = self.count
+		return result
+
 	@property
 	def type(self):
 		""" """
@@ -463,6 +521,83 @@ class PossibleValueSet(object):
 		""" """
 		self._values = value
 
+	@property
+	def count(self):
+		""" """
+		return self._count
+
+	@count.setter
+	def count(self, value):
+		self._count = value
+
+	@classmethod
+	def undetermined(self):
+		return PossibleValueSet()
+
+	@classmethod
+	def constant(self, value):
+		""" """
+		result = PossibleValueSet()
+		result.type = RegisterValueType.ConstantValue
+		result.value = value
+		return result
+
+	@classmethod
+	def constant_ptr(self, value):
+		result = PossibleValueSet()
+		result.type = RegisterValueType.ConstantPointerValue
+		result.value = value
+		return result
+
+	@classmethod
+	def stack_frame_offset(self, offset):
+		result = PossibleValueSet()
+		result.type = RegisterValueType.StackFrameOffset
+		result.value = value
+		return result
+
+	@classmethod
+	def signed_range_value(self, ranges):
+		result = PossibleValueSet()
+		result.value = 0
+		result.type = RegisterValueType.SignedRangeValue
+		result.ranges = ranges
+		result.count = len(ranges)
+		return result
+
+	@classmethod
+	def unsigned_range_value(self, ranges):
+		result = PossibleValueSet()
+		result.value = 0
+		result.type = RegisterValueType.UnsignedRangeValue
+		result.ranges = ranges
+		result.count = len(ranges)
+		return result
+
+	@classmethod
+	def in_set_of_values(self, values):
+		result = PossibleValueSet()
+		result.type = RegisterValueType.InSetOfValues
+		result.values = set(values)
+		result.count = len(values)
+		return result
+
+	@classmethod 
+	def not_in_set_of_values(self, values):
+		result = PossibleValueSet()
+		result.type = RegisterValueType.NotInSetOfValues
+		result.values = set(values)
+		result.count = len(values)
+		return result
+
+	@classmethod
+	def lookup_table_value(self, lookup_table, mapping):
+		result = PossibleValueSet()
+		result.type = RegisterValueType.LookupTableValue
+		result.table = lookup_table
+		result.mapping = mapping
+		return result
+
 
 class StackVariableReference(object):
 	def __init__(self, src_operand, t, name, var, ref_ofs, size):
@@ -2422,6 +2557,51 @@ class Function(object):
 			arch = self.arch
 		return core.BNIsCallInstruction(self.handle, arch.handle, addr)
 
+	def set_var_value(self, var, def_addr, value):
+		var_defs = self.mlil.get_var_definitions(var)
+		if var_defs is None:
+			raise ValueError("Could not get definition for Variable")
+		found = False
+		for site in var_defs:
+			if site.address == def_addr:
+				found = True
+				break
+		if not found:
+			raise ValueError("No definition for Variable found at given address")
+		def_site = core.BNArchitectureAndAddress()
+		def_site.arch = self.arch.handle
+		def_site.address = def_addr
+
+		var_data = core.BNVariable()
+		var_data.type = var.source_type
+		var_data.index = var.index
+		var_data.storage = var.storage
+		core.BNSetVariableValue(self.handle, var_data, def_site, value._to_api_object())
+
+	def clear_informed_var_value(self, var, def_addr):
+		var_defs = self.mlil.get_var_definitions(var)
+		if var_defs is None:
+			raise ValueError("Could not get definition for Variable")
+		found = False
+		for site in var_defs:
+			if site.address == def_addr:
+				found = True
+				break
+		if not found:
+			raise ValueError("No definition for Variable found at given address")
+		def_site = core.BNArchitectureAndAddress()
+		def_site.arch = self.arch.handle
+		def_site.address = def_addr
+
+		var_data = core.BNVariable()
+		var_data.type = var.source_type
+		var_data.index = var.index 
+		var_data.storage = var.storage
+		core.BNClearInformedVariableValue(self.handle, var_data, def_site)
+
+	def clear_informed_var_values(self):
+		core.BNClearInformedVariableValues(self.handle)
+
 	def request_debug_report(self, name):
 		core.BNRequestFunctionDebugReport(self.handle, name)
 		self.view.update_analysis()