Allow DataVariables to be added to Components

3 files changed, 115 insertions, 55 deletions

diff --git a/python/binaryview.py b/python/binaryview.py
index 8005630b..0ca357a4 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -252,6 +252,12 @@ class BinaryDataNotification:
 								   func: '_function.Function'):
 		pass
 
+	def component_data_var_added(self, view: 'BinaryView', _component: component.Component, var: 'DataVariable'):
+		pass
+
+	def component_data_var_removed(self, view: 'BinaryView', _component: component.Component, var: 'DataVariable'):
+		pass
+
 
 
 
@@ -495,6 +501,8 @@ class BinaryDataNotificationCallbacks:
 		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)
+		self._cb.componentDataVariableAdded = self._cb.componentDataVariableAdded.__class__(self._component_data_variable_added)
+		self._cb.componentDataVariableRemoved = self._cb.componentDataVariableRemoved.__class__(self._component_data_variable_removed)
 
 	def _register(self) -> None:
 		core.BNRegisterDataNotification(self._view.handle, self._cb)
@@ -840,6 +848,26 @@ class BinaryDataNotificationCallbacks:
 		except:
 			log_error(traceback.format_exc())
 
+	def _component_data_variable_added(self, ctxt, view: core.BNBinaryView, _component: core.BNComponent,
+									   var: core.BNDataVariable):
+		try:
+			component_handle = core.BNNewComponentReference(_component)
+			assert component_handle is not None, "core.BNNewComponentReference returned None"
+			result = component.Component(component_handle)
+			self._notify.component_data_var_added(self._view, result, DataVariable.from_core_struct(var[0], self._view))
+		except:
+			log_error(traceback.format_exc())
+
+	def _component_data_variable_removed(self, ctxt, view: core.BNBinaryView, _component: core.BNComponent,
+									var: core.BNDataVariable):
+		try:
+			component_handle = core.BNNewComponentReference(_component)
+			assert component_handle is not None, "core.BNNewComponentReference returned None"
+			result = component.Component(component_handle)
+			self._notify.component_data_var_removed(self._view, result, DataVariable.from_core_struct(var[0], self._view))
+		except:
+			log_error(traceback.format_exc())
+
 	@property
 	def view(self) -> 'BinaryView':
 		return self._view
@@ -6189,6 +6217,31 @@ class BinaryView:
 
 		raise TypeError("Removal is only supported with a Component or string representing its Guid")
 
+	def get_function_parent_components(self, function: 'function.Function') -> List['component.Component']:
+		_components = []
+		count = ctypes.c_ulonglong(0)
+		bn_components = core.BNGetFunctionParentComponents(self.handle, function.handle, count)
+		try:
+			for i in range(count.value):
+				_components.append(component.Component(core.BNNewComponentReference(bn_components[i])))
+		finally:
+			core.BNFreeComponents(bn_components, count.value)
+		return _components
+
+	def get_data_variable_parent_components(self, data_variable: 'DataVariable') -> List['component.Component']:
+		_components = []
+		count = ctypes.c_ulonglong(0)
+		bn_components = core.BNGetDataVariableParentComponents(self.handle, data_variable.address, count)
+		try:
+			for i in range(count.value):
+				_components.append(component.Component(core.BNNewComponentReference(bn_components[i])))
+		finally:
+			core.BNFreeComponents(bn_components, count.value)
+		return _components
+
+	def get_constant_data(self, addr: int) -> databuffer.DataBuffer:
+		return databuffer.DataBuffer(handle=core.BNGetConstantData(self.handle, addr))
+
 	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
index 21942297..210ec8c9 100644
--- a/python/component.py
+++ b/python/component.py
@@ -43,7 +43,8 @@ class Component:
         return f'<Component "{self.display_name}" "({self.guid[:8]}...")>'
 
     def __del__(self):
-        core.BNFreeComponent(self.handle)
+        if (hasattr(self, 'handle')):
+            core.BNFreeComponent(self.handle)
 
     def __str__(self):
         return self._sprawl_component(self)
@@ -82,6 +83,15 @@ class Component:
         """
         return core.BNComponentContainsFunction(self.handle, func.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 add_component(self, component: 'Component') -> bool:
         """
         Move component to this component. This will remove it from the old parent.
@@ -91,14 +101,14 @@ class Component:
         """
         return core.BNComponentAddComponent(self.handle, component.handle)
 
-    def remove_function(self, func: 'function.Function') -> bool:
+    def contains_component(self, component: 'Component') -> bool:
         """
-        Remove function from this component.
+        Check whether this component contains a component.
 
-        :param func: Function to remove
-        :return: True if function was successfully removed.
+        :param component: Component to check
+        :return: True if this component contains the component.
         """
-        return core.BNComponentRemoveFunctionReference(self.handle, func.handle)
+        return core.BNComponentContainsComponent(self.handle, component.handle)
 
     def remove_component(self, component: 'Component') -> bool:
         """
@@ -113,14 +123,14 @@ class Component:
 
         return self.view.root_component.add_component(component)
 
-    def contains_component(self, component: 'Component') -> bool:
-        """
-        Check whether this component contains a component.
+    def add_data_variable(self, data_variable):
+        return core.BNComponentAddDataVariable(self.handle, data_variable.address)
 
-        :param component: Component to check
-        :return: True if this component contains the component.
-        """
-        return core.BNComponentContainsComponent(self.handle, component.handle)
+    def contains_data_variable(self, data_variable):
+        return core.BNComponentContainsDataVariable(self.handle, data_variable.address)
+
+    def remove_data_variable(self, data_variable):
+        return core.BNComponentRemoveDataVariable(self.handle, data_variable.address)
 
     @property
     def display_name(self) -> str:
@@ -164,61 +174,67 @@ class Component:
         return None
 
     @property
-    def components(self) -> Iterator['Component']:
+    def components(self) -> List['Component']:
         """
 		``components`` is an iterator for all Components contained within this Component
 
-		:return: An iterator containing Components
-		:rtype: SubComponentIterator
+		:return: A list of components
 		: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)
+        count = ctypes.c_ulonglong(0)
+        bn_components = core.BNComponentGetContainedComponents(self.handle, count)
+        components = []
+        try:
+            for i in range(count.value):
+                components.append(Component(core.BNNewComponentReference(bn_components[i])))
+        finally:
+            core.BNFreeComponents(bn_components, count.value)
 
-        return iter(SubComponentIterator(self))
+        return components
 
     @property
-    def functions(self) -> Iterator['function.Function']:
+    def functions(self) -> List['function.Function']:
         """
 		``functions`` is an iterator for all Functions contained within this Component
 
-		:return: An iterator containing Components
-		:rtype: ComponentIterator
+		:return: A list of functions
 		: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)
+        count = ctypes.c_ulonglong(0)
+        bn_functions = core.BNComponentGetContainedFunctions(self.handle, count)
+        funcs = []
+        try:
+            for i in range(count.value):
+                bn_function = core.BNNewFunctionReference(bn_functions[i])
+                funcs.append(function.Function(self.view, bn_function))
+        finally:
+            core.BNFreeFunctionList(bn_functions, count.value)
+
+        return funcs
 
-        return iter(FunctionIterator(self.view, self))
+    @property
+    def data_variables(self):
+        data_vars = []
+
+        count = ctypes.c_ulonglong(0)
+        bn_data_vars = core.BNComponentGetContainedDataVariables(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_data_variables(self, recursive=False):
         """
diff --git a/python/function.py b/python/function.py
index 86662cf7..bce4f843 100644
--- a/python/function.py
+++ b/python/function.py
@@ -2233,15 +2233,6 @@ 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']: