bbot 2.0.1.4720rc0__py3-none-any.whl → 2.3.0.5401rc0__py3-none-any.whl

bbot/__init__.py

@@ -1,4 +1,4 @@
 # version placeholder (replaced by poetry-dynamic-versioning)
-__version__ = "v2.0.1.4720rc"
+__version__ = "v2.3.0.5401rc"
 
 from .scanner import Scanner, Preset

bbot/cli.py

@@ -29,7 +29,6 @@ scan_name = ""
 
 
 async def _main():
-
     import asyncio
     import traceback
     from contextlib import suppress
@@ -45,7 +44,6 @@ async def _main():
     global scan_name
 
     try:
-
         # start by creating a default scan preset
         preset = Preset(_log=True, name="bbot_cli_main")
         # parse command line arguments and merge into preset
@@ -81,7 +79,6 @@ async def _main():
 
         # if we're listing modules or their options
         if options.list_modules or options.list_module_options:
-
             # if no modules or flags are specified, enable everything
             if not (options.modules or options.output_modules or options.flags):
                 for module, preloaded in preset.module_loader.preloaded().items():
@@ -133,8 +130,8 @@ async def _main():
         ]
         if deadly_modules and not options.allow_deadly:
             log.hugewarning(f"You enabled the following deadly modules: {','.join(deadly_modules)}")
-            log.hugewarning(f"Deadly modules are highly intrusive")
-            log.hugewarning(f"Please specify --allow-deadly to continue")
+            log.hugewarning("Deadly modules are highly intrusive")
+            log.hugewarning("Please specify --allow-deadly to continue")
             return False
 
         # --current-preset
@@ -172,9 +169,8 @@ async def _main():
             log.trace(f"Command: {' '.join(sys.argv)}")
 
             if sys.stdin.isatty():
-
                 # warn if any targets belong directly to a cloud provider
-                for event in scan.target.events:
+                for event in scan.target.seeds.events:
                     if event.type == "DNS_NAME":
                         cloudcheck_result = scan.helpers.cloudcheck(event.host)
                         if cloudcheck_result:

bbot/core/config/files.py

@@ -10,7 +10,6 @@ bbot_code_dir = Path(__file__).parent.parent.parent
 
 
 class BBOTConfigFiles:
-
     config_dir = (Path.home() / ".config" / "bbot").resolve()
     defaults_filename = (bbot_code_dir / "defaults.yml").resolve()
     config_filename = (config_dir / "bbot.yml").resolve()

bbot/core/config/logger.py

@@ -1,3 +1,4 @@
+import os
 import sys
 import atexit
 import logging
@@ -5,9 +6,11 @@ from copy import copy
 import multiprocessing
 import logging.handlers
 from pathlib import Path
+from contextlib import suppress
 
 from ..helpers.misc import mkdir, error_and_exit
 from ...logger import colorize, loglevel_mapping
+from ..multiprocess import SHARED_INTERPRETER_STATE
 
 
 debug_format = logging.Formatter("%(asctime)s [%(levelname)s] %(name)s %(filename)s:%(lineno)s %(message)s")
@@ -64,17 +67,44 @@ class BBOTLogger:
 
         self.listener = None
 
-        self.process_name = multiprocessing.current_process().name
-        if self.process_name == "MainProcess":
+        # if we haven't set up logging yet, do it now
+        if "_BBOT_LOGGING_SETUP" not in os.environ:
+            os.environ["_BBOT_LOGGING_SETUP"] = "1"
             self.queue = multiprocessing.Queue()
             self.setup_queue_handler()
             # Start the QueueListener
             self.listener = logging.handlers.QueueListener(self.queue, *self.log_handlers.values())
             self.listener.start()
-            atexit.register(self.listener.stop)
+            atexit.register(self.cleanup_logging)
 
         self.log_level = logging.INFO
 
+    def cleanup_logging(self):
+        # Close the queue handler
+        with suppress(Exception):
+            self.queue_handler.close()
+
+        # Clean root logger
+        root_logger = logging.getLogger()
+        for handler in list(root_logger.handlers):
+            with suppress(Exception):
+                root_logger.removeHandler(handler)
+            with suppress(Exception):
+                handler.close()
+
+        # Clean all other loggers
+        for logger in logging.Logger.manager.loggerDict.values():
+            if hasattr(logger, "handlers"):  # Logger, not PlaceHolder
+                for handler in list(logger.handlers):
+                    with suppress(Exception):
+                        logger.removeHandler(handler)
+                    with suppress(Exception):
+                        handler.close()
+
+        # Stop queue listener
+        with suppress(Exception):
+            self.listener.stop()
+
     def setup_queue_handler(self, logging_queue=None, log_level=logging.DEBUG):
         if logging_queue is None:
             logging_queue = self.queue
@@ -86,7 +116,7 @@ class BBOTLogger:
 
         self.core_logger.setLevel(log_level)
         # disable asyncio logging for child processes
