feat(math): introduce Chi Square entropy in EntropyReport

Introduce another entropy measure based on Chi Square. This entropy
measure is introduced by modifying the EntropyReport class so that it
contains two EntropyMeasures:
- shannon: for Shannon entropy, which was already there
- chi_square: for Chi Square entropy, which we introduce

The format_entropy_plot has been adjusted to display two graphs. One for
Shannon, the other for Chi Square.

Author: qkaiser

tests/test_processing.py

@@ -2,6 +2,7 @@
 import sys
 import zipfile
 from pathlib import Path
+from statistics import mean
 from typing import Collection, List, Optional, Tuple, Type, TypeVar
 
 import attr
@@ -33,6 +34,7 @@
 )
 from unblob.report import (
     ChunkReport,
+    EntropyMeasurements,
     EntropyReport,
     ExtractDirectoryExistsReport,
     FileMagicReport,
@@ -199,7 +201,15 @@ def test_calculate_block_size(
 
 def test_format_entropy_plot_error():
     with pytest.raises(TypeError):
-        format_entropy_plot(percentages=[], block_size=1024)
+        format_entropy_plot(
+            EntropyReport(
+                shannon=EntropyMeasurements(percentages=[], block_size=1024),
+                chi_square=EntropyMeasurements(
+                    percentages=[],
+                    block_size=1024,
+                ),
+            )
+        )
 
 
 @pytest.mark.parametrize(
@@ -215,17 +225,20 @@ def test_format_entropy_plot_error():
 )
 def test_format_entropy_plot_no_exception(percentages: List[float], block_size: int):
     assert str(block_size) in format_entropy_plot(
-        percentages=percentages,
-        block_size=block_size,
+        EntropyReport(
+            shannon=EntropyMeasurements(
+                percentages=percentages, block_size=block_size, mean=mean(percentages)
+            ),
+            chi_square=EntropyMeasurements(
+                percentages=percentages, block_size=block_size, mean=mean(percentages)
+            ),
+        )
     )
 
 
 def test_calculate_entropy_no_exception():
     report = calculate_entropy(Path(sys.executable))
-    format_entropy_plot(
-        percentages=report.percentages,
-        block_size=report.block_size,
-    )
+    format_entropy_plot(report)
 
 
 @pytest.mark.parametrize(
@@ -434,17 +447,24 @@ def get_all(file_name, report_type: Type[ReportType]) -> List[ReportType]:
     assert (
         unknown_entropy is not None
     )  # removes pyright complaints for the below lines :(
-    assert unknown_entropy.percentages == [0.0, 75.0] + [100.0] * 62
-    assert unknown_entropy.block_size == 1024
-    assert round(unknown_entropy.mean, 2) == 98.05  # noqa: PLR2004
-    assert unknown_entropy.highest == 100.0  # noqa: PLR2004
-    assert unknown_entropy.lowest == 0.0
+    assert unknown_entropy.shannon.percentages == [0.0, 75.0] + [100.0] * 62
+    assert unknown_entropy.shannon.block_size == 1024
+    assert round(unknown_entropy.shannon.mean, 2) == 98.05  # noqa: PLR2004
+    assert unknown_entropy.shannon.highest == 100.0  # noqa: PLR2004
+    assert unknown_entropy.shannon.lowest == 0.0
 
     # we should have entropy calculated for files without extractions, except for empty files
     assert get_all("empty.txt", EntropyReport) == []
-    assert [EntropyReport(percentages=[100.0], block_size=1024, mean=100.0)] == get_all(
-        "0-255.bin", EntropyReport
-    )
+    assert [
+        EntropyReport(
+            shannon=EntropyMeasurements(
+                percentages=[100.0], block_size=1024, mean=100.0
+            ),
+            chi_square=EntropyMeasurements(
+                percentages=[0.0], block_size=1024, mean=0.0
+            ),
+        )
+    ] == get_all("0-255.bin", EntropyReport)
 
 
 @pytest.mark.parametrize(

unblob/processing.py

@@ -2,6 +2,7 @@
 import shutil
 from operator import attrgetter
 from pathlib import Path
+from statistics import mean
 from typing import Iterable, List, Optional, Sequence, Set, Tuple, Type, Union
 
 import attr
@@ -34,6 +35,7 @@
 from .pool import make_pool
 from .report import (
     CalculateMultiFileExceptionReport,
+    EntropyMeasurements,
     EntropyReport,
     ExtractDirectoryExistsReport,
     FileMagicReport,
@@ -563,8 +565,7 @@ def _calculate_entropy(self, path: Path) -> Optional[EntropyReport]:
                 logger.debug(
                     "Entropy chart",
                     # New line so that chart title will be aligned correctly in the next line
-                    chart="\n"
-                    + format_entropy_plot(report.percentages, report.block_size),
+                    chart="\n" + format_entropy_plot(report),
                     path=path,
                     _verbosity=3,
                 )
@@ -709,7 +710,8 @@ def calculate_entropy(path: Path) -> EntropyReport:
     can contain 0-8 bits of entropy. We normalize it for visualization to a
     0-100% scale, to make it easier to interpret the graph.
     """
-    percentages = []
+    shannon_percentages = []
+    chi_square_percentages = []
 
     # We could use the chunk size instead of another syscall,
     # but we rely on the actual file size written to the disk
@@ -725,31 +727,52 @@ def calculate_entropy(path: Path) -> EntropyReport:
         max_limit=1024 * 1024,
     )
 
-    entropy_sum = 0.0
     with File.from_path(path) as file:
         for chunk in iterate_file(file, 0, file_size, buffer_size=block_size):
-            entropy = mt.shannon_entropy(chunk)
-            entropy_percentage = round(entropy / 8 * 100, 2)
-            percentages.append(entropy_percentage)
-            entropy_sum += entropy * len(chunk)
-
-    report = EntropyReport(
-        percentages=percentages,
-        block_size=block_size,
-        mean=entropy_sum / file_size / 8 * 100,
+            shannon_entropy = mt.shannon_entropy(chunk)
+            shannon_entropy_percentage = round(shannon_entropy / 8 * 100, 2)
+            shannon_percentages.append(shannon_entropy_percentage)
+
+            chi_square_entropy = mt.chi_square(chunk)
+            chi_square_max = 256 * (len(chunk) - 1)
+            chisquare_entropy_percentage = round(
+                (chi_square_entropy / chi_square_max) * 100, 2
+            )
+            chi_square_percentages.append(chisquare_entropy_percentage)
+
+    entropy_report = EntropyReport(
+        shannon=EntropyMeasurements(
+            percentages=shannon_percentages,
+            block_size=block_size,
+            mean=mean(shannon_percentages),
+        ),
+        chi_square=EntropyMeasurements(
+            percentages=chi_square_percentages,
+            block_size=block_size,
+            mean=mean(chi_square_percentages),
+        ),
     )
 
     logger.debug(
-        "Entropy calculated",
+        "Shannon entropy calculated",
+        path=path,
+        size=file_size,
+        block_size=entropy_report.shannon.block_size,
+        mean=round(entropy_report.shannon.mean, 2),
+        highest=round(entropy_report.shannon.highest, 2),
+        lowest=round(entropy_report.shannon.lowest, 2),
+    )
+    logger.debug(
+        "Chi square entropy calculated",
         path=path,
         size=file_size,
-        block_size=report.block_size,
-        mean=round(report.mean, 2),
-        highest=round(report.highest, 2),
-        lowest=round(report.lowest, 2),
+        block_size=entropy_report.chi_square.block_size,
+        mean=round(entropy_report.chi_square.mean, 2),
+        highest=round(entropy_report.chi_square.highest, 2),
+        lowest=round(entropy_report.chi_square.lowest, 2),
     )
 
-    return report
+    return entropy_report
 
 
 def calculate_block_size(
@@ -763,23 +786,39 @@ def calculate_block_size(
     return block_size  # noqa: RET504
 
 
-def format_entropy_plot(percentages: List[float], block_size: int):
+def format_entropy_plot(report: EntropyReport):
+    # start from scratch
+    plt.clear_figure()
+    # go colorless
+    plt.clear_color()
+    plt.title("Shannon Entropy distribution")
+    plt.xlabel(f"{report.shannon.block_size} bytes")
+
+    plt.plot(report.shannon.percentages, label="Shannon (%)")
+    # 16 height leaves no gaps between the lines
+    plt.plot_size(100, 16)
+
+    # Draw ticks every 1Mb on the x axis.
+    plt.xticks(range(len(report.shannon.percentages) + 1))
+    # zoom into the area where our data lives
+    plt.ylim(report.shannon.lowest, report.shannon.highest + 1)
+
+    shannon_plot = plt.build()
+
     # start from scratch
     plt.clear_figure()
     # go colorless
     plt.clear_color()
-    plt.title("Entropy distribution")
-    # plt.xlabel(humanize.naturalsize(block_size))
-    plt.xlabel(f"{block_size} bytes")
-    plt.ylabel("entropy %")
+    plt.title("Chi Square Entropy distribution")
+    plt.xlabel(f"{report.chi_square.block_size} bytes")
 
-    plt.scatter(percentages, marker="dot")
+    plt.plot(report.chi_square.percentages, label="Chi square (%)")
     # 16 height leaves no gaps between the lines
     plt.plot_size(100, 16)
-    plt.ylim(0, 100)
     # Draw ticks every 1Mb on the x axis.
-    plt.xticks(range(len(percentages) + 1))
-    # Always show 0% and 100%
-    plt.yticks(range(0, 101, 10))
+    plt.xticks(range(len(report.chi_square.percentages) + 1))
+    # zoom into the area where our data lives
+    plt.ylim(report.chi_square.lowest, report.chi_square.highest + 1)
 
-    return plt.build()
+    chi_square_plot = plt.build()
+    return "\n" + shannon_plot + "\n" + chi_square_plot

unblob/report.py

@@ -191,7 +191,7 @@ class FileMagicReport(Report):
 
 
 @attr.define(kw_only=True, frozen=True)
-class EntropyReport(Report):
+class EntropyMeasurements:
     percentages: List[float]
     block_size: int
     mean: float
@@ -205,6 +205,12 @@ def lowest(self):
         return min(self.percentages)
 
 
+@attr.define(kw_only=True, frozen=True)
+class EntropyReport(Report):
+    shannon: EntropyMeasurements
+    chi_square: EntropyMeasurements
+
+
 @final
 @attr.define(kw_only=True, frozen=True)
 class ChunkReport(Report):