diff --git a/test/test_volatility.py b/test/test_volatility.py index 632a6bcde..960ac00d1 100644 --- a/test/test_volatility.py +++ b/test/test_volatility.py @@ -11,8 +11,15 @@ import sys import tempfile import contextlib import functools +import json +import logging from typing import List, Tuple +from test import WINDOWS_TESTS_DATA_DIR + +test_logger = logging.getLogger(__name__) + + # # HELPER FUNCTIONS # @@ -71,25 +78,100 @@ def runvolshell(img, volshell, python, volshellargs=None, globalargs=None): return runvol(args, volshell, python) +def load_test_data(plugin: str, test_key: str): + if plugin.startswith("windows."): + data_path = WINDOWS_TESTS_DATA_DIR / f"{plugin}.json" + # TODO: add Linux and macOS when any of these requires this API + else: + raise Exception(f"Cannot determine OS of plugin: {plugin}") + + if not data_path.exists(): + raise FileNotFoundError( + f"Test data not found for plugin {plugin} at {data_path}" + ) + + with open(data_path) as f: + # This will raise an explicit exception by itself on failures + return json.load(f)[test_key] + + +def dict_lower_strvalues(dict_to_convert: dict): + """Lower each value of type string of a dictionary + + Args: + dict_to_convert: The dictionary in which to lower the string values + Returns: + A copy of the dictionary with lowered string values + """ + + converted = {} + for key, value in dict_to_convert.items(): + if isinstance(value, str): + converted[key] = value.lower() + else: + converted[key] = value + return converted + + def match_output_row( - expected_row: dict, plugin_json_out: List[dict], exact_match: bool = False + expected_row: dict, + plugin_json_out: List[dict], + exact_match: bool = False, + case_sensitive: bool = True, + children_recursive: bool = False, ): - """Search each row of a plugin's JSON output for an expected row. Each row is a dict. + """Search each row in a plugin's JSON output for a matching row. + This method supports recursive comparisons using the "__children" key, making it useful for testing hierarchical plugins like windows.pstree. + It also maintains case sensitivity and exact matching behavior when traversing nested structures. Args: expected_row: The expected row to be found in the output plugin_json_out: The plugin's output in JSON format (typically obtained through -r json and json.loads) - exact_match: Whether to require exactly the expected row, no more no less, or to anticipate columns' addition by checking only + exact_match: Require exactly the expected row, no more no less, or anticipate columns' addition by checking only the expected row keys and values + case_sensitive: Operate case sensitive match for str values of both dictionaries or not + children_recursive: Perform a recursive match by inspecting "__children" keys of each expected_row + + Returns: + A boolean indicating whether a match was found or not """ + # Lower each string value of both dicts + if not case_sensitive: + expected_row = dict_lower_strvalues(expected_row) + plugin_json_out_tmp = [] + for row in plugin_json_out: + plugin_json_out_tmp.append(dict_lower_strvalues(row)) + plugin_json_out = plugin_json_out_tmp + if not exact_match: for row in plugin_json_out: if all( - expected_item in row.items() for expected_item in expected_row.items() + expected_item in row.items() + for expected_item in expected_row.items() + if not expected_item[0] == "__children" ): - return True + if ( + children_recursive + and "__children" in expected_row + and "__children" in row + ): + for children_expected_row in expected_row["__children"]: + if not match_output_row( + children_expected_row, + row["__children"], + case_sensitive=case_sensitive, + children_recursive=True, + ): + break + else: + # We matched all the children keys + return True + else: + # No recursion required and we already matched the row + return True else: + # No "__children" recursion here as we want to match the whole tree at once for row in plugin_json_out: if expected_row == row: return True @@ -97,6 +179,18 @@ def match_output_row( return False +def count_entries_flat(plugin_json_out: List[dict]): + """Count the number of entries as if -r json wasn't specified. Allows to get a non-hierarchical count, without running a plugin twice + (once with "-r json" and once without) while still preserving JSON features. + + Args: + plugin_json_out: The plugin's output in JSON format (typically obtained through -r json and json.loads) + """ + # Remove whitespaces between entries + # If a value contains {", it will be represented by {\" so no confusion + return json.dumps(plugin_json_out, separators=(",", ":")).count('{"') + + # # TESTS #