Move RBTree to the rb_root extension object for better integration

Author: gcmoreira

volatility3/framework/symbols/linux/__init__.py

@@ -46,6 +46,7 @@ def __init__(self, *args, **kwargs) -> None:
         # Might not exist in older kernels or the current symbols
         self.optional_set_type_class("mount", extensions.mount)
         self.optional_set_type_class("mnt_namespace", extensions.mnt_namespace)
+        self.optional_set_type_class("rb_root", extensions.rb_root)
 
         # Network
         self.set_type_class("net", extensions.net)
@@ -425,36 +426,3 @@ def get_module_from_volobj_type(
         kernel = context.modules[kernel_module_name]
 
         return kernel
-
-
-class RBTree(object):
-    """Simple Red-Black tree abstraction"""
-
-    def __init__(self, root):
-        self.root = root
-
-    def _walk_nodes(self, root_node) -> Iterator[int]:
-        """Traverses the Red-Black tree from the root node and yields a pointer to each
-        node in this tree.
-
-        Args:
-            root_node: A Red-Black tree node from which to start descending
-
-        Yields:
-            A pointer to every node descending from the specified root node
-        """
-        if not root_node:
-            return
-
-        yield root_node
-        yield from self._walk_nodes(root_node.rb_left)
-        yield from self._walk_nodes(root_node.rb_right)
-
-    def get_nodes(self) -> Iterator[int]:
-        """Yields a pointer to each node in the Red-Black tree
-
-        Yields:
-            A pointer to every node in the Red-Black tree
-        """
-
-        yield from self._walk_nodes(root_node=self.root.rb_node)

volatility3/framework/symbols/linux/extensions/__init__.py

@@ -1326,7 +1326,7 @@ def get_mount_points(
             vmlinux = linux.LinuxUtilities.get_module_from_volobj_type(
                 self._context, self
             )
-            for node in linux.RBTree(self.mounts).get_nodes():
+            for node in self.mounts.get_nodes():
                 mnt = linux.LinuxUtilities.container_of(
                     node, "mount", "mnt_list", vmlinux
                 )
@@ -1925,3 +1925,31 @@ def get_file_mode(self) -> str:
             The inode's file mode string
         """
         return stat.filemode(self.i_mode)
+
+
+class rb_root(objects.StructType):
+    def _walk_nodes(self, root_node) -> Iterator[int]:
+        """Traverses the Red-Black tree from the root node and yields a pointer to each
+        node in this tree.
+
+        Args:
+            root_node: A Red-Black tree node from which to start descending
+
+        Yields:
+            A pointer to every node descending from the specified root node
+        """
+        if not root_node:
+            return
+
+        yield root_node
+        yield from self._walk_nodes(root_node.rb_left)
+        yield from self._walk_nodes(root_node.rb_right)
+
+    def get_nodes(self) -> Iterator[int]:
+        """Yields a pointer to each node in the Red-Black tree
+
+        Yields:
+            A pointer to every node in the Red-Black tree
+        """
+
+        yield from self._walk_nodes(root_node=self.rb_node)