Added legends to entropy plots; added module prioritization.

src/binwalk/core/module.py

@@ -148,6 +148,9 @@ class Module(object):
 	# containing a valid file attribute.
 	AUTO_UPDATE_STATUS = True
 
+	# Modules with higher priorities are executed first
+	PRIORITY = 5
+
 	def __init__(self, dependency=False, **kwargs):
 		self.errors = []
 		self.results = []
@@ -339,13 +342,14 @@ class Module(object):
 	def footer(self):
 		self.config.display.footer()
 			
-	def main(self, status):
+	def main(self, parent):
 		'''
 		Responsible for calling self.init, initializing self.config.display, and calling self.run.
 
 		Returns the value returned from self.run.
 		'''
-		self.status = status
+		self.status = parent.status
+		self.modules = parent.loaded_modules
 
 		# Reset all dependency modules
 		for (dependency, module) in iterator(self.DEPENDS):
@@ -447,16 +451,16 @@ class Modules(object):
 
 		@attribute - The desired module attribute.
 
-		Returns a list of modules that contain the specified attribute.
+		Returns a list of modules that contain the specified attribute, in the order they should be executed.
 		'''
 		import binwalk.modules
-		modules = []
+		modules = {}
 
 		for (name, module) in inspect.getmembers(binwalk.modules):
 			if inspect.isclass(module) and hasattr(module, attribute):
-				modules.append(module)
+				modules[module] = module.PRIORITY
 
-		return modules
+		return sorted(modules, key=modules.get, reverse=True)
 
 	def help(self):
 		help_string = "\nBinwalk v%s\nCraig Heffner, http://www.binwalk.org\n" % binwalk.core.settings.Settings.VERSION
@@ -508,7 +512,7 @@ class Modules(object):
 		obj = self.load(module)
 
 		if isinstance(obj, binwalk.core.module.Module) and obj.enabled:
-			obj.main(status=self.status)
+			obj.main(parent=self)
 			self.status.clear()
 
 		# Add object to loaded_modules here, that way if a module has already been

src/binwalk/modules/entropy.py

 import os
 import math
 import binwalk.core.common
+from binwalk.core.compat import *
 from binwalk.core.module import Module, Option, Kwarg
 
 class Entropy(Module):
@@ -14,6 +15,8 @@ class Entropy(Module):
 	FILE_WIDTH = 1024
 	FILE_FORMAT = 'png'
 
+	COLORS = ['r', 'g', 'c', 'b', 'm']
+
 	TITLE = "Entropy"
 
 	CLI = [
@@ -29,31 +32,60 @@ class Entropy(Module):
 				   long='no-plot',
 				   kwargs={'do_plot' : False},
 				   description='Do not generate an entropy plot graph'),
+			Option(short='Q',
+				   long='no-legend',
+				   kwargs={'show_legend' : False},
+				   description='Omit the legend from the entropy plot graph (implied if multiple files are specified)'),
 	]
 
 	KWARGS = [
 			Kwarg(name='enabled', default=False),
 			Kwarg(name='save_plot', default=False),
 			Kwarg(name='do_plot', default=True),
+			Kwarg(name='show_legend', default=True),
 			Kwarg(name='block_size', default=1024),
 	]
 
+	# Run this module last so that it can process all other module's results and overlay them on the entropy graph
+	PRIORITY = 0
+
 	def init(self):
 		self.HEADER[-1] = "ENTROPY"
-
-		self.entropy_results = {}
 		self.algorithm = self.shannon
+		self.display_results = True
+		self.max_description_length = 0
+		self.markers = []
 
+		# Automatically save plots if there is more than one file specified
 		if len(self.config.target_files) > 1:
 			self.save_plot = True
+			self.show_legend = False
+
+		# Get a list of all other module's results to mark on the entropy graph
+		for (module, obj) in iterator(self.modules):
+			for result in obj.results:
+				if result.description:
+					description = result.description.split(',')[0]
+					if len(description) > self.max_description_length:
+						self.max_description_length = len(description)
+					self.markers.append((result.offset, description))
+
+		# If other modules have been run and they produced results, don't spam the terminal with entropy results
+		if self.markers:
+			self.display_results = False
 
 		if self.config.block:
 			self.block_size = self.config.block
 
 	def run(self):
 		for fp in self.config.target_files:
+
+			if self.display_results:
 				self.header()
+
 			self.calculate_file_entropy(fp)
+
+			if self.display_results:
 				self.footer()
 
 	def calculate_file_entropy(self, fp):
@@ -70,7 +102,7 @@ class Entropy(Module):
 			i = 0
 			while i < dlen:
 				entropy = self.algorithm(data[i:i+self.block_size])
-				r = self.result(offset=(file_offset + i), file=fp, entropy=entropy, description=("%f" % entropy))
+				r = self.result(offset=(file_offset + i), file=fp, entropy=entropy, description=("%f" % entropy), display=self.display_results)
 				i += self.block_size
 
 		if self.do_plot:
@@ -108,8 +140,10 @@ class Entropy(Module):
 		import pyqtgraph as pg
 		from pyqtgraph.Qt import QtCore, QtGui
 
+		i = 0
 		x = []
 		y = []
+		plotted_colors = {}
 
 		for r in self.results:
 			x.append(r.offset)
@@ -117,7 +151,29 @@ class Entropy(Module):
 
 		plt = pg.plot(title=fname, clear=True)
 
-		plt.plot(x, y, pen='y') #pen='b'
+		plt.plot(x, y, pen='y')
+
+		if self.show_legend and self.markers:
+			plt.addLegend(size=(self.max_description_length*10, 0))
+
+			for (offset, description) in self.markers:
+				# If this description has already been plotted at a different offset, we need to 
+				# use the same color for the marker, but set the description to None to prevent
+				# duplicate entries in the graph legend.
+				#
+				# Else, get the next color and use it to mark descriptions of this type.
+				if has_key(plotted_colors, description):
+					color = plotted_colors[description]
+					description = None
+				else:
+					color = self.COLORS[i]
+					plotted_colors[description] = color
+
+					i += 1
+					if i >= len(self.COLORS):
+						i = 0
+
+				plt.plot(x=[offset,offset], y=[0,1.1], name=description, pen=pg.mkPen(color, width=2.5))
 
 		if self.save_plot:
 			exporter = pg.exporters.ImageExporter.ImageExporter(plt.plotItem)