Various python types api bug fixes and unit tests

2 files changed, 145 insertions, 64 deletions

diff --git a/python/types.py b/python/types.py
index cb215070..b058cc2c 100644
--- a/python/types.py
+++ b/python/types.py
@@ -62,6 +62,8 @@ class QualifiedName:
 		self._name:List[str] = []
 		if isinstance(name, str):
 			self._name = [name]
+		elif isinstance(name, bytes):
+			self._name = [name.decode("utf-8")]
 		elif isinstance(name, self.__class__):
 			self._name = name._name
 		elif isinstance(name, (list, tuple)):
@@ -649,7 +651,7 @@ class TypeBuilder:
 		return ArrayBuilder.create(type, count)
 
 	@staticmethod
-	def function(ret:Optional['Type'], params:Optional[ParamsType]=None, calling_convention:'callingconvention.CallingConvention'=None,
+	def function(ret:Optional['Type']=None, params:Optional[ParamsType]=None, calling_convention:'callingconvention.CallingConvention'=None,
 		variable_arguments:BoolWithConfidenceType=BoolWithConfidence(False),
 		stack_adjust:SizeWithConfidenceType=0) -> 'FunctionBuilder':
 		"""
@@ -660,13 +662,7 @@ class TypeBuilder:
 		:param CallingConvention calling_convention: optional argument for the function calling convention
 		:param bool variable_arguments: optional boolean, true if the function has a variable number of arguments
 		"""
-		if ret is None:
-			ret = Type.void()
-
-		_params = params
-		if params is None:
-			_params = []
-		return FunctionBuilder.create(ret, calling_convention, _params, variable_arguments, stack_adjust)
+		return FunctionBuilder.create(ret, calling_convention, params, variable_arguments, stack_adjust)
 
 	@staticmethod
 	def structure(members:Optional[MembersType]=None, packed:_bool=False, type:StructureVariant=StructureVariant.StructStructureType) -> 'StructureBuilder':
@@ -872,12 +868,16 @@ class ArrayBuilder(TypeBuilder):
 
 class FunctionBuilder(TypeBuilder):
 	@classmethod
-	def create(cls, return_type:SomeType, calling_convention:Optional['callingconvention.CallingConvention']=None,
+	def create(cls, return_type:Optional[SomeType]=None, calling_convention:Optional['callingconvention.CallingConvention']=None,
 		params:Optional[ParamsType]=None, var_args:BoolWithConfidenceType=False,
 		stack_adjust:OffsetWithConfidenceType=0, platform:'_platform.Platform'=None,
 		confidence:int=core.max_confidence) -> 'FunctionBuilder':
 		param_buf = FunctionBuilder._to_core_struct(params)
-		ret_conf = return_type._to_core_struct()
+		if return_type is None:
+			ret_conf = Type.void()._to_core_struct()
+		else:
+			ret_conf = return_type._to_core_struct()
+
 		conv_conf = core.BNCallingConventionWithConfidence()
 		if calling_convention is None:
 			conv_conf.convention = None
@@ -958,7 +958,9 @@ class FunctionBuilder(TypeBuilder):
 		return BoolWithConfidence.from_core_struct(core.BNTypeBuilderHasVariableArguments(self._handle))
 
 	@staticmethod
-	def _to_core_struct(params:ParamsType):
+	def _to_core_struct(params:Optional[ParamsType]=None):
+		if params is None:
+			params = []
 		param_buf = (core.BNFunctionParameter * len(params))()
 		for i, param in enumerate(params):
 			core_param = param_buf[i]
@@ -1022,16 +1024,8 @@ class StructureBuilder(TypeBuilder):
 		assert builder_handle is not None, "Can't instantiate Structure with builder_handle set to None"
 		self.builder_handle = builder_handle
 
-	@classmethod
-	def create(cls, members:MembersType=None,
-		type:StructureVariant=StructureVariant.StructStructureType,
-		packed:bool=False,
-		width:Optional[int]=None, platform:'_platform.Platform'=None,
-		confidence:int=core.max_confidence) -> 'StructureBuilder':
-		structure_builder_handle = core.BNCreateStructureBuilderWithOptions(type, packed)
-		assert structure_builder_handle is not None, "core.BNCreateStructureBuilderWithOptions returned None"
-		if width is not None:
-			core.BNSetStructureBuilderWidth(structure_builder_handle, width)
+	@staticmethod
+	def _add_members_to_builder(structure_builder_handle, members) -> None:
 		if members is None:
 			members = []
 		for member in members:
@@ -1045,6 +1039,21 @@ class StructureBuilder(TypeBuilder):
 				core.BNAddStructureBuilderMember(structure_builder_handle, member._to_core_struct(), "", MemberAccess.NoAccess, MemberScope.NoScope)
 			else:
 				assert False, f"Structure member type {member} not supported"
+
+	def _add_members(self, members):
+		StructureBuilder._add_members_to_builder(self.builder_handle, members)
+
+	@classmethod
+	def create(cls, members:MembersType=None,
+		type:StructureVariant=StructureVariant.StructStructureType,
+		packed:bool=False,
+		width:Optional[int]=None, platform:'_platform.Platform'=None,
+		confidence:int=core.max_confidence) -> 'StructureBuilder':
+		structure_builder_handle = core.BNCreateStructureBuilderWithOptions(type, packed)
+		assert structure_builder_handle is not None, "core.BNCreateStructureBuilderWithOptions returned None"
+		if width is not None:
+			core.BNSetStructureBuilderWidth(structure_builder_handle, width)
+		StructureBuilder._add_members_to_builder(structure_builder_handle, members)
 		type_builder_handle = core.BNCreateStructureTypeBuilderWithBuilder(structure_builder_handle)
 		assert type_builder_handle is not None, "core.BNCreateStructureTypeBuilderWithBuilder returned None"
 		return cls(type_builder_handle, structure_builder_handle, platform, confidence)
@@ -1074,13 +1083,12 @@ class StructureBuilder(TypeBuilder):
 			core.BNFreeStructureMemberList(members, count.value)
 
 	@members.setter
-	def members(self, members:List[StructureMember]) -> None:
-		for i in range(len(self.members)):
-			core.BNRemoveStructureBuilderMember(self.builder_handle, i)
-
-		for member in members:
-			core.BNAddStructureBuilderMember(self.builder_handle, member.type._to_core_struct(),
-				member.name, ctypes.c_int(member.access), ctypes.c_int(member.scope))
+	def members(self, members:MembersType=None) -> None:
+		count = len(self.members)
+		# remove members in reverse order
+		for i in reversed(range(count)):
+			self.remove(i)
+		self._add_members(members)
 
 	@property
 	def packed(self) -> bool:
@@ -1122,7 +1130,7 @@ class StructureBuilder(TypeBuilder):
 		return f"<struct: size {self.width:#x}>"
 
 	def __getitem__(self, name:str) -> Optional[StructureMember]:
-		member = core.BNGetStructureBuilderMemberByName(self._handle, name)
+		member = core.BNGetStructureBuilderMemberByName(self.builder_handle, name)
 		if member is None:
 			return None
 		try:
@@ -1156,23 +1164,6 @@ class StructureBuilder(TypeBuilder):
 				return i
 		return None
 
-	def index_from(self, index:Optional[MemberIndex]=None, name:Optional[MemberName]=None, offset:Optional[MemberOffset]=None) -> MemberIndex:
-		if index is not None:
-			if index >= len(self.members):
-				raise IndexError("list index out of range")
-		elif name is not None:
-			index = self.index_by_name(name)
-			if index is None:
-				raise ValueError(f"StructureMember {name} doesn't exist")
-		elif offset is not None:
-			index = self.index_by_offset(offset)
-			if index is None:
-				raise ValueError(f"No StructureMember at offset exists")
-		else:
-			raise ValueError("One of the following must")
-
-		return index
-
 	def replace(self, index:int, type:SomeType, name:str="", overwrite_existing:bool=True):
 		core.BNReplaceStructureBuilderMember(self.builder_handle, index,
 			type._to_core_struct(), name, overwrite_existing)
@@ -1412,13 +1403,13 @@ class NamedTypeReferenceBuilder(TypeBuilder):
 			named_type.name, width, align, None, core.max_confidence, const, volatile)
 
 	@staticmethod
-	def named_type_from_type_and_id(type_id:str, name:QualifiedName, type:Optional['Type']) -> 'NamedTypeReferenceBuilder':
+	def named_type_from_type_and_id(type_id:str, name:QualifiedName, type:Optional['Type']=None) -> 'NamedTypeReferenceBuilder':
 		if type is None:
 			return NamedTypeReferenceBuilder.create(NamedTypeReferenceClass.UnknownNamedTypeClass, type_id, name)
 		elif type.type_class == TypeClass.StructureTypeClass:
-			if type.structure_type == StructureVariant.StructStructureType:
+			if type.type == StructureVariant.StructStructureType:
 				return NamedTypeReferenceBuilder.create(NamedTypeReferenceClass.StructNamedTypeClass, type_id, name)
-			elif type.structure_type == StructureVariant.UnionStructureType:
+			elif type.type == StructureVariant.UnionStructureType:
 				return NamedTypeReferenceBuilder.create(NamedTypeReferenceClass.UnionNamedTypeClass, type_id, name)
 			else:
 				return NamedTypeReferenceBuilder.create(NamedTypeReferenceClass.ClassNamedTypeClass, type_id, name)
@@ -1432,7 +1423,7 @@ class NamedTypeReferenceBuilder(TypeBuilder):
 		if type_class is None:
 			return NamedTypeReferenceBuilder.create(NamedTypeReferenceClass.UnknownNamedTypeClass, str(uuid.uuid4()), name)
 		else:
-			return NamedTypeReferenceBuilder.create(NamedTypeReferenceClass.TypedefNamedTypeClass, str(uuid.uuid4()), name)
+			return NamedTypeReferenceBuilder.create(type_class, str(uuid.uuid4()), name)
 
 	@staticmethod
 	def named_type_from_registered_type(view:'binaryview.BinaryView', name:QualifiedName) -> 'NamedTypeReferenceBuilder':
@@ -1766,7 +1757,7 @@ class Type:
 		return ArrayType.create(type, count)
 
 	@staticmethod
-	def function(ret:Optional['Type'], params:Optional[ParamsType]=None, calling_convention:'callingconvention.CallingConvention'=None,
+	def function(ret:Optional['Type']=None, params:Optional[ParamsType]=None, calling_convention:'callingconvention.CallingConvention'=None,
 		variable_arguments:BoolWithConfidenceType=False,
 		stack_adjust:OffsetWithConfidence=OffsetWithConfidence(0)) -> 'FunctionType':
 		"""
diff --git a/suite/api_test.py b/suite/api_test.py
index 696a0709..b2c13daf 100644
--- a/suite/api_test.py
+++ b/suite/api_test.py
@@ -10,7 +10,8 @@ from binaryninja.architecture import Architecture
 from binaryninja.pluginmanager import RepositoryManager
 from binaryninja.platform import Platform
 from binaryninja.function import Function
-from binaryninja.types import (Type, TypeBuilder, FunctionParameter, BoolWithConfidence, EnumerationBuilder, NamedTypeReferenceBuilder,
+from binaryninja.enums import (StructureVariant, NamedTypeReferenceClass)
+from binaryninja.types import (QualifiedName, Type, TypeBuilder, FunctionParameter, BoolWithConfidence, EnumerationBuilder, NamedTypeReferenceBuilder,
     IntegerBuilder, CharBuilder, FloatBuilder, WideCharBuilder, PointerBuilder, ArrayBuilder, FunctionBuilder, StructureBuilder,
 	StructureMember)
 
@@ -425,7 +426,7 @@ class TypeBuilderTest(unittest.TestCase):
 		self.plat = Platform['x86_64']
 		self.cc = self.plat.calling_conventions[0]
 
-	def test_builder_mutability_round_trip(self):
+	def test_IntegerBuilder(self):
 		ib = TypeBuilder.int(4)
 		ib.const = True
 		ib.volatile = False
@@ -438,6 +439,7 @@ class TypeBuilderTest(unittest.TestCase):
 		assert len(ib) == 4
 		assert ib == ib.immutable_copy().mutable_copy(), "IntegerBuilder failed to round trip mutability"
 
+	def test_CharBuilder(self):
 		b = TypeBuilder.char("my_char")
 		b.const = True
 		b.volatile = False
@@ -449,6 +451,7 @@ class TypeBuilderTest(unittest.TestCase):
 		assert b == b.immutable_copy().mutable_copy(), "CharBuilder failed to round trip mutability"
 
 
+	def test_FloatBuilder(self):
 		b = TypeBuilder.float(4, "half")
 		b.const = True
 		b.volatile = False
@@ -457,6 +460,7 @@ class TypeBuilderTest(unittest.TestCase):
 		assert b.alternate_name == "half"
 		assert b == b.immutable_copy().mutable_copy(), "FloatBuilder failed to round trip mutability"
 
+	def test_WideCharBuilder(self):
 		b = TypeBuilder.wide_char(4, "wchar32_t")
 		b.const = True
 		b.volatile = False
@@ -465,6 +469,8 @@ class TypeBuilderTest(unittest.TestCase):
 		assert b.alternate_name == "wchar32_t"
 		assert b == b.immutable_copy().mutable_copy(), "WideCharBuilder failed to round trip mutability"
 
+	def test_PointerBuilder(self):
+		ib = TypeBuilder.int(4)
 		b = TypeBuilder.pointer(self.arch, ib, 4)
 		b.const = True
 		b.volatile = False
@@ -474,14 +480,19 @@ class TypeBuilderTest(unittest.TestCase):
 		assert b == b.immutable_copy().mutable_copy(), "PointerBuilder failed to round trip mutability"
 		pb = b
 
+	def test_VoidBuilder(self):
 		b = TypeBuilder.void()
 		assert b == b.immutable_copy().mutable_copy(), "VoidBuilder failed to round trip mutability"
-		vb = b
 
+	def test_BoolBuilder(self):
 		b = TypeBuilder.bool()
 		assert b == b.immutable_copy().mutable_copy(), "VoidBuilder failed to round trip mutability"
-		bb = b
 
+	def test_FunctionBuilder(self):
+		bb = TypeBuilder.bool()
+		ib = TypeBuilder.int(4)
+		pb = TypeBuilder.pointer(self.arch, ib, 4)
+		vb = TypeBuilder.void()
 		b = TypeBuilder.function(vb, [FunctionParameter(pb, "arg1")], self.cc)
 		assert b.system_call_number is None
 		b.system_call_number = 1
@@ -515,17 +526,96 @@ class TypeBuilderTest(unittest.TestCase):
 		assert b.parameters[0].name == "arg1"
 		assert len(b.parameters) == 1
 
-		b = TypeBuilder.function(vb)
+		b = TypeBuilder.function()
 		assert len(b.parameters) == 0
 		assert b.return_value == TypeBuilder.void()
 		assert b == b.immutable_copy().mutable_copy(), "FunctionBuilder failed to round trip mutability"
 
-		# b = TypeBuilder.structure([(ib, "name")], False)
-		# assert b.alignment == 4
-		# assert b.
-		# b = TypeBuilder.structure([StructureMember()])
-
-		b = TypeBuilder.array(pb, 4)
-		assert len(b) == len(pb) * 4
+	def test_ArrayBuilder(self):
+		ib = TypeBuilder.int(4)
+		b = TypeBuilder.array(ib, 4)
+		assert len(b) == len(ib) * 4
 		assert b.count == 4
-		assert b.element_type == pb.immutable_copy()
+		assert b.element_type == ib.immutable_copy()
+		assert b == b.immutable_copy().mutable_copy(), "ArrayBuilder failed to round trip mutability"
+
+	def test_StructureBuilder(self):
+		ib = TypeBuilder.int(4)
+		b = TypeBuilder.structure([StructureMember(ib, "name", 0), Type.bool()])
+		b.members = [*b.members, StructureMember(ib, "name2", 8)]
+		assert not b.union
+		b.type = StructureVariant.UnionStructureType
+		assert b.union
+		b.type = StructureVariant.StructStructureType
+		assert b['name'].name == "name"
+		assert b['name'].type == ib.immutable_copy()
+
+		it = iter(b)
+		mem = next(it)
+		assert mem.name == "name"
+		assert mem.type == ib.immutable_copy()
+		mem = next(it)
+		assert mem.name == "field_4"
+		assert mem.type == Type.bool()
+		mem = next(it)
+		assert mem.name == "name2"
+		assert mem.type == ib.immutable_copy()
+
+		assert len(b) == 12
+		assert b.member_at_offset(0).name == "name"
+		assert b.member_at_offset(0).type == ib.immutable_copy()
+		assert b.member_at_offset(4).name == "field_4"
+		assert b.member_at_offset(4).type == Type.bool()
+		assert b.member_at_offset(8).name == "name2"
+		assert b.member_at_offset(8).type == ib.immutable_copy()
+		assert b.index_by_name("name") == 0
+		assert b.index_by_name("name2") == 2
+		assert b.index_by_offset(0) == 0
+		assert b.index_by_offset(4) == 1
+		assert b.index_by_offset(8) == 2
+		b.add_member_at_offset("foo", Type.int(4), 0x20)
+		mem = b.member_at_offset(0x20)
+		assert mem.name == "foo"
+		assert mem.type == Type.int(4)
+		assert b == b.immutable_copy().mutable_copy(), "StructureBuilder failed to round trip mutability"
+
+	def test_NamedTypeReferenceBuilder(self):
+		b = TypeBuilder.named_type_from_type("foobar", NamedTypeReferenceClass.UnknownNamedTypeClass)
+		assert b.name == "foobar"
+		assert b.id == b.type_id
+		assert b.named_type_class == NamedTypeReferenceClass.UnknownNamedTypeClass
+		assert b == b.immutable_copy().mutable_copy(), "NamedTypeReferenceBuilder failed to round trip mutability"
+
+		b = TypeBuilder.named_type_from_type_and_id("type_id", QualifiedName(b"name"), Type.int(4))
+		assert b.name == "name"
+		assert b.id == "type_id"
+		assert b.named_type_class == NamedTypeReferenceClass.TypedefNamedTypeClass
+
+		b = TypeBuilder.named_type_from_type_and_id("type_id", QualifiedName(b"name"))
+		assert b.name == "name"
+		assert b.id == "type_id"
+		assert b.named_type_class == NamedTypeReferenceClass.UnknownNamedTypeClass
+
+		enm = TypeBuilder.enumeration(self.arch, [("Member1", 0)], 4, False)
+		b = TypeBuilder.named_type_from_type_and_id("type_id", QualifiedName(b"name"), enm)
+		assert b.name == "name"
+		assert b.id == "type_id"
+		assert b.named_type_class == NamedTypeReferenceClass.EnumNamedTypeClass
+
+		str = TypeBuilder.structure([], True, StructureVariant.StructStructureType)
+		b = TypeBuilder.named_type_from_type_and_id("type_id", QualifiedName(b"name"), str)
+		assert b.name == "name"
+		assert b.id == "type_id"
+		assert b.named_type_class == NamedTypeReferenceClass.StructNamedTypeClass
+
+		str = TypeBuilder.structure([], True, StructureVariant.ClassStructureType)
+		b = TypeBuilder.named_type_from_type_and_id("type_id", QualifiedName(b"name"), str)
+		assert b.name == "name"
+		assert b.id == "type_id"
+		assert b.named_type_class == NamedTypeReferenceClass.ClassNamedTypeClass
+
+		str = TypeBuilder.structure([], True, StructureVariant.UnionStructureType)
+		b = TypeBuilder.named_type_from_type_and_id("type_id", QualifiedName(b"name"), str)
+		assert b.name == "name"
+		assert b.id == "type_id"
+		assert b.named_type_class == NamedTypeReferenceClass.UnionNamedTypeClass