Adding APIs for data variables

4 files changed, 142 insertions, 1 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 18a379c3..d1d8eeb5 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -665,6 +665,13 @@ namespace BinaryNinja
 		void Cancel();
 	};
 
+	struct DataVariable
+	{
+		uint64_t address;
+		Ref<Type> type;
+		bool autoDiscovered;
+	};
+
 	/*! BinaryView is the base class for creating views on binary data (e.g. ELF, PE, Mach-O).
 	    BinaryView should be subclassed to create a new BinaryView
 	*/
@@ -803,6 +810,14 @@ namespace BinaryNinja
 		void UpdateAnalysis();
 		void AbortAnalysis();
 
+		void DefineDataVariable(uint64_t addr, Type* type);
+		void DefineUserDataVariable(uint64_t addr, Type* type);
+		void UndefineDataVariable(uint64_t addr);
+		void UndefineUserDataVariable(uint64_t addr);
+
+		std::map<uint64_t, DataVariable> GetDataVariables();
+		bool GetDataVariableAtAddress(uint64_t addr, DataVariable& var);
+
 		std::vector<Ref<Function>> GetAnalysisFunctionList();
 		bool HasFunctions() const;
 		Ref<Function> GetAnalysisFunction(Platform* platform, uint64_t addr);
@@ -832,7 +847,6 @@ namespace BinaryNinja
 
 		void DefineImportedFunction(Symbol* importAddressSym, Function* func);
 
-
 		bool IsNeverBranchPatchAvailable(Architecture* arch, uint64_t addr);
 		bool IsAlwaysBranchPatchAvailable(Architecture* arch, uint64_t addr);
 		bool IsInvertBranchPatchAvailable(Architecture* arch, uint64_t addr);
@@ -854,9 +868,11 @@ namespace BinaryNinja
 		uint64_t GetNextFunctionStartAfterAddress(uint64_t addr);
 		uint64_t GetNextBasicBlockStartAfterAddress(uint64_t addr);
 		uint64_t GetNextDataAfterAddress(uint64_t addr);
+		uint64_t GetNextDataVariableAfterAddress(uint64_t addr);
 		uint64_t GetPreviousFunctionStartBeforeAddress(uint64_t addr);
 		uint64_t GetPreviousBasicBlockEndBeforeAddress(uint64_t addr);
 		uint64_t GetPreviousDataBeforeAddress(uint64_t addr);
+		uint64_t GetPreviousDataVariableBeforeAddress(uint64_t addr);
 
 		LinearDisassemblyPosition GetLinearDisassemblyPositionForAddress(uint64_t addr, DisassemblySettings* settings);
 		std::vector<LinearDisassemblyLine> GetPreviousLinearDisassemblyLines(LinearDisassemblyPosition& pos,
diff --git a/binaryninjacore.h b/binaryninjacore.h
index ac9e9da9..9af53442 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -930,6 +930,13 @@ extern "C"
 		size_t count, total;
 	};
 
+	struct BNDataVariable
+	{
+		uint64_t address;
+		BNType* type;
+		bool autoDiscovered;
+	};
+
 	BINARYNINJACOREAPI char* BNAllocString(const char* contents);
 	BINARYNINJACOREAPI void BNFreeString(char* str);
 
@@ -1417,9 +1424,11 @@ extern "C"
 	BINARYNINJACOREAPI uint64_t BNGetNextFunctionStartAfterAddress(BNBinaryView* view, uint64_t addr);
 	BINARYNINJACOREAPI uint64_t BNGetNextBasicBlockStartAfterAddress(BNBinaryView* view, uint64_t addr);
 	BINARYNINJACOREAPI uint64_t BNGetNextDataAfterAddress(BNBinaryView* view, uint64_t addr);
