Merge pull request #1861 from kyrre/feature/parquet-arrow-renderer

Feature/parquet arrow renderer

Author: ikelos

pyproject.toml

@@ -42,6 +42,8 @@ dev = [
     "types-jsonschema>=4.23.0,<5",
 ]
 
+arrow = ["pyarrow>=17.0.0"]
+
 test = [
     "volatility3[dev]",
     "pytest>=8.3.3,<9",

test/renderers/__init__.py

@@ -0,0 +1,109 @@
+import io
+import pytest
+from abc import ABC, abstractmethod
+from test import test_volatility
+
+HAS_PYARROW = False
+try:
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+    import pyarrow.compute as pc
+    HAS_PYARROW = True
+except ImportError:
+   # The user doesn't have pyarrow installed, but HAS_PYARROW will be false so just continue
+    pass
+
+
+@pytest.mark.skipif(not HAS_PYARROW, reason="pyarrow not installed")
+class TestArrowRendererBase(ABC):
+    """Base class for testing Arrow-based renderers.
+
+    Re-implements Windows and Linux plugin tests using PyArrow operations
+    instead of text-based assertions.
+    """
+
+    renderer_format = None  # Override in subclasses
+
+    @abstractmethod
+    def _get_table_from_output(self, output_bytes) -> "pa.Table":
+        """Parse output bytes into Arrow table. Override in subclasses."""
+
+    def test_windows_generic_pslist(self, volatility, python, image):
+        rc, out, _err = test_volatility.runvol_plugin(
+            "windows.pslist.PsList",
+            image,
+            volatility,
+            python,
+            globalargs=("-r", self.renderer_format),
+        )
+        assert rc == 0
+
+        table = self._get_table_from_output(out)
+        assert table.num_rows > 10
+
+        assert table.filter(pc.match_substring(pc.utf8_lower(table.column('ImageFileName')), "system")).num_rows > 0
+        assert table.filter(pc.match_substring(pc.utf8_lower(table.column('ImageFileName')), "csrss.exe")).num_rows > 0
+        assert table.filter(pc.match_substring(pc.utf8_lower(table.column('ImageFileName')), "svchost.exe")).num_rows > 0
+        assert table.filter(pc.greater(table.column('PID'), 0)).num_rows == table.num_rows
+
+    def test_linux_generic_pslist(self, volatility, python, image):
+        rc, out, _err = test_volatility.runvol_plugin(
+            "linux.pslist.PsList",
+            image,
+            volatility,
+            python,
+            globalargs=("-r", self.renderer_format),
+        )
+        assert rc == 0
+
+        table = self._get_table_from_output(out)
+        assert table.num_rows > 10
+
+        init_rows = table.filter(pc.match_substring(pc.utf8_lower(table.column('COMM')), "init"))
+        systemd_rows = table.filter(pc.match_substring(pc.utf8_lower(table.column('COMM')), "systemd"))
+        assert (init_rows.num_rows > 0) or (systemd_rows.num_rows > 0)
+
+        assert table.filter(pc.match_substring(pc.utf8_lower(table.column('COMM')), "watchdog")).num_rows > 0
+        assert table.filter(pc.greater(table.column('PID'), 0)).num_rows == table.num_rows
+
+    def test_windows_generic_handles(self, volatility, python, image):
+        rc, out, _err = test_volatility.runvol_plugin(
+            "windows.handles.Handles",
+            image,
+            volatility,
+            python,
+            globalargs=("-r", self.renderer_format),
+            pluginargs=("--pid", "4"),
+        )
+        assert rc == 0
+
+        table = self._get_table_from_output(out)
+        assert table.num_rows > 500
+        assert table.filter(pc.match_substring(pc.utf8_lower(table.column('Name')), "machine\\system")).num_rows > 0
+
+    def test_linux_generic_lsof(self, volatility, python, image):
+        rc, out, _err = test_volatility.runvol_plugin(
+            "linux.lsof.Lsof",
+            image,
+            volatility,
+            python,
+            globalargs=("-r", self.renderer_format),
+        )
+        assert rc == 0
+
+        table = self._get_table_from_output(out)
+        assert table.num_rows > 35
+
+class TestParquetRenderer(TestArrowRendererBase):
+    renderer_format = "parquet"
+
+    def _get_table_from_output(self, output_bytes):
+        return pq.read_table(io.BytesIO(output_bytes))
+
+
+class TestArrowRenderer(TestArrowRendererBase):
+    renderer_format = "arrow"
+
+    def _get_table_from_output(self, output_bytes):
+        return pa.ipc.open_stream(io.BytesIO(output_bytes)).read_all()
+

test/renderers/test_parquet_renderers.py

@@ -106,13 +106,6 @@ def run(self):
 
         volatility3.framework.require_interface_version(2, 0, 0)
 
-        renderers = dict(
-            [
-                (x.name.lower(), x)
-                for x in framework.class_subclasses(text_renderer.CLIRenderer)
-            ]
-        )
-
         # Load up system defaults
         delayed_logs, default_config = self.load_system_defaults("vol.json")
 
@@ -193,14 +186,6 @@ def run(self):
             default=False,
             action="store_true",
         )
-        parser.add_argument(
-            "-r",
-            "--renderer",
-            metavar="RENDERER",
-            help=f"Determines how to render the output ({', '.join(list(renderers))})",
-            default="quick",
-            choices=list(renderers),
-        )
         parser.add_argument(
             "-f",
             "--file",
@@ -270,11 +255,6 @@ def run(self):
         known_args = [arg for arg in sys.argv if arg != "--help" and arg != "-h"]
         partial_args, _ = parser.parse_known_args(known_args)
 
-        banner_output = sys.stdout
-        if renderers[partial_args.renderer].structured_output:
-            banner_output = sys.stderr
-        banner_output.write(f"Volatility 3 Framework {constants.PACKAGE_VERSION}\n")
-
         ### Start up logging
         if partial_args.log:
             file_logger = logging.FileHandler(partial_args.log)
@@ -346,6 +326,24 @@ def run(self):
 
         plugin_list = framework.list_plugins()
 
+        # Discover renderers after plugin directories are loaded
+        # This allows custom renderers to be found in plugin directories
+        renderers = dict(
+            [
+                (x.name.lower(), x)
+                for x in framework.class_subclasses(text_renderer.CLIRenderer)
+            ]
+        )
+
+        parser.add_argument(
+            "-r",
+            "--renderer",
+            metavar="RENDERER",
+            help=f"Determines how to render the output ({', '.join(list(renderers))})",
+            default="quick",
+            choices=list(renderers),
+        )
+
         seen_automagics = set()
         chosen_configurables_list = {}
         for amagic in automagics:
@@ -392,6 +390,13 @@ def run(self):
             # before all the plugins have been added
             argcomplete.autocomplete(parser)
         args = parser.parse_args()
+
+        # Display banner - redirect to stderr if using structured output
+        banner_output = sys.stdout
+        if renderers[args.renderer].structured_output:
+            banner_output = sys.stderr
+        banner_output.write(f"Volatility 3 Framework {constants.PACKAGE_VERSION}\n")
+
         if args.plugin is None:
             parser.error(
                 f"Please select a plugin to run (see '{self.CLI_NAME} --help' for options"