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

Feature/parquet arrow renderer

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/test_parquet_renderers.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()
+

volatility3/cli/__init__.py

@@ -106,13 +106,6 @@ class CommandLine:
 
         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 @@ class CommandLine:
             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 @@ class CommandLine:
         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 @@ class CommandLine:
 
         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 @@ class CommandLine:
             # 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"

volatility3/framework/plugins/renderers/parquet_renderer.py

@@ -0,0 +1,218 @@
+# This file is Copyright 2019 Volatility Foundation and licensed under the Volatility Software License 1.0
+# which is available at https://www.volatilityfoundation.org/license/vsl-v1.0
+#
+import datetime
+import logging
+import sys
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    TextIO,
+)
+from volatility3.framework import interfaces, renderers
+from volatility3.framework.renderers import format_hints
+from volatility3.cli import text_renderer
+
+vollog = logging.getLogger(__name__)
+
+ARROW_PRESENT = False
+try:
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    ARROW_PRESENT = True
+except ImportError:
+    vollog.debug("Arrow/Parquet libraries not found")
+
+
+class ArrowRenderer(text_renderer.CLIRenderer):
+    """Renderer that outputs Arrow IPC format data."""
+
+    name = "arrow"
+    structured_output = True
+    _version = (1, 0, 0)
+
+    def __init__(
+        self, options: Optional[List[interfaces.renderers.RenderOption]] = None
+    ) -> None:
+        super().__init__(options)
+
+        if not ARROW_PRESENT:
+            raise RuntimeError("Arrow output format requires the pyarrow package")
+
+        self._to_arrow_type = {
+            renderers.Disassembly: pa.utf8,
+            bool: pa.bool_,
+            int: pa.int64,
+            float: pa.float64,
+            str: pa.utf8,
+            datetime.datetime: lambda: pa.timestamp("ms"),
+            format_hints.Bin: pa.uint64,
+            format_hints.Hex: pa.uint64,
+            format_hints.MultiTypeData: pa.utf8,
+            format_hints.HexBytes: pa.binary,
+            renderers.LayerData: pa.binary,
+            bytes: pa.binary,
+        }
+
+        # indicates if the output from the plugin is nested, e.g., pstree
+        # which would then need to be flattened
+        self._is_tree_result = False
+        self._node_id_counter = 0
+
+    def get_render_options(self) -> List[interfaces.renderers.RenderOption]:
+        return []
+
+    def to_arrow_schema(self, grid: interfaces.renderers.TreeGrid) -> "pa.Schema":
+        fields = []
+        for column in grid.columns:
+            arrow_type = self._to_arrow_type[column.type]
+            fields.append(pa.field(column.name, arrow_type()))
+
+        # if the output is nested, e.g., windows.pstree
+        if self._is_tree_result:
+            fields.append(pa.field("_vol_id", pa.uint64()))
+            fields.append(pa.field("_vol_parent_id", pa.uint64()))
+
+        return pa.schema(fields)
+
+    def _flatten_tree_structure(self, nested: List[Dict]) -> List[Dict]:
+        """
+        Flattens a list of nested dicts using the `__children` key.
+
+        Each node gets a `_vol_id` and a `_vol_parent_id` to preserve
+        the original tree structure in a flat format suitable for tabular output.
+
+        Args:
+            nested: A list of dicts with optional `__children` lists (tree nodes).
+
+        Returns:
+            A flat list of dicts with `_vol_id` and `_vol_parent_id`.
+        """
+        rows = []
+        self._node_id_counter = 0
+
+        def _process_node(node: Dict, parent_id: Optional[int]):
+            current_id = self._node_id_counter
+            self._node_id_counter += 1
+
+            entry = {k: v for k, v in node.items() if k != "__children"}
+            entry["_vol_id"] = current_id
+            entry["_vol_parent_id"] = parent_id
+            rows.append(entry)
+
+            for child in node.get("__children", []):
+                _process_node(child, current_id)
+
+        for root in nested:
+            _process_node(root, None)
+
+        return rows
+
+    def output_result(self, schema: "pa.Schema", outfd: TextIO, result):
+        """Outputs the JSON data to a file in a particular format"""
+
+        if self._is_tree_result:
+            result = self._flatten_tree_structure(result)
+
+        t = pa.Table.from_pylist(result, schema=schema)
+        self.write_table(t, outfd)
+
+    def write_table(self, t: "pa.Table", outfd: TextIO) -> None:
+        buf = pa.BufferOutputStream()
+
+        writer = pa.ipc.new_stream(buf, t.schema)
+        writer.write_table(t)
+        writer.close()
+
+        # Get the buffer bytes and write to output
+        buf_bytes = buf.getvalue().to_pybytes()
+        outfd.buffer.write(buf_bytes)
+
+    def render(self, grid: interfaces.renderers.TreeGrid):
+        outfd = sys.stdout
+        final_output: Tuple[
+            Dict[str, List[interfaces.renderers.TreeNode]],
+            List[interfaces.renderers.TreeNode],
+        ] = ({}, [])
+
+        ignore_columns = self.ignored_columns(grid)
+
+        def visitor(
+            node: interfaces.renderers.TreeNode,
+            accumulator: Tuple[Dict[str, Dict[str, Any]], List[Dict[str, Any]]],
+        ) -> Tuple[Dict[str, Dict[str, Any]], List[Dict[str, Any]]]:
+            # Nodes always have a path value, giving them a path_depth of at least 1, we use max just in case
+            acc_map, final_tree = accumulator
+            node_dict: Dict[str, Any] = {"__children": []}
+            line = []
+            for column_index, column in enumerate(grid.columns):
+                if column in ignore_columns:
+                    continue
+
+                data = list(node.values)[column_index]
+
+                if isinstance(data, interfaces.renderers.BaseAbsentValue):
+                    data = None
+
+                if isinstance(data, renderers.Disassembly):
+                    data = text_renderer.display_disassembly(data)
+
+                if isinstance(data, renderers.LayerData):
+                    data = text_renderer.LayerDataRenderer().render_bytes(data)[0]
+
+                node_dict[column.name] = data
+                line.append(data)
+
+            if self.filter and self.filter.filter(line):
+                return accumulator
+
+            if node.parent:
+                acc_map[node.parent.path]["__children"].append(node_dict)
+                self._is_tree_result = True
+            else:
+                final_tree.append(node_dict)
+            acc_map[node.path] = node_dict
+
+            return (acc_map, final_tree)
+
+        if not grid.populated:
+            grid.populate(visitor, final_output)
+        else:
+            grid.visit(node=None, function=visitor, initial_accumulator=final_output)
+
+        schema = self.to_arrow_schema(grid)
+        self.output_result(schema, outfd, final_output[1])
+
+
+class ParquetRenderer(ArrowRenderer):
+    """Renderer that outputs Parquet format data."""
+
+    name = "parquet"
+    structured_output = True
+    _version = (1, 0, 0)
+
+    def get_render_options(self) -> List[interfaces.renderers.RenderOption]:
+        return []
+
+    def write_table(self, table: "pa.Table", outfd: TextIO) -> None:
+        """
+        Writes a table to stdout using the Parquet format.
+
+        Args:
+            t: The Arrow table to write
+            outfd: The output file descriptor
+
+        Returns:
+            Nothing
+        """
+        # Write DataFrame to a temporary file-like object
+        buf = pa.BufferOutputStream()
+        pq.write_table(table, buf, compression="snappy")
+
+        # Get the buffer as a bytes object
+        buf_bytes = buf.getvalue().to_pybytes()
+        outfd.buffer.write(buf_bytes)