Add the Components API

4 files changed, 479 insertions, 1 deletions

diff --git a/python/__init__.py b/python/__init__.py
index ef316c36..1e73f2cc 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -67,6 +67,7 @@ from .database import *
 from .secretsprovider import *
 from .typeparser import *
 from .typeprinter import *
+from .component import *
 # We import each of these by name to prevent conflicts between
 # log.py and the function 'log' which we don't import below
 from .log import (
diff --git a/python/binaryview.py b/python/binaryview.py
index 278cbecd..e68bf69d 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -48,6 +48,7 @@ from . import typelibrary
 from . import fileaccessor
 from . import databuffer
 from . import basicblock
+from . import component
 from . import lineardisassembly
 from . import metadata
 from . import highlight
@@ -228,6 +229,30 @@ class BinaryDataNotification:
 	def section_removed(self, view: 'BinaryView', section: 'Section') -> None:
 		pass
 
+	def component_added(self,  view: 'BinaryView', _component: component.Component) -> None:
+		pass
+
+	def component_removed(self, view: 'BinaryView', formerParent: component.Component,
+						  _component: component.Component) -> None:
+		pass
+
+	def component_name_updated(self,  view: 'BinaryView', previous_name: str, _component: component.Component) -> None:
+		pass
+
+	def component_moved(self, view: 'BinaryView', formerParent: component.Component, newParent: component.Component,
+						_component: component.Component) -> None:
+		pass
+
+	def component_function_added(self, view: 'BinaryView', _component: component.Component, func: '_function.Function'):
+		pass
+
+	def component_function_removed(self, view: 'BinaryView', _component: component.Component,
+								   func: '_function.Function'):
+		pass
+
+
+
+
 class StringReference:
 	_decodings = {
 	    StringType.AsciiString: "ascii", StringType.Utf8String: "utf-8", StringType.Utf16String: "utf-16",
@@ -462,7 +487,12 @@ class BinaryDataNotificationCallbacks:
 		self._cb.sectionAdded = self._cb.sectionAdded.__class__(self._section_added)
 		self._cb.sectionUpdated = self._cb.sectionUpdated.__class__(self._section_updated)
 		self._cb.sectionRemoved = self._cb.sectionRemoved.__class__(self._section_removed)
-
+		self._cb.componentNameUpdated = self._cb.componentNameUpdated.__class__(self._component_name_updated)
+		self._cb.componentAdded = self._cb.componentAdded.__class__(self._component_added)
+		self._cb.componentRemoved = self._cb.componentRemoved.__class__(self._component_removed)
+		self._cb.componentMoved = self._cb.componentMoved.__class__(self._component_moved)
+		self._cb.componentFunctionAdded = self._cb.componentFunctionAdded.__class__(self._component_function_added)
+		self._cb.componentFunctionRemoved = self._cb.componentFunctionRemoved.__class__(self._component_function_removed)
 
 	def _register(self) -> None:
 		core.BNRegisterDataNotification(self._view.handle, self._cb)
@@ -736,6 +766,78 @@ class BinaryDataNotificationCallbacks:
 		except:
 			log_error(traceback.format_exc())
 
+	def _component_added(self, ctxt, view: core.BNBinaryView, _component: core.BNComponent):
+		try:
+			component_handle = core.BNNewComponentReference(_component)
+			assert component_handle is not None, "core.BNNewComponentReference returned None"
+			result = component.Component(component_handle)
+			self._notify.component_added(self._view, result)
+		except:
+			log_error(traceback.format_exc())
+
+	def _component_removed(self, ctxt, view: core.BNBinaryView, formerParent: core.BNComponent, _component: core.BNComponent):
+		try:
+			formerParent_handle = core.BNNewComponentReference(formerParent)
+			assert formerParent_handle is not None, "core.BNNewComponentReference returned None"
+			formerParentResult = component.Component(formerParent_handle)
+			component_handle = core.BNNewComponentReference(_component)
+			assert component_handle is not None, "core.BNNewComponentReference returned None"
+			result = component.Component(component_handle)
+			self._notify.component_removed(self._view, formerParentResult, result)
+		except:
+			log_error(traceback.format_exc())
+
+	def _component_name_updated(self, ctxt, view: core.BNBinaryView, previous_name: str, _component: core.BNComponent):
+		try:
+			component_handle = core.BNNewComponentReference(_component)
+			assert component_handle is not None, "core.BNNewComponentReference returned None"
+			result = component.Component(component_handle)
+			self._notify.component_name_updated(self._view, previous_name, result)
+		except:
+			log_error(traceback.format_exc())
+
+	def _component_moved(self, ctxt, view: core.BNBinaryView, formerParent: core.BNComponent,
+						 newParent: core.BNComponent, _component: core.BNComponent):
+		try:
+			formerParent_handle = core.BNNewComponentReference(formerParent)
+			assert formerParent_handle is not None, "core.BNNewComponentReference returned None"
+			formerParentResult = component.Component(formerParent_handle)
+			newParent_handle = core.BNNewComponentReference(newParent)
+			assert newParent_handle is not None, "core.BNNewComponentReference returned None"
+			newParentResult = component.Component(newParent_handle)
+			component_handle = core.BNNewComponentReference(_component)
+			assert component_handle is not None, "core.BNNewComponentReference returned None"
+			result = component.Component(component_handle)
+			self._notify.component_moved(self._view, formerParentResult, newParentResult, result)
+		except:
+			log_error(traceback.format_exc())
+
+	def _component_function_added(self, ctxt, view: core.BNBinaryView, _component: core.BNComponent,
+								  func: '_function.Function'):
+		try:
+			component_handle = core.BNNewComponentReference(_component)
+			assert component_handle is not None, "core.BNNewComponentReference returned None"
+			result = component.Component(component_handle)
+			function_handle = core.BNNewFunctionReference(func)
+			assert function_handle is not None, "core.BNNewFunctionReference returned None"
+			function = _function.Function(self._view, function_handle)
+			self._notify.component_function_added(self._view, result, function)
+		except:
+			log_error(traceback.format_exc())
+
+	def _component_function_removed(self, ctxt, view: core.BNBinaryView, _component: core.BNComponent,
+								func: '_function.Function'):
+		try:
+			component_handle = core.BNNewComponentReference(_component)
+			assert component_handle is not None, "core.BNNewComponentReference returned None"
+			result = component.Component(component_handle)
+			function_handle = core.BNNewFunctionReference(func)
+			assert function_handle is not None, "core.BNNewFunctionReference returned None"
+			function = _function.Function(self._view, function_handle)
+			self._notify.component_function_removed(self._view, result, function)
+		except:
+			log_error(traceback.format_exc())
+
 	@property
 	def view(self) -> 'BinaryView':
 		return self._view
@@ -5827,6 +5929,105 @@ class BinaryView:
 	def notify_data_removed(self, offset: int, length: int) -> None:
 		core.BNNotifyDataRemoved(self.handle, offset, length)
 
+	def get_component(self, guid: str) -> Optional[component.Component]:
+		"""
+		Lookup a Component by its Guid
+
+		:param guid: Guid of the component to look up
+		:return: The Component with that Guid
+		"""
+		bn_component = core.BNGetComponentByGuid(self.handle, guid)
+		if bn_component is None:
+			return None
+		return component.Component(bn_component)
+
+	def get_component_by_path(self, path: str) -> Optional[component.Component]:
+		"""
+		Lookup a Component by its pathname
+
+		:note: This is a convenience method, and for performance-sensitive lookups, GetComponentByGuid is very highly recommended.
+
+		Lookups are done based on the .display_name of the Component.
+
+		All lookups are absolute from the root component, and are case-sensitive. Pathnames are delimited with "/"
+
+		:param path: Pathname of the desired Component
+		:return: The Component at that pathname
+		:Example:
+
+			>>> c = bv.create_component(name="MyComponent")
+			>>> c2 = bv.create_component(name="MySubComponent", parent=c)
+			>>> bv.get_component_by_path("/MyComponent/MySubComponent") == c2
+			True
+			>>> c3 = bv.create_component(name="MySubComponent", parent=c)
+			>>> c3
+			<Component "MySubComponent (1)" "(20712aff...")>
+			>>> bv.get_component_by_path("/MyComponent/MySubComponent (1)") == c3
+			True
+		"""
+		if not isinstance(path, str):
+			raise TypeError("Pathname must be a string")
+		bn_component = core.BNGetComponentByPath(self.handle, path)
+		if bn_component is None:
+			return None
+		return component.Component(bn_component)
+
+	@property
+	def root_component(self) -> component.Component:
+		"""
+		The root component for the BinaryView (read-only)
+
+		This Component cannot be removed, and houses all unparented Components.
+
+		:return: The root component
+		"""
+		return component.Component(core.BNGetRootComponent(self.handle))
+
+	def create_component(self, name: Optional[str] = None, parent: Union[component.Component, str, None] = None) -> component.Component:
+		"""
+		Create a new component with an optional name and parent.
+
+		The `parent` argument can be either a Component or the Guid of a component that the created component will be
+			added as a child of
+
+		:param name: Optional name to create the component with
+		:param parent: Optional parent to which the component will be added
+		:return: The created component
+		"""
+
+		if parent:
+			if isinstance(parent, component.Component):
+				if name:
+					return component.Component(core.BNCreateComponentWithParentAndName(self.handle, parent.guid, name))
+				else:
+					return component.Component(core.BNCreateComponentWithParent(self.handle, parent.guid))
+			elif isinstance(parent, str):
+				if name:
+					return component.Component(core.BNCreateComponentWithParentAndName(self.handle, parent, name))
+				else:
+					return component.Component(core.BNCreateComponentWithParent(self.handle, parent))
+			else:
+				raise TypeError("parent can only be a Component object or string GUID representing one")
+		else:
+			if name:
+				return component.Component(core.BNCreateComponentWithName(self.handle, name))
+			else:
+				return component.Component(core.BNCreateComponent(self.handle))
+
+	def remove_component(self, _component: Union[component.Component, str]) -> bool:
+		"""
+		Remove a component from the tree entirely.
+
+		:param _component: Component to remove
+		:return: Whether the removal was successful
+		"""
+		if isinstance(_component, component.Component):
+			return core.BNRemoveComponent(self.handle, _component.handle)
+		elif isinstance(_component, str):
+			return core.BNRemoveComponentByGuid(self.handle, _component)
+
+		raise TypeError("Removal is only supported with a Component or string representing its Guid")
+
 	def get_strings(self, start: Optional[int] = None, length: Optional[int] = None) -> List['StringReference']:
 		"""
 		``get_strings`` returns a list of strings defined in the binary in the optional virtual address range:
diff --git a/python/component.py b/python/component.py
new file mode 100644
index 00000000..21942297
--- /dev/null
+++ b/python/component.py
@@ -0,0 +1,266 @@
+
+import ctypes
+import inspect
+from typing import Generator, Optional, List, Tuple, Union, Mapping, Any, Dict, Iterator
+from dataclasses import dataclass
+
+from . import binaryview
+
+from . import function
+from . import _binaryninjacore as core
+from . import types
+
+
+class Component:
+    """
+    Components are objects that can contain Functions and other Components.
+
+    They can be queried for information about the functions contained within them.
+
+    Components have a Guid, which persistent across saves and loads of the database, and should be
+    used for retrieving components when such is required and a reference to the Component cannot be held.
+
+    """
+    def __init__(self, handle=None):
+
+        assert handle is not None, "Cannot create component directly, run `bv.create_component?`"
+
+        self.handle = handle
+
+        self.guid = core.BNComponentGetGuid(self.handle)
+
+    def __eq__(self, other):
+        if not isinstance(other, Component):
+            return NotImplemented
+        return core.BNComponentsEqual(self.handle, other.handle)
+
+    def __ne__(self, other):
+        if not isinstance(other, Component):
+            return NotImplemented
+        return core.BNComponentsNotEqual(self.handle, other.handle)
+
+    def __repr__(self):
+        return f'<Component "{self.display_name}" "({self.guid[:8]}...")>'
+
+    def __del__(self):
+        core.BNFreeComponent(self.handle)
+
+    def __str__(self):
+        return self._sprawl_component(self)
+
+    def _sprawl_component(self, c, depth=1, out=None):
+        """
+        Recursive quick function to print out the component's tree of items
+
+        :param c: Current cycle's component. On initial call, pass `self`
+        :param depth: Current tree depth.
+        :param out: Current text
+        :return:
+        """
+        _out = ([repr(c)] if not out else out.split('\n')) + [('  ' * depth + repr(f)) for f in c.functions]
+        _out += ['  ' * (depth+1) + repr(i) for i in (c.get_referenced_data_variables() + c.get_referenced_types())]
+        for i in c.components:
+            _out.append('  ' * depth + repr(i))
+            _out = self._sprawl_component(i, depth+1, '\n'.join(_out)).split('\n')
+        return '\n'.join(_out)
+
+    def add_function(self, func: 'function.Function') -> bool:
+        """
+        Add function to this component.
+
+        :param func: Function to add
+        :return: True if function was successfully added.
+        """
+        return core.BNComponentAddFunctionReference(self.handle, func.handle)
+
+    def contains_function(self, func: 'function.Function') -> bool:
+        """
+        Check whether this component contains a function.
+
+        :param func: Function to check
+        :return: True if this component contains the function.
+        """
+        return core.BNComponentContainsFunction(self.handle, func.handle)
+
+    def add_component(self, component: 'Component') -> bool:
+        """
+        Move component to this component. This will remove it from the old parent.
+
+        :param component: Component to add to this component.
+        :return: True if the component was successfully moved to this component
+        """
+        return core.BNComponentAddComponent(self.handle, component.handle)
+
+    def remove_function(self, func: 'function.Function') -> bool:
+        """
+        Remove function from this component.
+
+        :param func: Function to remove
+        :return: True if function was successfully removed.
+        """
+        return core.BNComponentRemoveFunctionReference(self.handle, func.handle)
+
+    def remove_component(self, component: 'Component') -> bool:
+        """
+        Remove a component from the current component, moving it to the root.
+
+        This function has no effect when used from the root component.
+        Use `BinaryView.remove_component` to Remove a component from the tree entirely.
+
+        :param component: Component to remove
+        :return:
+        """
+
+        return self.view.root_component.add_component(component)
+
+    def contains_component(self, component: 'Component') -> bool:
+        """
+        Check whether this component contains a component.
+
+        :param component: Component to check
+        :return: True if this component contains the component.
+        """
+        return core.BNComponentContainsComponent(self.handle, component.handle)
+
+    @property
+    def display_name(self) -> str:
+        """Original Name of the component (read-only)"""
+        return core.BNComponentGetDisplayName(self.handle)
+
+    @property
+    def name(self) -> str:
+        """Original name set for this component
+
+        :note: The `.display_name` property should be used for `bv.get_component_by_path()` lookups.
+
+        This can differ from the .display_name property if one of its sibling components has the same .original_name; In that
+        case, .name will be an automatically generated unique name (e.g. "MyComponentName (1)") while .original_name will
+        remain what was originally set (e.g. "MyComponentName")
+
+        If this component has a duplicate name and is moved to a component where none of its siblings share its name,
+        the .name property will return the original "MyComponentName"
+        """
+        return core.BNComponentGetOriginalName(self.handle)
+
+    @name.setter
+    def name(self, _name):
+        core.BNComponentSetName(self.handle, _name)
+
+    @property
+    def parent(self) -> Optional['Component']:
+        """
+        The component that contains this component, if it exists.
+        """
+        bn_component = core.BNComponentGetParent(self.handle)
+        if bn_component is not None:
+            return Component(bn_component)
+        return None
+
+    @property
+    def view(self):
+        bn_binaryview = core.BNComponentGetView(self.handle)
+        if bn_binaryview is not None:
+            return binaryview.BinaryView(handle=bn_binaryview)
+        return None
+
+    @property
+    def components(self) -> Iterator['Component']:
+        """
+		``components`` is an iterator for all Components contained within this Component
+
+		:return: An iterator containing Components
+		:rtype: SubComponentIterator
+		:Example:
+
+			>>> for subcomp in component.components:
+			...  print(repr(component))
+        """
+
+        @dataclass
+        class SubComponentIterator:
+            comp: Component
+
+            def __iter__(self):
+                count = ctypes.c_ulonglong(0)
+                bn_components = core.BNComponentGetContainedComponents(self.comp.handle, count)
+                try:
+                    for i in range(count.value):
+                        yield Component(core.BNNewComponentReference(bn_components[i]))
+                finally:
+                    core.BNFreeComponents(bn_components, count.value)
+
+        return iter(SubComponentIterator(self))
+
+    @property
+    def functions(self) -> Iterator['function.Function']:
+        """
+		``functions`` is an iterator for all Functions contained within this Component
+
+		:return: An iterator containing Components
+		:rtype: ComponentIterator
+		:Example:
+
+			>>> for func in component.functions:
+			...  print(func.name)
+        """
+        @dataclass
+        class FunctionIterator:
+            view: 'binaryview.BinaryView'
+            comp: Component
+
+            def __iter__(self):
+                count = ctypes.c_ulonglong(0)
+                bn_functions = core.BNComponentGetContainedFunctions(self.comp.handle, count)
+                try:
+                    for i in range(count.value):
+                        bn_function = core.BNNewFunctionReference(bn_functions[i])
+                        yield function.Function(self.view, bn_function)
+                finally:
+                    core.BNFreeFunctionList(bn_functions, count.value)
+
+        return iter(FunctionIterator(self.view, self))
+
+    def get_referenced_data_variables(self, recursive=False):
+        """
+        Get data variables referenced by this component
+
+        :param recursive: Optional; Get all DataVariables referenced by this component and subcomponents.
+        :return: List of DataVariables
+        """
+        data_vars = []
+        count = ctypes.c_ulonglong(0)
+        if recursive:
+            bn_data_vars = core.BNComponentGetReferencedDataVariablesRecursive(self.handle, count)
+        else:
+            bn_data_vars = core.BNComponentGetReferencedDataVariables(self.handle, count)
+        try:
+            for i in range(count.value):
+                bn_data_var = bn_data_vars[i]
+                data_var = binaryview.DataVariable.from_core_struct(bn_data_var, self.view)
+                data_vars.append(data_var)
+        finally:
+            core.BNFreeDataVariables(bn_data_vars, count.value)
+        return data_vars
+
+    def get_referenced_types(self, recursive=False):
+        """
+        Get Types referenced by this component
+
+        :param recursive: Optional; Get all Types referenced by this component and subcomponents.
+        :return: List of Types
+        """
+        _types = []
+        count = ctypes.c_ulonglong(0)
+
+        if recursive:
+            bn_types = core.BNComponentGetReferencedTypesRecursive(self.handle, count)
+        else:
+            bn_types = core.BNComponentGetReferencedTypes(self.handle, count)
+
+        try:
+            for i in range(count.value):
+                _types.append(types.Type(core.BNNewTypeReference(bn_types[i])))
+        finally:
+            core.BNComponentFreeReferencedTypes(bn_types, count.value)
+
+        return _types
diff --git a/python/function.py b/python/function.py
index 93ed8b49..d093c076 100644
--- a/python/function.py
+++ b/python/function.py
@@ -61,6 +61,7 @@ from .variable import (
 )
 from . import decorators
 from .enums import RegisterValueType
+from . import component
 
 ExpressionIndex = int
 InstructionIndex = int
@@ -2196,6 +2197,15 @@ class Function:
 		core.BNFreeTagList(tags, count.value)
 		return result
 
+	@property
+	def parent_components(self):
+		_components = []
+		count = ctypes.c_ulonglong(0)
+		bn_components = core.BNGetFunctionParentComponents(self.handle, count)
+		for i in range(count.value):
+			_components.append(component.Component(self.view, bn_components[i]))
+		return _components
+
 	def get_lifted_il_at(
 	    self, addr: int, arch: Optional['architecture.Architecture'] = None
 	) -> Optional['lowlevelil.LowLevelILInstruction']: