Fixed multiple status server bugs.

src/binwalk/core/module.py

@@ -6,6 +6,7 @@
 import io
 import os
 import sys
+import time
 import inspect
 import argparse
 import traceback
@@ -610,7 +611,7 @@ class Modules(object):
         self.arguments = []
         self.executed_modules = {}
         self.default_dependency_modules = {}
-        self.status = Status(completed=0, total=0, fp=None)
+        self.status = Status(completed=0, total=0, fp=None, running=False, shutdown=False, finished=False)
         self.status_server_started = False
         self.status_service = None
 
@@ -741,21 +742,29 @@ class Modules(object):
         '''
         Runs a specific module.
         '''
-        obj = self.load(module, kwargs)
+        try:
+            obj = self.load(module, kwargs)
 
-        if isinstance(obj, binwalk.core.module.Module) and obj.enabled:
-            obj.main()
-            self.status.clear()
+            if isinstance(obj, binwalk.core.module.Module) and obj.enabled:
+                obj.main()
+                self.status.clear()
 
-        # If the module is not being loaded as a dependency, add it to the executed modules dictionary.
-        # This is used later in self.execute to determine which objects should be returned.
-        if not dependency:
-            self.executed_modules[module] = obj
+            # If the module is not being loaded as a dependency, add it to the executed modules dictionary.
+            # This is used later in self.execute to determine which objects should be returned.
+            if not dependency:
+                self.executed_modules[module] = obj
 
-            # The unload method tells the module that we're done with it, and gives it a chance to do
-            # any cleanup operations that may be necessary. We still retain the object instance in self.executed_modules.
-            obj._unload_dependencies()
-            obj.unload()
+                # The unload method tells the module that we're done with it, and gives it a chance to do
+                # any cleanup operations that may be necessary. We still retain the object instance in self.executed_modules.
+                obj._unload_dependencies()
+                obj.unload()
+        except KeyboardInterrupt as e:
+            # Tell the status server to shut down, and give it time to clean up.
+            if self.status.running:
+                self.status.shutdown = True
+                while not self.status.finished:
+                    time.sleep(0.1)
+            raise e
 
         return obj
 

src/binwalk/core/statuserver.py

@@ -3,6 +3,7 @@
 
 import sys
 import time
+import errno
 import threading
 import binwalk.core.compat
 
@@ -18,6 +19,9 @@ class StatusRequestHandler(SocketServer.BaseRequestHandler):
         message_format = "%s     %3d%%     [ %d / %d ]"
         last_status_message_len = 0
         status_message = ''
+        message_sent = False
+
+        self.server.binwalk.status.running = True
 
         while True:
             time.sleep(0.1)
@@ -27,20 +31,33 @@ class StatusRequestHandler(SocketServer.BaseRequestHandler):
                 self.request.send(binwalk.core.compat.str2bytes(' ' * last_status_message_len))
                 self.request.send(binwalk.core.compat.str2bytes('\b' * last_status_message_len))
 
-                percentage = ((float(self.server.binwalk.status.completed) / float(self.server.binwalk.status.total)) * 100)
-                status_message = message_format % (self.server.binwalk.status.fp.path,
-                                                   percentage,
-                                                   self.server.binwalk.status.completed,
-                                                   self.server.binwalk.status.total)
-                last_status_message_len = len(status_message)
+                if self.server.binwalk.status.shutdown:
+                    self.server.binwalk.status.finished = True
+                    break
+
+                if self.server.binwalk.status.total != 0:
+                    percentage = ((float(self.server.binwalk.status.completed) / float(self.server.binwalk.status.total)) * 100)
+                    status_message = message_format % (self.server.binwalk.status.fp.path,
+                                                       percentage,
+                                                       self.server.binwalk.status.completed,
+                                                       self.server.binwalk.status.total)
+                elif not message_sent:
+                    status_message = "No status information available at this time!"
+                else:
+                    continue
 
+                last_status_message_len = len(status_message)
                 self.request.send(binwalk.core.compat.str2bytes(status_message))
+                message_sent = True
+            except IOError as e:
+                if e.errno == errno.EPIPE:
+                    break
+            except Exception as e:
+                binwalk.core.common.debug('StatusRequestHandler exception: ' + str(e) + '\n')
             except KeyboardInterrupt as e:
                 raise e
-            except Exception as e:
-                #sys.stderr.write(str(e) + "\n")
-                pass
 
+        self.server.binwalk.status.running = False
         return
 
 class ThreadedStatusServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer):

src/binwalk/modules/entropy.py

@@ -117,7 +117,10 @@ class Entropy(Module):
         # Need to invoke the pyqtgraph stuff via a separate process, as calling pg.exit
         # is pretty much required. pg.exit calls os._exit though, and we don't want to
         # exit out of the main process (especially if being run via the API).
-        if not binwalk.core.common.MSWindows():
+        #
+        # TODO: This is a nasty hack. Find a better way. Until then, just know that
+        #       pg.exit will quite the current process.
+        if False: #not binwalk.core.common.MSWindows():
             p = multiprocessing.Process(target=self._run)
             p.start()
             p.join()
@@ -135,7 +138,7 @@ class Entropy(Module):
             try:
                 import pyqtgraph as pg
             except ImportError as e:
-                binwalk.core.common.warning("pyqtgraph not found, visual entropy graphing will be disabled")
+                binwalk.core.common.warning("Failed to import pyqtgraph module, visual entropy graphing will be disabled")
                 self.do_plot = False
 
         for fp in iter(self.next_file, None):

src/binwalk/modules/hexdiff.py

@@ -118,6 +118,13 @@ class HexDiff(Module):
         loop_count = 0
         sep_count = 0
 
+        # Figure out the maximum diff size (largest file size)
+        self.status.total = 0
+        for i in range(0, len(target_files)):
+            if target_files[i].size > self.status.total:
+                self.status.total = target_files[i].size
+                self.status.fp = target_files[i]
+
         while True:
             line = ""
             done_files = 0
@@ -168,6 +175,7 @@ class HexDiff(Module):
 
             last_line = line
             loop_count += 1
+            self.status.completed += self.block
 
     def init(self):
         # To mimic expected behavior, if all options are False, we show everything