Add the ability to automatically create struct members

5 files changed, 261 insertions, 19 deletions

diff --git a/python/architecture.py b/python/architecture.py
index 19ef7831..e23f39d7 100644
--- a/python/architecture.py
+++ b/python/architecture.py
@@ -2808,3 +2808,107 @@ class ReferenceSource(object):
 	@address.setter
 	def address(self, value):
 		self._address = value
+
+
+class TypeFieldReference(object):
+	def __init__(self, func, arch, addr, size):
+		self._function = func
+		self._arch = arch
+		self._address = addr
+		self._size = size
+
+	def __repr__(self):
+		if self._arch:
+			return "<ref: %s@%#x, size: %#x>" % (self._arch.name, self._address, self._size)
+		else:
+			return "<ref: %#x, size: %#x>" % (self._address, self._size)
+
+	def __eq__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		return (self.function, self.arch, self.address, self._size) ==\
+			(other.address, other.function, other.arch, other.size)
+
+	def __ne__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		return not (self == other)
+
+	def __lt__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		if self.address < other.address:
+			return True
+		elif self.address > other.address:
+			return False
+		else:
+			return self.size < other.size
+
+	def __gt__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		if self.address > other.address:
+			return True
+		elif self.address < other.address:
+			return False
+		else:
+			return self.size > other.size
+
+	def __ge__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		if self.address > other.address:
+			return True
+		elif self.address < other.address:
+			return False
+		else:
+			return self.size >= other.size
+
+	def __le__(self, other):
+		if not isinstance(other, self.__class__):
+			return NotImplemented
+		if self.address < other.address:
+			return True
+		elif self.address > other.address:
+			return False
+		else:
+			return self.size <= other.size
+
+	def __hash__(self):
+		return hash((self._function, self._arch, self._address, self._size))
+
+	@property
+	def function(self):
+		""" """
+		return self._function
+
+	@function.setter
+	def function(self, value):
+		self._function = value
+
+	@property
+	def arch(self):
+		""" """
+		return self._arch
+
+	@arch.setter
+	def arch(self, value):
+		self._arch = value
+
+	@property
+	def address(self):
+		""" """
+		return self._address
+
+	@address.setter
+	def address(self, value):
+		self._address = value
+
+	@property
+	def size(self):
+		""" """
+		return self._size
+
+	@size.setter
+	def address(self, value):
+		self._size = value
\ No newline at end of file
diff --git a/python/binaryview.py b/python/binaryview.py
index 7fc37287..84dde18c 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -3505,12 +3505,12 @@ class BinaryView(object):
 
 	def get_code_refs_for_type_field(self, name, offset):
 		"""
-		``get_code_refs_for_type`` returns a list of ReferenceSource objects (xrefs or cross-references) that reference the provided type field.
+		``get_code_refs_for_type`` returns a list of TypeFieldReference objects (xrefs or cross-references) that reference the provided type field.
 
 		:param QualifiedName name: name of type to query for references
 		:param int offset: offset of the field, relative to the type
 		:return: List of References for the given type
-		:rtype: list(ReferenceSource)
+		:rtype: list(TypeFieldReference)
 		:Example:
 
 			>>> bv.get_code_refs_for_type_field('A', 0x8)
@@ -3533,8 +3533,9 @@ class BinaryView(object):
 			else:
 				arch = None
 			addr = refs[i].addr
-			result.append(binaryninja.architecture.ReferenceSource(func, arch, addr))
-		core.BNFreeCodeReferences(refs, count.value)
+			size = refs[i].size
+			result.append(binaryninja.architecture.TypeFieldReference(func, arch, addr, size))
+		core.BNFreeTypeFieldReferences(refs, count.value)
 		return result
 
 
@@ -3730,9 +3731,9 @@ class BinaryView(object):
 		core.BNRemoveUserDataReference(self.handle, from_addr, to_addr)
 
 
-	def get_all_type_fields_referenced_by_code(self, name):
+	def get_all_fields_referenced(self, name):
 		"""
-		``get_all_type_fields_referenced_by_code`` returns a list of offsets in the QualifiedName
+		``get_all_fields_referenced`` returns a list of offsets in the QualifiedName
 		specified by name, which are referenced by code.
 
 		:param QualifiedName name: name of type to query for references
@@ -3740,14 +3741,14 @@ class BinaryView(object):
 		:rtype: list(integer)
 		:Example:
 
