Commit b771ae0d by devttys0

Added --no-plot and --save-plot to entropy.py.

parent 4ae41878
import os
import math
import binwalk.core.common
from binwalk.core.module import Module, Option, Kwarg
class Entropy(Module):
......@@ -9,6 +11,9 @@ class Entropy(Module):
XUNITS = 'B'
YUNITS = 'E'
FILE_WIDTH = 1024
FILE_FORMAT = 'png'
TITLE = "Entropy"
CLI = [
......@@ -16,16 +21,32 @@ class Entropy(Module):
long='entropy',
kwargs={'enabled' : True},
description='Calculate file entropy'),
Option(short='J',
long='save-plot',
kwargs={'save_plot' : True},
description='Save plot as a PNG (implied if multiple files are specified)'),
Option(short='N',
long='no-plot',
kwargs={'do_plot' : False},
description='Do not generate an entropy plot graph'),
]
KWARGS = [
Kwarg(name='enabled', default=False),
Kwarg(name='save_plot', default=False),
Kwarg(name='do_plot', default=True),
Kwarg(name='block_size', default=1024),
]
def init(self):
self.HEADER[-1] = "ENTROPY"
self.entropy_results = {}
self.algorithm = self.shannon
if len(self.config.target_files) > 1:
self.save_plot = True
if self.config.block:
self.block_size = self.config.block
......@@ -52,6 +73,7 @@ class Entropy(Module):
r = self.result(offset=(file_offset + i), file=fp, entropy=entropy, description=("%f" % entropy))
i += self.block_size
if self.do_plot:
self.plot_entropy(fp.name)
def shannon(self, data):
......@@ -81,7 +103,7 @@ class Entropy(Module):
return e
def plot_entropy(self, title):
def plot_entropy(self, fname):
import numpy as np
import pyqtgraph as pg
from pyqtgraph.Qt import QtCore, QtGui
......@@ -93,10 +115,15 @@ class Entropy(Module):
x.append(r.offset)
y.append(r.entropy)
plt = pg.plot(title=title, clear=True)
plt = pg.plot(title=fname, clear=True)
plt.plot(x, y, pen='y') #pen='b'
if self.save_plot:
exporter = pg.exporters.ImageExporter.ImageExporter(plt.plotItem)
exporter.parameters()['width'] = self.FILE_WIDTH
exporter.export(binwalk.core.common.unique_file_name(os.path.basename(fname), self.FILE_FORMAT))
else:
plt.setLabel('left', self.YLABEL, units=self.YUNITS)
plt.setLabel('bottom', self.XLABEL, units=self.XUNITS)
QtGui.QApplication.instance().exec_()
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment