Merge pull request #751 from blacklanternsecurity/kill-modules

Added Ability to Kill Modules During Scan

bbot/cli.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import os
+import re
 import sys
 import asyncio
 import logging
@@ -307,14 +308,25 @@ 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:
                                 keyboard_input = input()
                                 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:

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)

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: