Merge pull request #64 from Wenzel/handle_keyboardinterrupt

Handle keyboardinterrupt

checksec/__main__.py

@@ -76,38 +76,51 @@ def main(args):
         output_cls = JSONOutput
 
     with output_cls() as check_output:
-        # we need to consume the iterator once to get the total
-        # for the progress bar
-        check_output.enumerating_tasks_start()
-        count = sum(1 for i in walk_lief_parsable_list(filepath_list, recursive))
-        check_output.enumerating_tasks_stop(count)
-        with ProcessPoolExecutor(max_workers=workers) as pool:
-            check_output.processing_tasks_start()
-            future_to_checksec = {
-                pool.submit(checksec_file, filepath): filepath
-                for filepath in walk_lief_parsable_list(filepath_list, recursive)
-            }
-            for future in as_completed(future_to_checksec):
-                filepath = future_to_checksec[future]
+        try:
+            # we need to consume the iterator once to get the total
+            # for the progress bar
+            check_output.enumerating_tasks_start()
+            count = sum(1 for i in walk_lief_parsable_list(filepath_list, recursive))
+            check_output.enumerating_tasks_stop(count)
+            with ProcessPoolExecutor(max_workers=workers) as pool:
                 try:
-                    data = future.result()
-                except FileNotFoundError:
-                    logging.debug("%s does not exist", filepath)
-                except ErrorParsingFailed:
-                    logging.debug("%s LIEF parsing failed")
-                except NotImplementedError:
-                    logging.debug("%s: Not an ELF/PE. Skipping", filepath)
-                else:
-                    check_output.add_checksec_result(filepath, data)
-                finally:
-                    check_output.checksec_result_end()
-
-        check_output.print()
+                    check_output.processing_tasks_start()
+                    future_to_checksec = {
+                        pool.submit(checksec_file, filepath): filepath
+                        for filepath in walk_lief_parsable_list(filepath_list, recursive)
+                    }
+                    for future in as_completed(future_to_checksec):
+                        filepath = future_to_checksec[future]
+                        try:
+                            data = future.result()
+                        except FileNotFoundError:
+                            logging.debug("%s does not exist", filepath)
+                        except ErrorParsingFailed:
+                            logging.debug("%s LIEF parsing failed")
+                        except NotImplementedError:
+                            logging.debug("%s: Not an ELF/PE. Skipping", filepath)
+                        else:
+                            check_output.add_checksec_result(filepath, data)
+                        finally:
+                            check_output.checksec_result_end()
+                except KeyboardInterrupt:
+                    # remove progress bars before waiting for ProcessPoolExecutor to shutdown
+                    check_output.__exit__(None, None, None)
+                    logging.info("Shutdown Process Pool ...")
+                    pool.shutdown(wait=True)
+                    raise
+        except KeyboardInterrupt:
+            pass
+        else:
+            check_output.print()
 
 
 def entrypoint():
     args = docopt(__doc__)
-    main(args)
+    try:
+        main(args)
+    except KeyboardInterrupt:
+        pass
 
 
 if __name__ == "__main__":

checksec/output.py

@@ -110,6 +110,15 @@ def __init__(self):
 
         self.process_task_id = None
 
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # cleanup the Rich progress bars
+        if self.enumerate_bar is not None:
+            self.enumerate_bar.stop()
+        if self.process_bar is not None:
+            self.process_bar.stop()
+        if self.display_res_bar is not None:
+            self.display_res_bar.stop()
+
     def enumerating_tasks_start(self):
         # start progress bar
         self.enumerate_bar.start()
@@ -118,6 +127,7 @@ def enumerating_tasks_start(self):
     def enumerating_tasks_stop(self, total: int):
         super().enumerating_tasks_stop(total)
         self.enumerate_bar.stop()
+        self.enumerate_bar = None
 
     def processing_tasks_start(self):
         # init progress bar
@@ -287,6 +297,7 @@ def checksec_result_end(self):
 
     def print(self):
         self.process_bar.stop()
+        self.process_bar = None
 
         if self.table_elf.row_count > 0:
             with self.display_res_bar:
@@ -299,6 +310,9 @@ def print(self):
                 self.console.print(self.table_pe)
                 self.display_res_bar.remove_task(task_id)
 
+        self.display_res_bar.stop()
+        self.display_res_bar = None
+
 
 class JSONOutput(AbstractChecksecOutput):
     def __init__(self):

setup.py

@@ -8,7 +8,7 @@
 
 setuptools.setup(
     name="checksec.py",
-    version="0.3.8",
+    version="0.3.9",
     author="Mathieu Tarral",
     author_email="[email protected]",
     description="Checksec tool implemented in Python",