enhance and extend testing helper functions

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
 #