-        if self.process_name != "MainProcess":
+        if not SHARED_INTERPRETER_STATE.is_main_process:
             logging.getLogger("asyncio").setLevel(logging.ERROR)
 
     def addLoggingLevel(self, levelName, levelNum, methodName=None):

bbot/core/core.py

@@ -6,6 +6,7 @@ from contextlib import suppress
 from omegaconf import OmegaConf
 
 from bbot.errors import BBOTError
+from .multiprocess import SHARED_INTERPRETER_STATE
 
 
 DEFAULT_CONFIG = None
@@ -41,9 +42,23 @@ class BBOTCore:
         self.logger
         self.log = logging.getLogger("bbot.core")
 
+        self._prep_multiprocessing()
+
+    def _prep_multiprocessing(self):
         import multiprocessing
+        from .helpers.process import BBOTProcess
+
+        if SHARED_INTERPRETER_STATE.is_main_process:
+            # if this is the main bbot process, set the logger and queue for the first time
+            from functools import partialmethod
 
-        self.process_name = multiprocessing.current_process().name
+            BBOTProcess.__init__ = partialmethod(
+                BBOTProcess.__init__, log_level=self.logger.log_level, log_queue=self.logger.queue
+            )
+
+        # this makes our process class the default for process pools, etc.
+        mp_context = multiprocessing.get_context("spawn")
+        mp_context.Process = BBOTProcess
 
     @property
     def home(self):
@@ -91,7 +106,7 @@ class BBOTCore:
         if DEFAULT_CONFIG is None:
             self.default_config = self.files_config.get_default_config()
             # ensure bbot home dir
-            if not "home" in self.default_config:
+            if "home" not in self.default_config:
                 self.default_config["home"] = "~/.bbot"
         return DEFAULT_CONFIG
 
@@ -187,12 +202,14 @@ class BBOTCore:
         if os.environ.get("BBOT_TESTING", "") == "True":
             process = self.create_thread(*args, **kwargs)
         else:
-            if self.process_name == "MainProcess":
+            if SHARED_INTERPRETER_STATE.is_scan_process:
                 from .helpers.process import BBOTProcess
 
                 process = BBOTProcess(*args, **kwargs)
             else:
-                raise BBOTError(f"Tried to start server from process {self.process_name}")
+                import multiprocessing
+
+                raise BBOTError(f"Tried to start server from process {multiprocessing.current_process().name}")
         process.daemon = True
         return process
 

bbot/core/engine.py

@@ -10,6 +10,7 @@ import traceback
 import contextlib
 import contextvars
 import zmq.asyncio
+import multiprocessing
 from pathlib import Path
 from concurrent.futures import CancelledError
 from contextlib import asynccontextmanager, suppress
@@ -17,6 +18,7 @@ from contextlib import asynccontextmanager, suppress
 from bbot.core import CORE
 from bbot.errors import BBOTEngineError
 from bbot.core.helpers.async_helpers import get_event_loop
+from bbot.core.multiprocess import SHARED_INTERPRETER_STATE
 from bbot.core.helpers.misc import rand_string, in_exception_chain
 
 
@@ -264,10 +266,8 @@ class EngineClient(EngineBase):
         return [s for s in self.CMDS if isinstance(s, str)]
 
     def start_server(self):
-        import multiprocessing
-
         process_name = multiprocessing.current_process().name
-        if process_name == "MainProcess":
+        if SHARED_INTERPRETER_STATE.is_scan_process:
             kwargs = dict(self.server_kwargs)
             # if we're in tests, we use a single event loop to avoid weird race conditions
             # this allows us to more easily mock http, etc.
@@ -641,8 +641,9 @@ class EngineServer(EngineBase):
             except BaseException as e:
                 if isinstance(e, (TimeoutError, asyncio.exceptions.TimeoutError)):
                     self.log.warning(f"{self.name}: Timeout after {timeout:,} seconds in finished_tasks({tasks})")
-                    for task in tasks:
+                    for task in list(tasks):
                         task.cancel()
+                        self._await_cancelled_task(task)
                 else:
                     if not in_exception_chain(e, (KeyboardInterrupt, asyncio.CancelledError)):
                         self.log.error(f"{self.name}: Unhandled exception in finished_tasks({tasks}): {e}")
@@ -664,9 +665,9 @@ class EngineServer(EngineBase):
                 child_task.cancel()
 
         for task in [parent_task] + list(child_tasks):
-            await self._cancel_task(task)
+            await self._await_cancelled_task(task)
 
-    async def _cancel_task(self, task):
+    async def _await_cancelled_task(self, task):
         try:
             await asyncio.wait_for(task, timeout=10)
         except (TimeoutError, asyncio.exceptions.TimeoutError):
@@ -682,5 +683,5 @@ class EngineServer(EngineBase):
         for client_id in list(self.tasks):
             await self.cancel_task(client_id)
         for client_id, tasks in self.child_tasks.items():
-            for task in tasks:
-                await self._cancel_task(task)
+            for task in list(tasks):
+                await self._await_cancelled_task(task)