Add support for deriving structures from other structures

3 files changed, 201 insertions, 4 deletions

diff --git a/python/binaryview.py b/python/binaryview.py
index 5c198986..3d2acbdb 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -4488,6 +4488,35 @@ class BinaryView:
 		finally:
 			core.BNFreeDataReferences(refs)
 
+	def get_data_refs_from_for_type_field(self, name: '_types.QualifiedNameType', offset: int) -> List[int]:
+		"""
+		``get_data_refs_from_for_type_field`` returns a list of virtual addresses of data which are referenced by the type ``name``.
+
+		Only data referenced by structures with the ``__data_var_refs`` attribute are included.
+
+		:param QualifiedName name: name of type to query for references
+		:param int offset: offset of the field, relative to the type
+		:return: list of integers
+		:rtype: list(integer)
+		:Example:
+
+			>>> bv.get_data_refs_from_for_type_field('A', 0x8)
+			[4203812]
+			>>>
+		"""
+		count = ctypes.c_ulonglong(0)
+		_name = _types.QualifiedName(name)._to_core_struct()
+		refs = core.BNGetDataReferencesFromForTypeField(self.handle, _name, offset, count)
+		assert refs is not None, "core.BNGetDataReferencesFromForTypeField returned None"
+
+		result = []
+		try:
+			for i in range(0, count.value):
+				result.append(refs[i])
+			return result
+		finally:
+			core.BNFreeDataReferences(refs)
+
 	def get_type_refs_for_type(self, name: '_types.QualifiedNameType') -> List['_types.TypeReferenceSource']:
 		"""
 		``get_type_refs_for_type`` returns a list of TypeReferenceSource objects (xrefs or cross-references) that reference the provided QualifiedName.
diff --git a/python/typeprinter.py b/python/typeprinter.py
index 3ad1e9f2..2177ba40 100644
--- a/python/typeprinter.py
+++ b/python/typeprinter.py
@@ -504,8 +504,15 @@ class CoreTypePrinter(TypePrinter):
 			type_ = types.Type.create(handle=core.BNNewTypeReference(core_lines[i].type), platform=data.platform)
 			root_type = types.Type.create(handle=core.BNNewTypeReference(core_lines[i].rootType), platform=data.platform)
 			root_type_name = core.pyNativeStr(core_lines[i].rootTypeName)
-			line = types.TypeDefinitionLine(core_lines[i].lineType, tokens, type_, root_type, root_type_name,
-									  core_lines[i].offset, core_lines[i].fieldIndex)
+			if core_lines[i].baseType:
+				const_conf = types.BoolWithConfidence.get_core_struct(False, 0)
+				volatile_conf = types.BoolWithConfidence.get_core_struct(False, 0)
+				handle = core.BNCreateNamedTypeReference(core_lines[i].baseType, 0, 1, const_conf, volatile_conf)
+				base_type = types.NamedTypeReferenceType(handle, data.platform)
+			else:
+				base_type = None
+			line = types.TypeDefinitionLine(core_lines[i].lineType, tokens, type_, root_type, root_type_name, base_type,
+									  core_lines[i].baseOffset, core_lines[i].offset, core_lines[i].fieldIndex)
 			lines.append(line)
 		core.BNFreeTypeDefinitionLineList(core_lines, count.value)
 		return lines
diff --git a/python/types.py b/python/types.py
index 5c480127..ee47c841 100644
--- a/python/types.py
+++ b/python/types.py
@@ -220,6 +220,8 @@ class TypeDefinitionLine:
 	type: 'Type'
 	root_type: 'Type'
 	root_type_name: str
+	base_type: Optional['NamedTypeReferenceType']
+	base_offset: int
 	offset: int
 	field_index: int
 
@@ -235,8 +237,15 @@ class TypeDefinitionLine:
 		type_ = Type.create(handle=core.BNNewTypeReference(struct.type), platform=platform)
 		root_type = Type.create(handle=core.BNNewTypeReference(struct.rootType), platform=platform)
 		root_type_name = core.pyNativeStr(struct.rootTypeName)
-		return TypeDefinitionLine(struct.lineType, tokens, type_, root_type, root_type_name,
-								  struct.offset, struct.fieldIndex)
+		if struct.baseType:
+			const_conf = BoolWithConfidence.get_core_struct(False, 0)
+			volatile_conf = BoolWithConfidence.get_core_struct(False, 0)
+			handle = core.BNCreateNamedTypeReference(struct.baseType, 0, 1, const_conf, volatile_conf)
+			base_type = NamedTypeReferenceType(handle, platform)
+		else:
+			base_type = None
+		return TypeDefinitionLine(struct.lineType, tokens, type_, root_type, root_type_name, base_type,
+								  struct.baseOffset, struct.offset, struct.fieldIndex)
 
 	def _to_core_struct(self):
 		struct = core.BNTypeDefinitionLine()
@@ -246,6 +255,11 @@ class TypeDefinitionLine:
 		struct.type = core.BNNewTypeReference(self.type.handle)
 		struct.rootType = core.BNNewTypeReference(self.root_type.handle)
 		struct.rootTypeName = self.root_type_name
+		if self.base_type is None:
+			struct.baseType = None
+		else:
+			struct.baseType = core.BNNewNamedTypeReference(self.base_type.ntr_handle)
+		struct.baseOffset = self.base_offset
 		struct.offset = self.offset
 		struct.fieldIndex = self.field_index
 		return struct
@@ -1151,6 +1165,56 @@ class StructureMember:
 		return len(self.type)
 
 
+@dataclass
+class InheritedStructureMember:
+	base: 'NamedTypeReferenceType'
+	base_offset: int
+	member: StructureMember
+	member_index: int
+
+	def __repr__(self):
+		if self.base is None:
+			return f"<member index {self.member_index}: {repr(self.member)}>"
+		return f"<inherited from {self.base.name} @ {self.base_offset:#x} index {self.member_index}: {repr(self.member)}>"
+
+	def __len__(self):
+		return len(self.member.type)
+
+
+@dataclass
+class BaseStructure:
+	type: 'NamedTypeReferenceType'
+	offset: int
+	width: int
+
+	def __init__(self, type: Union['NamedTypeReferenceType', 'StructureType'], offset: int, width: int = 0):
+		if isinstance(type, StructureType):
+			self.type = type.registered_name
+			self.offset = offset
+			self.width = type.width
+		else:
+			self.type = type
+			self.offset = offset
+			self.width = width
+
+	def __repr__(self):
+		return f"<base: {repr(self.type)}, offset {self.offset:#x}, width: {self.width:#x}>"
+
+	def _to_core_struct(self):
+		result = core.BNBaseStructure()
+		result.type = self.type.ntr_handle
+		result.offset = self.offset
+		result.width = self.width
+		return result
+
+	@staticmethod
+	def _from_core_struct(core_obj: core.BNBaseStructure, platform: Optional['_platform.Platform'] = None):
+		const_conf = BoolWithConfidence.get_core_struct(False, 0)
+		volatile_conf = BoolWithConfidence.get_core_struct(False, 0)
+		handle = core.BNCreateNamedTypeReference(core_obj.type, 0, 1, const_conf, volatile_conf)
+		return BaseStructure(NamedTypeReferenceType(handle, platform), core_obj.offset, core_obj.width)
+
+
 class StructureBuilder(TypeBuilder):
 	def __init__(
 	    self, handle: core.BNTypeBuilderHandle, builder_handle: core.BNStructureBuilderHandle,
@@ -1237,6 +1301,27 @@ class StructureBuilder(TypeBuilder):
 		self._add_members(members)
 
 	@property
+	def base_structures(self) -> List[BaseStructure]:
+		"""Base structure list. Offsets that are not defined by this structure will be filled
+		    in by the fields of the base structure(s)."""
+		count = ctypes.c_ulonglong()
+		bases = core.BNGetBaseStructuresForStructureBuilder(self.builder_handle, count)
+		try:
+			result = []
+			for i in range(0, count.value):
+				result.append(BaseStructure._from_core_struct(bases[i], self.platform))
+			return result
+		finally:
+			core.BNFreeBaseStructureList(bases, count.value)
+
+	@base_structures.setter
+	def base_structures(self, value: List[BaseStructure]) -> None:
+		bases = (core.BNBaseStructure * len(value))()
+		for i in range(0, len(value)):
+			bases[i] = value[i]._to_core_struct()
+		core.BNSetBaseStructuresForStructureBuilder(self.builder_handle, bases, len(value))
+
+	@property
 	def packed(self) -> bool:
 		return core.BNIsStructureBuilderPacked(self.builder_handle)
 
@@ -1261,10 +1346,26 @@ class StructureBuilder(TypeBuilder):
 		core.BNSetStructureBuilderWidth(self.builder_handle, value)
 
 	@property
+	def pointer_offset(self) -> int:
+		return core.BNGetStructureBuilderPointerOffset(self.builder_handle)
+
+	@pointer_offset.setter
+	def pointer_offset(self, value: int) -> None:
+		core.BNSetStructureBuilderPointerOffset(self.builder_handle, value)
+
+	@property
 	def union(self) -> bool:
 		return core.BNIsStructureBuilderUnion(self.builder_handle)
 
 	@property
+	def propagate_data_var_refs(self) -> bool:
+		return core.BNStructureBuilderPropagatesDataVariableReferences(self.builder_handle)
+
+	@propagate_data_var_refs.setter
+	def propagate_data_var_refs(self, value: bool) -> None:
+		core.BNSetStructureBuilderPropagatesDataVariableReferences(self.builder_handle, value)
+
+	@property
 	def type(self) -> StructureVariant:
 		return StructureVariant(core.BNGetStructureBuilderType(self.builder_handle))
 
@@ -2329,11 +2430,35 @@ class StructureType(Type):
 		return result
 
 	@property
+	def base_structures(self) -> List[BaseStructure]:
+		"""Base structure list (read-only). Offsets that are not defined by this structure will be filled
+		    in by the fields of the base structure(s)."""
+		count = ctypes.c_ulonglong()
+		bases = core.BNGetBaseStructuresForStructure(self.struct_handle, count)
+		try:
+			result = []
+			for i in range(0, count.value):
+				result.append(BaseStructure._from_core_struct(bases[i], self.platform))
+			return result
+		finally:
+			core.BNFreeBaseStructureList(bases, count.value)
+
+	@property
 	def width(self):
 		"""Structure width"""
 		return core.BNGetStructureWidth(self.struct_handle)
 
 	@property
+	def pointer_offset(self):
+		"""
+		Structure pointer offset. Pointers to this structure will implicitly
+		have this offset subtracted from the pointer to arrive at the start of the structure.
+		Effectively, the pointer offset becomes the new start of the structure, and fields
+		before it are accessed using negative offsets from the pointer.
+		"""
+		return core.BNGetStructurePointerOffset(self.struct_handle)
+
+	@property
 	def alignment(self):
 		"""Structure alignment"""
 		return core.BNGetStructureAlignment(self.struct_handle)
@@ -2343,9 +2468,45 @@ class StructureType(Type):
 		return core.BNIsStructurePacked(self.struct_handle)
 
 	@property
+	def propagate_data_var_refs(self) -> bool:
+		"""Whether structure field references propagate the references to data variable field values"""
+		return core.BNStructurePropagatesDataVariableReferences(self.struct_handle)
+
+	@property
 	def type(self) -> StructureVariant:
 		return StructureVariant(core.BNGetStructureType(self.struct_handle))
 
+	def members_including_inherited(self, view: 'binaryview.BinaryView') -> List[InheritedStructureMember]:
+		"""Returns structure member list, including those inherited by base structures"""
+		count = ctypes.c_ulonglong()
+		members = core.BNGetStructureMembersIncludingInherited(self.struct_handle, view.handle, count)
+		assert members is not None, "core.BNGetInheritedStructureMembers returned None"
+		try:
+			result = []
+			for i in range(0, count.value):
+				if members[i].base:
+					const_conf = BoolWithConfidence.get_core_struct(False, 0)
+					volatile_conf = BoolWithConfidence.get_core_struct(False, 0)
+					handle = core.BNCreateNamedTypeReference(members[i].base, 0, 1, const_conf, volatile_conf)
+					base_type = NamedTypeReferenceType(handle, self.platform)
+				else:
+					base_type = None
+				result.append(
+				    InheritedStructureMember(
+						base_type,
+						members[i].baseOffset,
+				        StructureMember(
+							Type.create(core.BNNewTypeReference(members[i].member.type), confidence=members[i].member.typeConfidence),
+							members[i].member.name, members[i].member.offset, MemberAccess(members[i].member.access),
+							MemberScope(members[i].member.scope)
+						),
+						members[i].memberIndex
+				    )
+				)
+		finally:
+			core.BNFreeInheritedStructureMemberList(members, count.value)
+		return result
+
 	def with_replaced_structure(self, from_struct, to_struct) -> 'StructureType':
 		return StructureType(
 		    core.BNStructureWithReplacedStructure(self.struct_handle, from_struct.handle, to_struct.handle)