+	BINARYNINJACOREAPI uint64_t BNGetNextDataVariableAfterAddress(BNBinaryView* view, uint64_t addr);
 	BINARYNINJACOREAPI uint64_t BNGetPreviousFunctionStartBeforeAddress(BNBinaryView* view, uint64_t addr);
 	BINARYNINJACOREAPI uint64_t BNGetPreviousBasicBlockEndBeforeAddress(BNBinaryView* view, uint64_t addr);
 	BINARYNINJACOREAPI uint64_t BNGetPreviousDataBeforeAddress(BNBinaryView* view, uint64_t addr);
+	BINARYNINJACOREAPI uint64_t BNGetPreviousDataVariableBeforeAddress(BNBinaryView* view, uint64_t addr);
 
 	BINARYNINJACOREAPI BNLinearDisassemblyPosition BNGetLinearDisassemblyPositionForAddress(BNBinaryView* view,
 		uint64_t addr, BNDisassemblySettings* settings);
@@ -1430,6 +1439,14 @@ extern "C"
 		BNLinearDisassemblyPosition* pos, BNDisassemblySettings* settings, size_t* count);
 	BINARYNINJACOREAPI void BNFreeLinearDisassemblyLines(BNLinearDisassemblyLine* lines, size_t count);
 
+	BINARYNINJACOREAPI void BNDefineDataVariable(BNBinaryView* view, uint64_t addr, BNType* type);
+	BINARYNINJACOREAPI void BNDefineUserDataVariable(BNBinaryView* view, uint64_t addr, BNType* type);
+	BINARYNINJACOREAPI void BNUndefineDataVariable(BNBinaryView* view, uint64_t addr);
+	BINARYNINJACOREAPI void BNUndefineUserDataVariable(BNBinaryView* view, uint64_t addr);
+	BINARYNINJACOREAPI BNDataVariable* BNGetDataVariables(BNBinaryView* view, size_t* count);
+	BINARYNINJACOREAPI void BNFreeDataVariables(BNDataVariable* vars, size_t count);
+	BINARYNINJACOREAPI bool BNGetDataVariableAtAddress(BNBinaryView* view, uint64_t addr, BNDataVariable* var);
+
 	// Disassembly settings
 	BINARYNINJACOREAPI BNDisassemblySettings* BNCreateDisassemblySettings(void);
 	BINARYNINJACOREAPI BNDisassemblySettings* BNNewDisassemblySettingsReference(BNDisassemblySettings* settings);
diff --git a/binaryview.cpp b/binaryview.cpp
index 864b4687..6d96c8ea 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -764,6 +764,67 @@ void BinaryView::AbortAnalysis()
 }
 
 
