Fixed bugs in entropy legend

src/binwalk/__init__.py

@@ -304,6 +304,7 @@ class Binwalk(object):
 			target_files = [target_files]
 
 		Plotter3D(target_files, offset=offset, length=length, weight=weight, verbose=verbose).plot()
+		#PlotFiles(target_files, offset=offset, length=length, weight=weight, verbose=verbose)
 
 	def scan(self, target_files, offset=0, length=0, show_invalid_results=False, callback=None, start_callback=None, end_callback=None, base_dir=None, matryoshka=1, plugins_whitelist=[], plugins_blacklist=[]):
 		'''

src/binwalk/entropy.py

@@ -42,6 +42,7 @@ class PlotEntropy(object):
 
 		i = 0
 		descriptions = {}
+		plotted_colors = {}
 		max_description_length = None
 
 		for (offset, results) in file_results:
@@ -70,14 +71,31 @@ class PlotEntropy(object):
 		#if average:
 		#	plt.addLine(y=average, pen='r')
 
-		if file_results:
-			for (offset, descs) in iterator(descriptions):
-				for description in descs:
-					plt.plot(x=[offset,offset], y=[0,1.1], name=description, pen=pg.mkPen(self.COLORS[i], width=2.5))
+		if descriptions:
+			ordered_offsets = get_keys(descriptions)
+			ordered_offsets.sort()
+
+			for offset in ordered_offsets:
+				for description in descriptions[offset]:
+
+					# 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 save:
 			exporter = pg.exporters.ImageExporter.ImageExporter(plt.plotItem)
 			exporter.parameters()['width'] = self.FILE_WIDTH

src/binwalk/plotter.py

@@ -131,9 +131,7 @@ class Plotter(object):
 
 		return scatter_plot
 
-	def plot(self):
-		from pyqtgraph.Qt import QtCore, QtGui
-
+	def plot(self, wait=True):
 		self.window.show()
 
 		for file_name in self.files:
@@ -147,6 +145,12 @@ class Plotter(object):
 
 			self.window.addItem(self._generate_plot(plot_points, data_weights))
 
+		if wait:
+			self.wait()
+
+	def wait(self):
+		from pyqtgraph.Qt import QtCore, QtGui
+
 		t = QtCore.QTimer()
 		t.start(50)
 		QtGui.QApplication.instance().exec_()
@@ -179,6 +183,23 @@ class Plotter2D(Plotter):
 		elif self.plane_count == 2:
 			return (ord(data[0]), ord(data[1]), 0)
 		
+class PlotFiles(object):
+
+	def __init__(self, files, offset=0, length=0, weight=None, verbose=False, overlay=False):
+		
+		if overlay:
+			Plotter3D(files, offset=offset, length=length, weight=weight, verbose=verbose).plot(wait=True)
+		else:
+			objs = []
+
+			for f in files:
+				p = Plotter3D(files, offset=offset, length=length, weight=weight, verbose=verbose)
+				p.plot(wait=False)
+				objs.append(p)
+
+			for obj in objs:
+				obj.wait()
+
 
 
 if __name__ == '__main__':