Test that handlers tolerate a non-empty unknown chunk prefix/suffix

Author: e3krisztian

tests/test_handlers.py

@@ -7,10 +7,13 @@
 - The expected output in the __output__ folder.
 """
 
+import hashlib
 import inspect
+from collections import Counter
 from pathlib import Path
 from typing import Type
 
+import attr
 import pytest
 
 from unblob import handlers
@@ -39,6 +42,111 @@ def test_all_handlers(
     check_output_is_the_same(output_dir, extraction_config.extract_root)
 
 
+BLOCK_SHIFTING_PREFIX = bytes([0]) + b"unique unknown prefix" + bytes([0])
+PADDING_CHECK_SUFFIX = bytes([0] * 511) + b"unique unknown suffix" + bytes([0])
+
+
+@pytest.mark.parametrize(
+    "input_dir, output_dir", gather_integration_tests(TEST_DATA_PATH)
+)
+@pytest.mark.parametrize(
+    "prefix, suffix",
+    [
+        # pytest.param(b"", b"", id="no-extras"),
+        pytest.param(BLOCK_SHIFTING_PREFIX, b"", id="block-shifted"),
+        pytest.param(b"", PADDING_CHECK_SUFFIX, id="padding-check"),
+        pytest.param(
+            BLOCK_SHIFTING_PREFIX,
+            PADDING_CHECK_SUFFIX,
+            id="block-shifted-padding-check",
+        ),
+    ],
+)
+def test_all_handlers_chunk_stability(
+    input_dir: Path,
+    output_dir: Path,
+    extraction_config: ExtractionConfig,
+    tmp_path: Path,
+    prefix: bytes,
+    suffix: bytes,
+):
+    """Test that handlers tolerate a non-empty unknown chunk prefix/suffix"""
+    altered_input_file = tmp_path / "input_file"
+
+    for input_file in input_dir.iterdir():
+        altered_input_file.write_bytes(prefix + input_file.read_bytes() + suffix)
+
+        config = attr.evolve(
+            extraction_config,
+            extract_root=extraction_config.extract_root / input_file.name,
+        )
+        reports = process_file(config, altered_input_file)
+        check_result(reports)
+
+        check_output_do_not_change_much_due_to_extras(
+            input_file,
+            expected=output_dir / (input_file.name + config.extract_suffix),
+            actual=config.extract_root,
+        )
+
+
+def hash_bytes(data: bytes) -> str:
+    return hashlib.sha512(data).hexdigest()
+
+
+def hash_dir(root: Path, print=lambda *_: None) -> Counter:
+    """Hash all files under an unblob extraction directory :root:, ignoring test prefix/suffix.
+
+    Directory structures having the same set of files will result in the same output,
+    even if the file names are different or the directory structure is different.
+
+    Test prefix/suffix is excluded from hash calculation, so that unknown chunks extended
+    with them will produce the same hash.
+
+    Returns: count of each hashes found
+    """
+    hash_counter = Counter()
+    for path in sorted(root.rglob("*")):
+        if not path.is_file() or path.name == ".gitkeep":
+            continue
+
+        content = path.read_bytes()
+        # ignore newly introduced unknown chunk files
+        if content in (BLOCK_SHIFTING_PREFIX, PADDING_CHECK_SUFFIX):
+            continue
+
+        # remove extras introduced before hashing
+        if content.startswith(BLOCK_SHIFTING_PREFIX):
+            content = content[len(BLOCK_SHIFTING_PREFIX) :]
+        if content.endswith(PADDING_CHECK_SUFFIX):
+            content = content[: -len(PADDING_CHECK_SUFFIX)]
+
+        hash = hash_bytes(content)
+        hash_counter[hash] += 1
+        # Output for debugging failures
+        print(f"    {path} =\n        {hash}")
+    return hash_counter
+
+
+def check_output_do_not_change_much_due_to_extras(
+    original_input_file: Path, expected: Path, actual: Path
+):
+    print(f"{expected=}")
+    expected_counts = hash_dir(expected, print=print)
+    print(f"{actual=}")
+    actual_counts = hash_dir(actual, print=print)
+
+    # original input will show up in the extraction of the modified input due to the extra unknown chunks
+    # but it might not show up in the expected output if it was a "whole file chunk"
+    hash_original_input = hash_bytes(original_input_file.read_bytes())
+    if hash_original_input not in expected_counts:
+        print(f"Warn: hash of original input: {hash_original_input} not in expected")
+        assert actual_counts[hash_original_input] <= 1
+        del actual_counts[hash_original_input]
+
+    assert expected_counts == actual_counts
+
+
 @pytest.mark.parametrize(
     "handler",
     (pytest.param(handler, id=handler.NAME) for handler in handlers.BUILTIN_HANDLERS),

unblob/testing.py

@@ -41,7 +41,7 @@ def gather_integration_tests(test_data_path: Path):
 @pytest.fixture
 def extraction_config(tmp_path: Path):
     config = ExtractionConfig(
-        extract_root=tmp_path,
+        extract_root=tmp_path / "extract_root",
         entropy_depth=0,
         keep_extracted_chunks=True,
     )