mirror of
https://github.com/volatilityfoundation/volatility3.git
synced 2026-09-13 21:27:39 +02:00
Ensure the arrow/parquet classes aren't created unless the libraries are there
This commit is contained in:
Mike Auty
@@ -9,12 +9,13 @@ from typing import (
|
||||
Dict,
|
||||
List,
|
||||
Optional,
|
||||
Tuple,
|
||||
TextIO,
|
||||
Tuple,
|
||||
)
|
||||
|
||||
from volatility3.cli import text_renderer
|
||||
from volatility3.framework import interfaces, renderers
|
||||
from volatility3.framework.renderers import format_hints
|
||||
from volatility3.cli import text_renderer
|
||||
|
||||
vollog = logging.getLogger(__name__)
|
||||
|
||||
@@ -27,192 +28,191 @@ try:
|
||||
except ImportError:
|
||||
vollog.debug("Arrow/Parquet libraries not found")
|
||||
|
||||
if ARROW_PRESENT:
|
||||
|
||||
class ArrowRenderer(text_renderer.CLIRenderer):
|
||||
"""Renderer that outputs Arrow IPC format data."""
|
||||
class ArrowRenderer(text_renderer.CLIRenderer):
|
||||
"""Renderer that outputs Arrow IPC format data."""
|
||||
|
||||
name = "arrow"
|
||||
structured_output = True
|
||||
_version = (1, 0, 0)
|
||||
name = "arrow"
|
||||
structured_output = True
|
||||
_version = (1, 0, 0)
|
||||
|
||||
def __init__(
|
||||
self, options: Optional[List[interfaces.renderers.RenderOption]] = None
|
||||
) -> None:
|
||||
super().__init__(options)
|
||||
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.binary,
|
||||
format_hints.HexBytes: pa.binary,
|
||||
renderers.LayerData: pa.binary,
|
||||
bytes: pa.binary,
|
||||
}
|
||||
|
||||
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.binary,
|
||||
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
|
||||
|
||||
# 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 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()))
|
||||
|
||||
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()))
|
||||
|
||||
# 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)
|
||||
|
||||
return pa.schema(fields)
|
||||
def _flatten_tree_structure(self, nested: List[Dict]) -> List[Dict]:
|
||||
"""
|
||||
Flattens a list of nested dicts using the `__children` key.
|
||||
|
||||
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.
|
||||
|
||||
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).
|
||||
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
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)
|
||||
|
||||
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 child in node.get("__children", []):
|
||||
_process_node(child, current_id)
|
||||
for root in nested:
|
||||
_process_node(root, None)
|
||||
|
||||
for root in nested:
|
||||
_process_node(root, None)
|
||||
return rows
|
||||
|
||||
return rows
|
||||
def output_result(self, schema: "pa.Schema", outfd: TextIO, result):
|
||||
"""Outputs the JSON data to a file in a particular format"""
|
||||
|
||||
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)
|
||||
|
||||
if self._is_tree_result:
|
||||
result = self._flatten_tree_structure(result)
|
||||
t = pa.Table.from_pylist(result, schema=schema)
|
||||
self.write_table(t, outfd)
|
||||
|
||||
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()
|
||||
|
||||
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()
|
||||
|
||||
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)
|
||||
|
||||
# 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],
|
||||
] = ({}, [])
|
||||
|
||||
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)
|
||||
|
||||
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
|
||||
|
||||
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]
|
||||
|
||||
data = list(node.values)[column_index]
|
||||
if isinstance(data, interfaces.renderers.BaseAbsentValue):
|
||||
data = None
|
||||
|
||||
if isinstance(data, interfaces.renderers.BaseAbsentValue):
|
||||
data = None
|
||||
if isinstance(data, renderers.Disassembly):
|
||||
data = text_renderer.display_disassembly(data)
|
||||
|
||||
if isinstance(data, renderers.Disassembly):
|
||||
data = text_renderer.display_disassembly(data)
|
||||
if isinstance(data, renderers.LayerData):
|
||||
data = text_renderer.LayerDataRenderer().render_bytes(data)[0]
|
||||
|
||||
if isinstance(data, renderers.LayerData):
|
||||
data = text_renderer.LayerDataRenderer().render_bytes(data)[0]
|
||||
node_dict[column.name] = data
|
||||
line.append(data)
|
||||
|
||||
node_dict[column.name] = data
|
||||
line.append(data)
|
||||
if self.filter and self.filter.filter(line):
|
||||
return accumulator
|
||||
|
||||
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
|
||||
|
||||
if node.parent:
|
||||
acc_map[node.parent.path]["__children"].append(node_dict)
|
||||
self._is_tree_result = True
|
||||
return (acc_map, final_tree)
|
||||
|
||||
if not grid.populated:
|
||||
grid.populate(visitor, final_output)
|
||||
else:
|
||||
final_tree.append(node_dict)
|
||||
acc_map[node.path] = node_dict
|
||||
grid.visit(
|
||||
node=None, function=visitor, initial_accumulator=final_output
|
||||
)
|
||||
|
||||
return (acc_map, final_tree)
|
||||
schema = self.to_arrow_schema(grid)
|
||||
self.output_result(schema, outfd, final_output[1])
|
||||
|
||||
if not grid.populated:
|
||||
grid.populate(visitor, final_output)
|
||||
else:
|
||||
grid.visit(node=None, function=visitor, initial_accumulator=final_output)
|
||||
class ParquetRenderer(ArrowRenderer):
|
||||
"""Renderer that outputs Parquet format data."""
|
||||
|
||||
schema = self.to_arrow_schema(grid)
|
||||
self.output_result(schema, outfd, final_output[1])
|
||||
name = "parquet"
|
||||
structured_output = True
|
||||
_version = (1, 0, 0)
|
||||
|
||||
def get_render_options(self) -> List[interfaces.renderers.RenderOption]:
|
||||
return []
|
||||
|
||||
class ParquetRenderer(ArrowRenderer):
|
||||
"""Renderer that outputs Parquet format data."""
|
||||
def write_table(self, table: "pa.Table", outfd: TextIO) -> None:
|
||||
"""
|
||||
Writes a table to stdout using the Parquet format.
|
||||
|
||||
name = "parquet"
|
||||
structured_output = True
|
||||
_version = (1, 0, 0)
|
||||
Args:
|
||||
t: The Arrow table to write
|
||||
outfd: The output file descriptor
|
||||
|
||||
def get_render_options(self) -> List[interfaces.renderers.RenderOption]:
|
||||
return []
|
||||
Returns:
|
||||
Nothing
|
||||
"""
|
||||
# Write DataFrame to a temporary file-like object
|
||||
buf = pa.BufferOutputStream()
|
||||
pq.write_table(table, buf, compression="snappy")
|
||||
|
||||
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)
|
||||
# Get the buffer as a bytes object
|
||||
buf_bytes = buf.getvalue().to_pybytes()
|
||||
outfd.buffer.write(buf_bytes)
|
||||
|
||||
Reference in New Issue
Block a user