-			>>> bv.get_all_type_fields_referenced_by_code('A')
-			['<type D, offset 0x8, direct>', '<type C, offset 0x10, indirect>']
+			>>> bv.get_all_fields_referenced('A')
+			[0, 8, 16, 24, 32, 40]
 			>>>
 
 		"""
 		count = ctypes.c_ulonglong(0)
 		name = types.QualifiedName(name)._get_core_struct()
-		refs = core.BNGetAllFieldsReferencedByCode(self.handle, name, count)
+		refs = core.BNGetAllFieldsReferenced(self.handle, name, count)
 
 		result = []
 		for i in range(0, count.value):
@@ -3756,6 +3757,136 @@ class BinaryView(object):
 		core.BNFreeDataReferences(refs, count.value)
 		return result
 
+	def get_all_sizes_referenced(self, name):
+		"""
+		``get_all_sizes_referenced`` returns a map from field offset to a list of sizes of
+		the accesses to it.
+
+		:param QualifiedName name: name of type to query for references
+		:return: A map from field offset to the	size of the code accesses to it
+		:rtype: map
+		:Example:
+
+			>>> bv.get_all_sizes_referenced('B')
+			{0: [1, 8], 8: [8], 16: [1, 8]}
+			>>>
+
+		"""
+		count = ctypes.c_ulonglong(0)
+		name = types.QualifiedName(name)._get_core_struct()
+		refs = core.BNGetAllSizesReferenced(self.handle, name, count)
+
+		result = {}
+		for i in range(0, count.value):
+			result[refs[i].offset] = []
+			for j in range(0, refs[i].count):
+				result[refs[i].offset].append(refs[i].sizes[j])
+
+		core.BNFreeTypeFieldReferenceSizeInfo(refs, count.value)
+		return result
+
+	def get_all_types_referenced(self, name):
+		"""
+		``get_all_types_referenced`` returns a map from field offset to a related to the
+		type field access.
+
+		:param QualifiedName name: name of type to query for references
+		:return: A map from field offset to a list of incoming types written to it
+		:rtype: map
+		:Example:
+
+			>>> bv.get_all_types_referenced('B')
+			{0: [<type: char, 0% confidence>], 8: [<type: int64_t, 0% confidence>],
+			16: [<type: char, 0% confidence>, <type: bool>]}
+			>>>
+
+		"""
+		count = ctypes.c_ulonglong(0)
+		name = types.QualifiedName(name)._get_core_struct()
+		refs = core.BNGetAllTypesReferenced(self.handle, name, count)
+
+		result = {}
+		for i in range(0, count.value):
+			result[refs[i].offset] = []
+			for j in range(0, refs[i].count):
+				typeObj = types.Type(core.BNNewTypeReference(refs[i].types[j].type),\
+					confidence = refs[i].types[j].confidence)
+				result[refs[i].offset].append(typeObj)
+
+		core.BNFreeTypeFieldReferenceTypeInfo(refs, count.value)
+		return result
+
+	def get_sizes_referenced(self, name, offset):
+		"""
+		``get_sizes_referenced`` returns a list of sizes of the accesses to it.
+
+		:param QualifiedName name: name of type to query for references
+		:param int offset: offset of the field
+		:return: a list of sizes of the accesses to it.
+		:rtype: list
+		:Example:
+
+			>>> bv.get_sizes_referenced('B', 16)
+			[1, 8]
+			>>>
+
+		"""
+		count = ctypes.c_ulonglong(0)
+		name = types.QualifiedName(name)._get_core_struct()
+		refs = core.BNGetSizesReferenced(self.handle, name, offset, count)
+
+		result = []
+		for i in range(0, count.value):
+			result.append(refs[i])
+
+		core.BNFreeTypeFieldReferenceSizes(refs, count.value)
+		return result
+
+	def get_types_referenced(self, name, offset):
+		"""
+		``get_types_referenced`` returns a list of types related to the type field access.
+
+		:param QualifiedName name: name of type to query for references
+		:param int offset: offset of the field
+		:return: a list of types related to the type field access.
+		:rtype: list
+		:Example:
+
+			>>> bv.get_types_referenced('B', 0x10)
+			[<type: bool>, <type: char, 0% confidence>]
+			>>>
+		"""
+		count = ctypes.c_ulonglong(0)
+		name = types.QualifiedName(name)._get_core_struct()
+		refs = core.BNGetTypesReferenced(self.handle, name, offset, count)
+
+		result = []
+		for i in range(0, count.value):
+			typeObj = types.Type(core.BNNewTypeReference(refs[i].type),\
+				confidence = refs[i].confidence)
+			result.append(typeObj)
+
+		core.BNFreeTypeFieldReferenceTypes(refs, count.value)
+		return result
+
+	def create_structure_from_offset_access(self, name):
+		newMemberAdded = ctypes.c_bool(False)
+		name = types.QualifiedName(name)._get_core_struct()
+		struct = core.BNCreateStructureFromOffsetAccess(self.handle, name, newMemberAdded)
+		if struct is None:
+			return None
+		return types.Structure(struct)
+
+	@classmethod
+	def create_structure_member_from_access(self, name, offset):
+		name = types.QualifiedName(name)._get_core_struct()
+		result = core.BNCreateStructureMemberFromAccess(self.handle, name, offset)
+		if not result.type:
+			return None
+
+		return types.Type(core.BNNewTypeReference(result.type),\
+			confidence = result.confidence)
+
 	def get_callers(self, addr):
 		"""
 		``get_callers`` returns a list of ReferenceSource objects (xrefs or cross-references) that call the provided virtual address.
