From 9cb214b2c020d3db94f97d5e7bb910e1f89ee98b Mon Sep 17 00:00:00 2001
From: TheTechromancer <thetechromancer@protonmail.com>
Date: Wed, 20 Sep 2023 14:27:20 -0400
Subject: [PATCH] Added ability to kill modules during scan

---
 bbot/cli.py             | 14 +++++++++++++-
 bbot/modules/base.py    |  9 +++++++--
 bbot/scanner/manager.py |  4 ++++
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/bbot/cli.py b/bbot/cli.py
index 4069867158..675f320df4 100755
--- a/bbot/cli.py
+++ b/bbot/cli.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import os
+import re
 import sys
 import asyncio
 import logging
@@ -307,6 +308,7 @@ async def _main():
 
                     def keyboard_listen():
                         allowed_errors = 10
+                        kill_regex = re.compile(r"kill (?P<module>[a-z0-9_]+)")
                         while 1:
                             keyboard_input = "a"
                             try:
@@ -314,7 +316,17 @@ def keyboard_listen():
                                 allowed_errors = 10
                             except Exception:
                                 allowed_errors -= 1
-                            if not keyboard_input:
+                            if keyboard_input:
+                                log.verbose(f'Got keyboard input: "{keyboard_input}"')
+                                kill_match = kill_regex.match(keyboard_input)
+                                if kill_match:
+                                    module = kill_match.group("module")
+                                    if module in scanner.modules:
+                                        log.hugewarning(f'Killing module: "{module}"')
+                                        scanner.manager.kill_module(module, message="killed by user")
+                                    else:
+                                        log.warning(f'Invalid module: "{module}"')
+                            else:
                                 toggle_log_level(logger=log)
                                 scanner.manager.modules_status(_log=True)
                             if allowed_errors <= 0:
diff --git a/bbot/modules/base.py b/bbot/modules/base.py
index 2057212b63..3bb206f983 100644
--- a/bbot/modules/base.py
+++ b/bbot/modules/base.py
@@ -324,7 +324,7 @@ async def _setup(self):
     async def _worker(self):
         async with self.scan.acatch(context=self._worker):
             try:
-                while not self.scan.stopping:
+                while not self.scan.stopping and not self.errored:
                     # hold the reigns if our outgoing queue is full
                     if self._qsize > 0 and self.outgoing_event_queue.qsize() >= self._qsize:
                         await asyncio.sleep(0.1)
@@ -505,7 +505,7 @@ def queue_outgoing_event(self, event, **kwargs):
         except AttributeError:
             self.debug(f"Not in an acceptable state to queue outgoing event")
 
-    def set_error_state(self, message=None):
+    def set_error_state(self, message=None, clear_outgoing_queue=False):
         if not self.errored:
             log_msg = f"Setting error state for module {self.name}"
             if message is not None:
@@ -522,6 +522,11 @@ def set_error_state(self, message=None):
                 # if there are leftover objects in the queue, the scan will hang.
                 self._incoming_event_queue = False
 
+            if clear_outgoing_queue:
+                with suppress(asyncio.queues.QueueEmpty):
+                    while 1:
+                        self.outgoing_event_queue.get_nowait()
+
     # override in the module to define different values to comprise the hash
     def get_per_host_hash(self, event):
         parsed = getattr(event, "parsed", None)
diff --git a/bbot/scanner/manager.py b/bbot/scanner/manager.py
index d085a60b0d..1ed44d37ff 100644
--- a/bbot/scanner/manager.py
+++ b/bbot/scanner/manager.py
@@ -346,6 +346,10 @@ async def _worker_loop(self):
         except Exception:
             log.critical(traceback.format_exc())
 
+    def kill_module(self, module_name, message=None):
+        module = self.scan.modules[module_name]
+        module.set_error_state(message=message, clear_outgoing_queue=True)
+
     @property
     def modules_by_priority(self):
         if not self._modules_by_priority: