From 223791dbb8a1c84bfb3f4aa05f6b15924b1c5eae Mon Sep 17 00:00:00 2001 From: Mathieu Tarral Date: Wed, 30 Sep 2020 15:39:55 +0200 Subject: [PATCH 1/3] output: add exit handler for Rich output --- checksec/output.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/checksec/output.py b/checksec/output.py index 06587f0..24169d3 100644 --- a/checksec/output.py +++ b/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): From c39df2ff8bc7f9a4eb93f8c7dc1947d18bd9c195 Mon Sep 17 00:00:00 2001 From: Mathieu Tarral Date: Wed, 30 Sep 2020 16:14:22 +0200 Subject: [PATCH 2/3] checksec: handle KeyboardInterrupt --- checksec/__main__.py | 67 ++++++++++++++++++++++++++------------------ 1 file changed, 40 insertions(+), 27 deletions(-) diff --git a/checksec/__main__.py b/checksec/__main__.py index 166d06d..f8ffb62 100644 --- a/checksec/__main__.py +++ b/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__": From c06352e6a2113a4d4f7b6c62bb4017de14480958 Mon Sep 17 00:00:00 2001 From: Mathieu Tarral Date: Wed, 30 Sep 2020 16:15:07 +0200 Subject: [PATCH 3/3] release v0.3.9 --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index a65ccdf..9e98735 100644 --- a/setup.py +++ b/setup.py @@ -8,7 +8,7 @@ setuptools.setup( name="checksec.py", - version="0.3.8", + version="0.3.9", author="Mathieu Tarral", author_email="mathieu.tarral@protonmail.com", description="Checksec tool implemented in Python",