From fc7778d0a9e73de9090cdbc3d6b6f83efc212c10 Mon Sep 17 00:00:00 2001
From: Alex Kenion <alexk@wordfence.com>
Date: Tue, 5 Sep 2023 11:37:10 -0400
Subject: [PATCH] Added user-friendly messaging when terminal size is too small
 for progress output and improved process termination

---
 wordfence/cli/scan/scan.py    | 47 +++++++++++++++++++++--------------
 wordfence/scanning/scanner.py |  7 ++++++
 2 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/wordfence/cli/scan/scan.py b/wordfence/cli/scan/scan.py
index 0587f0ff..7e8cafe4 100644
--- a/wordfence/cli/scan/scan.py
+++ b/wordfence/cli/scan/scan.py
@@ -18,7 +18,7 @@
 from wordfence.version import __version__
 from .reporting import Report, ReportFormat
 from .configure import Configurer
-from .progress import ProgressDisplay, reset_terminal
+from .progress import ProgressDisplay, ProgressException, reset_terminal
 
 
 screen_handler: Optional[logging.Handler] = None
@@ -223,12 +223,12 @@ def execute(self) -> int:
                         'output'
                     )
                 return 1
-            scanner = scanning.scanner.Scanner(options)
+            self.scanner = scanning.scanner.Scanner(options)
             if progress:
                 use_log_events = True
             else:
                 use_log_events = False
-            scanner.scan(
+            self.scanner.scan(
                     report.add_result,
                     progress.handle_update if progress is not None else None,
                     progress.scan_finished_handler if progress is not None
@@ -243,25 +243,22 @@ def execute(self) -> int:
                     print(progress.results_message)
         return 0
 
-
-def handle_repeated_interrupt(signal_number: int, stack) -> None:
-    revert_progress_changes()
-    if parent_process() is None:
-        log.warning('Scan command terminating immediately...')
-        reset_terminal()
-    os._exit(130)
+    def terminate(self) -> None:
+        if hasattr(self, 'scanner') and self.scanner is not None:
+            self.scanner.terminate()
 
 
-def handle_interrupt(signal_number: int, stack) -> None:
-    revert_progress_changes()
-    if parent_process() is None:
-        log.info('Scan command interrupted, stopping...')
-        reset_terminal()
-    signal.signal(signal.SIGINT, handle_repeated_interrupt)
-    sys.exit(130)
+def initialize_interrupt_handling(command: ScanCommand) -> None:
 
+    def handle_interrupt(signal_number: int, stack) -> None:
+        revert_progress_changes()
+        if parent_process() is None:
+            log.info('Scan command interrupted, stopping...')
+            command.terminate()
+            reset_terminal()
+        sys.exit(130)
 
-signal.signal(signal.SIGINT, handle_interrupt)
+    signal.signal(signal.SIGINT, handle_interrupt)
 
 
 def print_error(message: str) -> None:
@@ -301,12 +298,15 @@ def main(config) -> int:
         configurer.check_config()
         if not config.configure:
             command = ScanCommand(config)
+            initialize_interrupt_handling(command)
             command.execute()
         return 0
     except api.licensing.LicenseRequiredException:
         reset_terminal_with_error('A valid Wordfence CLI license is required')
         return 1
     except BaseException as exception:
+        if command is not None:
+            command.terminate()
         reset_terminal()
         if isinstance(exception, ExceptionContainer):
             if config.debug:
@@ -316,5 +316,14 @@ def main(config) -> int:
         if config.debug:
             raise exception
         else:
-            print_error(f'Error: {exception}')
+            if isinstance(exception, ProgressException):
+                print_error(
+                        'The current terminal size is inadequate for '
+                        'displaying progress output for the current scan '
+                        'options'
+                    )
+            elif isinstance(exception, SystemExit):
+                raise exception
+            else:
+                print_error(f'Error: {exception}')
         return 1
diff --git a/wordfence/scanning/scanner.py b/wordfence/scanning/scanner.py
index c0cf7034..b9572e6c 100644
--- a/wordfence/scanning/scanner.py
+++ b/wordfence/scanning/scanner.py
@@ -647,6 +647,7 @@ class Scanner:
     def __init__(self, options: Options):
         self.options = options
         self.failed = 0
+        self.active = []
 
     def _handle_worker_error(self, error: Exception):
         self.failed += 1
@@ -668,6 +669,7 @@ def scan(
                 use_log_events=use_log_events,
                 event_queue=event_queue if use_log_events else None
             )
+        self.active.append(file_locator_process)
         file_locator_process.start()
         for path in self.options.paths:
             file_locator_process.add_path(path)
@@ -691,6 +693,7 @@ def scan(
                     scanned_content_limit=self.options.scanned_content_limit,
                     use_log_events=use_log_events
                 ) as worker_pool:
+            self.active.append(worker_pool)
             if self.options.path_source is not None:
                 log.debug('Reading input paths...')
                 while True:
@@ -705,3 +708,7 @@ def scan(
         scan_finished_handler = scan_finished_handler if scan_finished_handler\
             else default_scan_finished_handler
         scan_finished_handler(metrics, timer)
+
+    def terminate(self) -> None:
+        for active in self.active:
+            active.terminate()