diff --git a/python/function.py b/python/function.py
index cb8a869c..0ec6142d 100644
--- a/python/function.py
+++ b/python/function.py
@@ -2364,7 +2364,7 @@ class Function(object):
 		name = types.QualifiedName(name)._get_core_struct()
 		core.BNRemoveUserTypeReference(self.handle, from_arch.handle, from_addr, name)
 
-	def add_user_type_field_ref(self, from_addr, name, offset, from_arch=None):
+	def add_user_type_field_ref(self, from_addr, name, offset, from_arch = None, size = 0):
 		"""
 		``add_user_type_field_ref`` places a user-defined type field cross-reference from the
 		instruction at the given address and architecture to the specified type. If the specified
@@ -2375,10 +2375,11 @@ class Function(object):
 		:param QualifiedName name: name of the referenced type
 		:param int offset: offset of the field, relative to the type
 		:param Architecture from_arch: (optional) architecture of the source instruction
+		:param int size: (optional) the size of the access
 		:rtype: None
 		:Example:
 
-			>>> current_function.add_user_type_field_ref(here, 'A'. 0x8)
+			>>> current_function.add_user_type_field_ref(here, 'A', 0x8)
 
 		"""
 
@@ -2386,9 +2387,10 @@ class Function(object):
 			from_arch = self.arch
 
 		name = types.QualifiedName(name)._get_core_struct()
-		core.BNAddUserTypeFieldReference(self.handle, from_arch.handle, from_addr, name, offset)
+		core.BNAddUserTypeFieldReference(self.handle, from_arch.handle, from_addr, name,\
+			offset, size)
 
-	def remove_user_type_field_ref(self, from_addr, name, offset, from_arch=None):
+	def remove_user_type_field_ref(self, from_addr, name, offset, from_arch = None, size = 0):
 		"""
 		``remove_user_type_field_ref`` removes a user-defined type field cross-reference.
 		If the given address is not contained within this function, or if there is no
@@ -2398,6 +2400,7 @@ class Function(object):
 		:param QualifiedName name: name of the referenced type
 		:param int offset: offset of the field, relative to the type
 		:param Architecture from_arch: (optional) architecture of the source instruction
+		:param int size: (optional) the size of the access
 		:rtype: None
 		:Example:
 
@@ -2409,7 +2412,8 @@ class Function(object):
 			from_arch = self.arch
 
 		name = types.QualifiedName(name)._get_core_struct()
-		core.BNRemoveUserTypeFieldReference(self.handle, from_arch.handle, from_addr, name, offset)
+		core.BNRemoveUserTypeFieldReference(self.handle, from_arch.handle, from_addr, name,\
+			offset, size)
 
 	def get_low_level_il_at(self, addr, arch=None):
 		"""
diff --git a/python/generator.cpp b/python/generator.cpp
index f78324cf..4901b07a 100644
--- a/python/generator.cpp
+++ b/python/generator.cpp
@@ -344,7 +344,8 @@ int main(int argc, char* argv[])
 		{
 			if ((arg.type->GetClass() == PointerTypeClass) &&
 				(arg.type->GetChildType()->GetWidth() == 1) &&
-				(arg.type->GetChildType()->IsSigned()))
+				(arg.type->GetChildType()->IsSigned()) &&
+				(!arg.type->GetChildType()->IsBool()))
 				{
 					stringArgument = true;
 					break;
diff --git a/python/types.py b/python/types.py
index 39378963..af62e5b4 100644
--- a/python/types.py
+++ b/python/types.py
@@ -19,6 +19,8 @@
 # IN THE SOFTWARE.
 
 from __future__ import absolute_import
+
+from binaryninja.log import log_warn
 max_confidence = 255
 
 import ctypes
@@ -1452,26 +1454,26 @@ class Structure(object):
 		tc.confidence = t.confidence
 		core.BNAddStructureBuilderMember(self._handle, tc, name)
 
-	def insert(self, offset, t, name = ""):
+	def insert(self, offset, t, name = "", overwriteExisting = True):
 		if not self._mutable:
 			raise AttributeError("Finalized Structure object is immutable, use mutable_copy()")
 		tc = core.BNTypeWithConfidence()
 		tc.type = t.handle
 		tc.confidence = t.confidence
-		core.BNAddStructureBuilderMemberAtOffset(self._handle, tc, name, offset)
+		core.BNAddStructureBuilderMemberAtOffset(self._handle, tc, name, offset, overwriteExisting)
 
 	def remove(self, i):
 		if not self._mutable:
 			raise AttributeError("Finalized Structure object is immutable, use mutable_copy()")
 		core.BNRemoveStructureBuilderMember(self._handle, i)
 
-	def replace(self, i, t, name = ""):
+	def replace(self, i, t, name = "", overwriteExisting = True):
 		if not self._mutable:
 			raise AttributeError("Finalized Structure object is immutable, use mutable_copy()")
 		tc = core.BNTypeWithConfidence()
 		tc.type = t.handle
 		tc.confidence = t.confidence
-		core.BNReplaceStructureBuilderMember(self._handle, i, tc, name)
+		core.BNReplaceStructureBuilderMember(self._handle, i, tc, name, overwriteExisting)
 
 	def mutable_copy(self):
 		if self._mutable: