Linux: change pstree to a basic plugin rather than inheriting from pslist

Author: eve-mem

volatility3/framework/plugins/linux/pstree.py

@@ -2,18 +2,49 @@
 # which is available at https://www.volatilityfoundation.org/license/vsl-v1.0
 #
 
+from volatility3.framework import interfaces, renderers
+from volatility3.framework.configuration import requirements
+from volatility3.framework.renderers import format_hints
 from volatility3.plugins.linux import pslist
 
 
-class PsTree(pslist.PsList):
+class PsTree(interfaces.plugins.PluginInterface):
     """Plugin for listing processes in a tree based on their parent process
     ID."""
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._tasks = {}
-        self._levels = {}
-        self._children = {}
+    _required_framework_version = (2, 0, 0)
+
+    @classmethod
+    def get_requirements(cls):
+        # Since we're calling the plugin, make sure we have the plugin's requirements
+        return [
+            requirements.ModuleRequirement(
+                name="kernel",
+                description="Linux kernel",
+                architectures=["Intel32", "Intel64"],
+            ),
+            requirements.PluginRequirement(
+                name="pslist", plugin=pslist.PsList, version=(2, 2, 0)
+            ),
+            requirements.ListRequirement(
+                name="pid",
+                description="Filter on specific process IDs",
+                element_type=int,
+                optional=True,
+            ),
+            requirements.BooleanRequirement(
+                name="threads",
+                description="Include user threads",
+                optional=True,
+                default=False,
+            ),
+            requirements.BooleanRequirement(
+                name="decorate_comm",
+                description="Show `user threads` comm in curly brackets, and `kernel threads` comm in square brackets",
+                optional=True,
+                default=False,
+            ),
+        ]
 
     def find_level(self, pid: int) -> None:
         """Finds how deep the PID is in the tasks hierarchy.
@@ -40,32 +71,26 @@ def find_level(self, pid: int) -> None:
 
     def _generator(
         self,
-        pid_filter,
-        include_threads: bool = False,
-        decorate_com: bool = False,
-        dump: bool = False,
+        tasks: list,
+        decorate_comm: bool = False,
     ):
         """Generates the tasks hierarchy tree.
 
         Args:
-            pid_filter: A function which takes a process object and returns True if the process should be ignored/filtered
-            include_threads: If True, the output will also show the user threads
-                             If False, only the thread group leaders will be shown
-                             Defaults to False.
+            tasks: A list of task objects to be displayed
             decorate_comm: If True, it decorates the comm string of
                             - User threads: in curly brackets,
                             - Kernel threads: in square brackets
                            Defaults to False.
         Yields:
             Each rows
         """
-        vmlinux = self.context.modules[self.config["kernel"]]
-        for proc in self.list_tasks(
-            self.context,
-            vmlinux.name,
-            filter_func=pid_filter,
-            include_threads=include_threads,
-        ):
+
+        self._tasks = {}
+        self._levels = {}
+        self._children = {}
+
+        for proc in tasks:
             self._tasks[proc.pid] = proc
 
         # Build the child/level maps
@@ -75,13 +100,10 @@ def _generator(
         def yield_processes(pid):
             task = self._tasks[pid]
 
-            row = self._get_task_fields(task, decorate_com)
-            if dump:
-                file_output = self._get_file_output(task)
-            else:
-                file_output = "Disabled"
-            row = self._get_task_fields(task, decorate_com)
-            row += (file_output,)  # also add the file output column
+            row = pslist.PsList.get_task_fields(task, decorate_comm)
+            # update the first element, the offset, in the row tuple to use format_hints.Hex
+            # as a simple int is returned from get_task_fields.
+            row = (format_hints.Hex(row[0]),) + row[1:]
 
             tid = task.pid
             yield (self._levels[tid] - 1, row)
@@ -92,3 +114,27 @@ def yield_processes(pid):
         for pid, level in self._levels.items():
             if level == 1:
                 yield from yield_processes(pid)
+
+    def run(self):
+        filter_func = pslist.PsList.create_pid_filter(self.config.get("pid", None))
+        include_threads = self.config.get("threads")
+        decorate_comm = self.config.get("decorate_comm")
+
+        return renderers.TreeGrid(
+            [
+                ("OFFSET (V)", format_hints.Hex),
+                ("PID", int),
+                ("TID", int),
+                ("PPID", int),
+                ("COMM", str),
+            ],
+            self._generator(
+                pslist.PsList.list_tasks(
+                    self.context,
+                    self.config["kernel"],
+                    filter_func=filter_func,
+                    include_threads=include_threads,
+                ),
+                decorate_comm=decorate_comm,
+            ),
+        )