+void BinaryView::DefineDataVariable(uint64_t addr, Type* type)
+{
+	BNDefineDataVariable(m_object, addr, type->GetObject());
+}
+
+
+void BinaryView::DefineUserDataVariable(uint64_t addr, Type* type)
+{
+	BNDefineUserDataVariable(m_object, addr, type->GetObject());
+}
+
+
+void BinaryView::UndefineDataVariable(uint64_t addr)
+{
+	BNUndefineDataVariable(m_object, addr);
+}
+
+
+void BinaryView::UndefineUserDataVariable(uint64_t addr)
+{
+	BNUndefineUserDataVariable(m_object, addr);
+}
+
+
+map<uint64_t, DataVariable> BinaryView::GetDataVariables()
+{
+	size_t count;
+	BNDataVariable* vars = BNGetDataVariables(m_object, &count);
+
+	map<uint64_t, DataVariable> result;
+	for (size_t i = 0; i < count; i++)
+	{
+		DataVariable var;
+		var.address = vars[i].address;
+		var.type = new Type(BNNewTypeReference(vars[i].type));
+		var.autoDiscovered = vars[i].autoDiscovered;
+		result[var.address] = var;
+	}
+
+	BNFreeDataVariables(vars, count);
+	return result;
+}
+
+
+bool BinaryView::GetDataVariableAtAddress(uint64_t addr, DataVariable& var)
+{
+	var.address = 0;
+	var.type = nullptr;
+	var.autoDiscovered = false;
+
+	BNDataVariable result;
+	if (!BNGetDataVariableAtAddress(m_object, addr, &result))
+		return false;
+
+	var.address = result.address;
+	var.type = new Type(result.type);
+	var.autoDiscovered = result.autoDiscovered;
+	return true;
+}
+
+
 vector<Ref<Function>> BinaryView::GetAnalysisFunctionList()
 {
 	size_t count;
diff --git a/python/__init__.py b/python/__init__.py
index 390c4efd..0d9c15a2 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -696,6 +696,15 @@ class LinearDisassemblyLine(object):
 	def __repr__(self):
 		return repr(self.contents)
 
+class DataVariable(object):
+	def __init__(self, addr, var_type, auto_discovered):
+		self.address = addr
+		self.type = var_type
+		self.auto_discovered = auto_discovered
+
+	def __repr__(self):
+		return "<var 0x%x: %s>" % (self.address, str(self.type))
+
 class BinaryView(object):
 	name = None
 	"""Binary View name"""
@@ -994,6 +1003,20 @@ class BinaryView(object):
 		"""Iterator for all lines in the linear disassembly of the view"""
 		return self.get_linear_disassembly(None)
 
+	@property
+	def data_vars(self):
+		"""List of data variables (read-only)"""
+		count = ctypes.c_ulonglong(0)
+		var_list = core.BNGetDataVariables(self.handle, count)
+		result = {}
+		for i in xrange(0, count.value):
+			addr = var_list[i].address
+			var_type = Type(core.BNNewTypeReference(var_list[i].type))
+			auto_discovered = var_list[i].autoDiscovered
+			result[addr] = DataVariable(addr, var_type, auto_discovered)
+		core.BNFreeDataVariables(var_list, count.value)
+		return result
+
 	def __len__(self):
 		return int(core.BNGetViewLength(self.handle))
 
@@ -1367,6 +1390,24 @@ class BinaryView(object):
 	def abort_analysis(self):
 		core.BNAbortAnalysis(self.handle)
 
+	def define_data_var(self, addr, var_type):
+		core.BNDefineDataVariable(self.handle, addr, var_type.handle)
+
+	def define_user_data_var(self, addr, var_type):
+		core.BNDefineUserDataVariable(self.handle, addr, var_type.handle)
+
+	def undefine_data_var(self, addr):
+		core.BNUndefineDataVariable(self.handle, addr)
+
+	def undefine_user_data_var(self, addr):
+		core.BNUndefineUserDataVariable(self.handle, addr)
+
+	def get_data_var_at(self, addr):
+		var = core.BNDataVariable()
+		if not core.BNGetDataVariableAtAddress(self.handle, addr, var):
+			return None
+		return DataVariable(var.address, Type(var.type), var.autoDiscovered)
+
 	def get_function_at(self, platform, addr):
 		func = core.BNGetAnalysisFunction(self.handle, platform.handle, addr)
 		if func is None:
@@ -1552,6 +1593,9 @@ class BinaryView(object):
 	def get_next_data_after(self, addr):
 		return core.BNGetNextDataAfterAddress(self.handle, addr)
 
+	def get_next_data_var_after(self, addr):
+		return core.BNGetNextDataVariableAfterAddress(self.handle, addr)
+
 	def get_previous_function_start_before(self, addr):
 		return core.BNGetPreviousFunctionStartBeforeAddress(self.handle, addr)
 
@@ -1561,6 +1605,9 @@ class BinaryView(object):
 	def get_previous_data_before(self, addr):
 		return core.BNGetPreviousDataBeforeAddress(self.handle, addr)
 
+	def get_previous_data_var_before(self, addr):
+		return core.BNGetPreviousDataVariableBeforeAddress(self.handle, addr)
+
 	def get_linear_disassembly_position_at(self, addr, settings):
 		if settings is not None:
 			settings = settings.handle