add whole-file entropy calculation to triage

1 files changed, 23 insertions, 2 deletions

diff --git a/python/examples/triage/entropy.py b/python/examples/triage/entropy.py
index a2ab41a6..8dbd3900 100644
--- a/python/examples/triage/entropy.py
+++ b/python/examples/triage/entropy.py
@@ -2,7 +2,7 @@ import math
 import threading
 from PySide6.QtWidgets import QWidget
 from PySide6.QtGui import QImage, QColor, QPainter
-from PySide6.QtCore import Qt, QSize, QTimer
+from PySide6.QtCore import Qt, QSize, QTimer, Signal
 import binaryninjaui
 from binaryninjaui import ViewFrame, UIContext
 from binaryninja.enums import ThemeColor
@@ -14,11 +14,22 @@ class EntropyThread(threading.Thread):
 		self.image = image
 		self.block_size = block_size
 		self.updated = False
+		self.average_entropy = 0.0
+		self.sample_count = 0
+		self.lock = threading.Lock()
 
 	def run(self):
 		width = self.image.width()
 		for i in range(0, width):
-			v = int(self.data.get_entropy(self.data.start + i * self.block_size, self.block_size)[0] * 255)
+			entropy_result = self.data.get_entropy(self.data.start + i * self.block_size, self.block_size)
+			entropy_value = entropy_result[0] if entropy_result else 0.0
+			v = int(entropy_value * 255)
+
+			# Update running average
+			with self.lock:
+				self.sample_count += 1
+				self.average_entropy += (entropy_value - self.average_entropy) / self.sample_count
+
 			if v >= 240:
 				color = binaryninjaui.getThemeColor(ThemeColor.YellowStandardHighlightColor)
 				self.image.setPixelColor(i, 0, color)
@@ -29,8 +40,14 @@ class EntropyThread(threading.Thread):
 				self.image.setPixelColor(i, 0, color)
 			self.updated = True
 
+	def get_average_entropy(self):
+		with self.lock:
+			return self.average_entropy
+
 
 class EntropyWidget(QWidget):
+	entropyUpdated = Signal(float)
+
 	def __init__(self, parent, view, data):
 		super(EntropyWidget, self).__init__(parent)
 		self.view = view
@@ -72,6 +89,7 @@ class EntropyWidget(QWidget):
 		if self.thread.updated:
 			self.thread.updated = False
 			self.update()
+			self.entropyUpdated.emit(self.thread.get_average_entropy())
 
 	def mousePressEvent(self, event):
 		if event.button() != Qt.LeftButton:
@@ -88,3 +106,6 @@ class EntropyWidget(QWidget):
 			self.setToolTip(f"0x{addr:x}")
 		else:
 			self.setToolTip(f"File offset: 0x{offset:x}")
+
+	def get_average_entropy(self):
+		return self.thread.get_average_entropy()