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

bbot/scanner/scanner.py

@@ -10,11 +10,11 @@ from datetime import datetime
 from collections import OrderedDict
 
 from bbot import __version__
-
 from bbot.core.event import make_event
 from .manager import ScanIngress, ScanEgress
 from bbot.core.helpers.misc import sha1, rand_string
 from bbot.core.helpers.names_generator import random_name
+from bbot.core.multiprocess import SHARED_INTERPRETER_STATE
 from bbot.core.helpers.async_helpers import async_to_sync_gen
 from bbot.errors import BBOTError, ScanError, ValidationError
 
@@ -115,25 +115,37 @@ class Scanner:
             dispatcher (Dispatcher, optional): Dispatcher object to use. Defaults to new Dispatcher.
             **kwargs (list[str], optional): Additional keyword arguments (passed through to `Preset`).
         """
+        self._root_event = None
+        self._finish_event = None
+        self.start_time = None
+        self.end_time = None
+        self.duration = None
+        self.duration_human = None
+        self.duration_seconds = None
+
+        self._success = False
+
         if scan_id is not None:
-            self.id = str(id)
+            self.id = str(scan_id)
         else:
             self.id = f"SCAN:{sha1(rand_string(20)).hexdigest()}"
 
-        preset = kwargs.pop("preset", None)
+        custom_preset = kwargs.pop("preset", None)
         kwargs["_log"] = True
 
         from .preset import Preset
 
-        if preset is None:
-            preset = Preset(*targets, **kwargs)
-        else:
-            if not isinstance(preset, Preset):
-                raise ValidationError(f'Preset must be of type Preset, not "{type(preset).__name__}"')
-        self.preset = preset.bake(self)
+        base_preset = Preset(*targets, **kwargs)
+
+        if custom_preset is not None:
+            if not isinstance(custom_preset, Preset):
+                raise ValidationError(f'Preset must be of type Preset, not "{type(custom_preset).__name__}"')
+            base_preset.merge(custom_preset)
+
+        self.preset = base_preset.bake(self)
 
         # scan name
-        if preset.scan_name is None:
+        if self.preset.scan_name is None:
             tries = 0
             while 1:
                 if tries > 5:
@@ -148,12 +160,16 @@ class Scanner:
                     break
                 tries += 1
         else:
-            scan_name = str(preset.scan_name)
-        self.name = scan_name
+            scan_name = str(self.preset.scan_name)
+        self.name = scan_name.replace("/", "_")
+
+        # make sure the preset has a description
+        if not self.preset.description:
+            self.preset.description = self.name
 
         # scan output dir
-        if preset.output_dir is not None:
-            self.home = Path(preset.output_dir).resolve() / self.name
+        if self.preset.output_dir is not None:
+            self.home = Path(self.preset.output_dir).resolve() / self.name
         else:
             self.home = self.preset.bbot_home / "scans" / self.name
 
@@ -198,8 +214,8 @@ class Scanner:
             )
 
         # url file extensions
-        self.url_extension_blacklist = set(e.lower() for e in self.config.get("url_extension_blacklist", []))
-        self.url_extension_httpx_only = set(e.lower() for e in self.config.get("url_extension_httpx_only", []))
+        self.url_extension_blacklist = {e.lower() for e in self.config.get("url_extension_blacklist", [])}
+        self.url_extension_httpx_only = {e.lower() for e in self.config.get("url_extension_httpx_only", [])}
 
         # url querystring behavior
         self.url_querystring_remove = self.config.get("url_querystring_remove", True)
@@ -232,6 +248,8 @@ class Scanner:
         self._dns_strings = None
         self._dns_regexes = None
         self._dns_regexes_yara = None
+        self._dns_yara_rules_uncompiled = None
+        self._dns_yara_rules = None
 
         self.__log_handlers = None
         self._log_handler_backup = []
@@ -241,6 +259,9 @@ class Scanner:
         Creates the scan's output folder, loads its modules, and calls their .setup() methods.
         """
 
+        # update the master PID
+        SHARED_INTERPRETER_STATE.update_scan_pid()
+
         self.helpers.mkdir(self.home)
         if not self._prepped:
             # save scan preset
@@ -248,7 +269,7 @@ class Scanner:
                 f.write(self.preset.to_yaml())
 
             # log scan overview
-            start_msg = f"Scan with {len(self.preset.scan_modules):,} modules seeded with {len(self.target):,} targets"
+            start_msg = f"Scan seeded with {len(self.seeds):,} targets"
             details = []
             if self.whitelist != self.target:
                 details.append(f"{len(self.whitelist):,} in whitelist")
@@ -272,7 +293,9 @@ class Scanner:
                 self.debug(
                     f"Setting intercept module {intercept_module.name}._incoming_event_queue to previous intercept module {prev_intercept_module.name}.outgoing_event_queue"
                 )
-                intercept_module._incoming_event_queue = prev_intercept_module.outgoing_event_queue
+                interqueue = asyncio.Queue()
+                intercept_module._incoming_event_queue = interqueue
+                prev_intercept_module._outgoing_event_queue = interqueue
 
             # abort if there are no output modules
             num_output_modules = len([m for m in self.modules.values() if m._type == "output"])
@@ -304,18 +327,18 @@ class Scanner:
 
     async def async_start(self):
         """ """
-        failed = True
-        scan_start_time = datetime.now()
+        self.start_time = datetime.now()
+        self.root_event.data["started_at"] = self.start_time.isoformat()
         try:
             await self._prep()
 
             self._start_log_handlers()
-            self.trace(f'Ran BBOT {__version__} at {scan_start_time}, command: {" ".join(sys.argv)}')
+            self.trace(f'Ran BBOT {__version__} at {self.start_time}, command: {" ".join(sys.argv)}')
             self.trace(f"Target: {self.preset.target.json}")
             self.trace(f"Preset: {self.preset.to_dict(redact_secrets=True)}")
 
             if not self.target:
-                self.warning(f"No scan targets specified")
+                self.warning("No scan targets specified")
 
             # start status ticker
             self.ticker_task = asyncio.create_task(
@@ -325,7 +348,7 @@ class Scanner:
             self.status = "STARTING"
 
             if not self.modules:
-                self.error(f"No modules loaded")
+                self.error("No modules loaded")
                 self.status = "FAILED"
                 return
             else:
@@ -339,7 +362,8 @@ class Scanner:
 
             # distribute seed events
             self.init_events_task = asyncio.create_task(
-                self.ingress_module.init_events(self.target.events), name=f"{self.name}.ingress_module.init_events()"
+                self.ingress_module.init_events(self.target.seeds.events),
+                name=f"{self.name}.ingress_module.init_events()",
             )
 
             # main scan loop
@@ -361,16 +385,19 @@ class Scanner:
                 if self._finished_init and self.modules_finished:
                     new_activity = await self.finish()
                     if not new_activity:
+                        self._success = True
+                        scan_finish_event = await self._mark_finished()
+                        yield scan_finish_event
                         break
 
                 await asyncio.sleep(0.1)
 
-            failed = False
+            self._success = True
 
         except BaseException as e:
             if self.helpers.in_exception_chain(e, (KeyboardInterrupt, asyncio.CancelledError)):
                 self.stop()
-                failed = False
+                self._success = True
             else:
                 try:
                     raise
@@ -394,26 +421,47 @@ class Scanner:
             await self._report()
             await self._cleanup()
 
-            log_fn = self.hugesuccess
-            if self.status == "ABORTING":
-                self.status = "ABORTED"
-                log_fn = self.hugewarning
-            elif failed:
-                self.status = "FAILED"
-                log_fn = self.critical
-            else:
-                self.status = "FINISHED"
-
-            scan_run_time = datetime.now() - scan_start_time
-            scan_run_time = self.helpers.human_timedelta(scan_run_time)
-            log_fn(f"Scan {self.name} completed in {scan_run_time} with status {self.status}")
-
             await self.dispatcher.on_finish(self)
 
             self._stop_log_handlers()
 
+    async def _mark_finished(self):
+        log_fn = self.hugesuccess
+        if self.status == "ABORTING":
+            status = "ABORTED"
+            log_fn = self.hugewarning
+        elif not self._success:
+            status = "FAILED"
+            log_fn = self.critical
+        else:
+            status = "FINISHED"
+
+        self.end_time = datetime.now()
+        self.duration = self.end_time - self.start_time
+        self.duration_seconds = self.duration.total_seconds()
+        self.duration_human = self.helpers.human_timedelta(self.duration)
+
+        status_message = f"Scan {self.name} completed in {self.duration_human} with status {status}"
+
+        scan_finish_event = self.finish_event(status_message, status)
+
+        # queue final scan event with output modules
+        output_modules = [m for m in self.modules.values() if m._type == "output" and m.name != "python"]
+        for m in output_modules:
+            await m.queue_event(scan_finish_event)
+        # wait until output modules are flushed
+        while 1:
+            modules_finished = all(m.finished for m in output_modules)
+            if modules_finished:
+                break
+            await asyncio.sleep(0.05)
+
+        self.status = status
+        log_fn(status_message)
+        return scan_finish_event
+
     def _start_modules(self):
-        self.verbose(f"Starting module worker loops")
+        self.verbose("Starting module worker loops")
         for module in self.modules.values():
             module.start()
 
@@ -437,17 +485,17 @@ class Scanner:
             Soft-failed modules are not set to an error state but are also removed if `remove_failed` is True.
         """
         await self.load_modules()
-        self.verbose(f"Setting up modules")
+        self.verbose("Setting up modules")
         succeeded = []
         hard_failed = []
         soft_failed = []
 
         async for task in self.helpers.as_completed([m._setup() for m in self.modules.values()]):
             module, status, msg = await task
-            if status == True:
+            if status is True:
                 self.debug(f"Setup succeeded for {module.name} ({msg})")
                 succeeded.append(module.name)
-            elif status == False:
+            elif status is False:
                 self.warning(f"Setup hard-failed for {module.name}: {msg}")
                 self.modules[module.name].set_error_state()
                 hard_failed.append(module.name)
@@ -489,11 +537,11 @@ class Scanner:
         """
         if not self._modules_loaded:
             if not self.preset.modules:
-                self.warning(f"No modules to load")
+                self.warning("No modules to load")
                 return
 
             if not self.preset.scan_modules:
-                self.warning(f"No scan modules to load")
+                self.warning("No scan modules to load")
 
             # install module dependencies
             succeeded, failed = await self.helpers.depsinstaller.install(*self.preset.modules)
@@ -637,7 +685,7 @@ class Scanner:
 
             if modules_errored:
                 self.verbose(
-                    f'{self.name}: Modules errored: {len(modules_errored):,} ({", ".join([m for m in modules_errored])})'
+                    f'{self.name}: Modules errored: {len(modules_errored):,} ({", ".join(list(modules_errored))})'
                 )
 
             num_queued_events = self.num_queued_events
@@ -674,7 +722,7 @@ class Scanner:
                     memory_usage = module.memory_usage
                     module_memory_usage.append((module.name, memory_usage))
                 module_memory_usage.sort(key=lambda x: x[-1], reverse=True)
-                self.debug(f"MODULE MEMORY USAGE:")
+                self.debug("MODULE MEMORY USAGE:")
                 for module_name, usage in module_memory_usage:
                     self.debug(f"    - {module_name}: {self.helpers.bytes_to_human(usage)}")
 
@@ -721,12 +769,12 @@ class Scanner:
             # Trigger .finished() on every module and start over
             log.info("Finishing scan")
             for module in self.modules.values():
-                finished_event = self.make_event(f"FINISHED", "FINISHED", dummy=True, tags={module.name})
+                finished_event = self.make_event("FINISHED", "FINISHED", dummy=True, tags={module.name})
                 await module.queue_event(finished_event)
             self.verbose("Completed finish()")
             return True
-        # Return False if no new events were generated since last time
         self.verbose("Completed final finish()")
+        # Return False if no new events were generated since last time
         return False
 
     def _drain_queues(self):
@@ -849,6 +897,10 @@ class Scanner:
     def target(self):
         return self.preset.target
 
+    @property
+    def seeds(self):
+        return self.preset.seeds
+
     @property
     def whitelist(self):
         return self.preset.whitelist
@@ -945,12 +997,25 @@ class Scanner:
             }
             ```
         """
-        root_event = self.make_event(data=self.json, event_type="SCAN", dummy=True)
+        if self._root_event is None:
+            self._root_event = self.make_root_event(f"Scan {self.name} started at {self.start_time}")
+        self._root_event.data["status"] = self.status
+        return self._root_event
+
+    def finish_event(self, context=None, status=None):
+        if self._finish_event is None:
+            if context is None or status is None:
+                raise ValueError("Must specify context and status")
+            self._finish_event = self.make_root_event(context)
+            self._finish_event.data["status"] = status
+        return self._finish_event
+
+    def make_root_event(self, context):
+        root_event = self.make_event(data=self.json, event_type="SCAN", dummy=True, context=context)
         root_event._id = self.id
         root_event.scope_distance = 0
         root_event.parent = root_event
         root_event.module = self._make_dummy_module(name="TARGET", _type="TARGET")
-        root_event.discovery_context = f"Scan {self.name} started at {root_event.timestamp}"
         return root_event
 
     @property
@@ -959,14 +1024,13 @@ class Scanner:
         A list of DNS hostname strings generated from the scan target
         """
         if self._dns_strings is None:
-            dns_targets = set(t.host for t in self.target if t.host and isinstance(t.host, str))
-            dns_whitelist = set(t.host for t in self.whitelist if t.host and isinstance(t.host, str))
-            dns_targets.update(dns_whitelist)
-            dns_targets = sorted(dns_targets, key=len)
-            dns_targets_set = set()
+            dns_whitelist = {t.host for t in self.whitelist if t.host and isinstance(t.host, str)}
+            dns_whitelist = sorted(dns_whitelist, key=len)
+            dns_whitelist_set = set()
             dns_strings = []
-            for t in dns_targets:
-                if not any(x in dns_targets_set for x in self.helpers.domain_parents(t, include_self=True)):
+            for t in dns_whitelist:
+                if not any(x in dns_whitelist_set for x in self.helpers.domain_parents(t, include_self=True)):
+                    dns_whitelist_set.add(t)
                     dns_strings.append(t)
             self._dns_strings = dns_strings
         return self._dns_strings
@@ -1011,21 +1075,67 @@ class Scanner:
         Returns a list of DNS hostname regexes formatted specifically for compatibility with YARA rules.
         """
         if self._dns_regexes_yara is None:
-            self._dns_regexes_yara = self._generate_dns_regexes(r"(([a-z0-9-]+\.)+")
+            self._dns_regexes_yara = self._generate_dns_regexes(r"(([a-z0-9-]+\.)*")
         return self._dns_regexes_yara
 
+    @property
+    def dns_yara_rules_uncompiled(self):
+        if self._dns_yara_rules_uncompiled is None:
+            regexes_component_list = []
+            for i, r in enumerate(self.dns_regexes_yara):
+                regexes_component_list.append(rf"$dns_name_{i} = /\b{r.pattern}/ nocase")
+            if regexes_component_list:
+                regexes_component = " ".join(regexes_component_list)
+                self._dns_yara_rules_uncompiled = f'rule hostname_extraction {{meta: description = "matches DNS hostname pattern derived from target(s)" strings: {regexes_component} condition: any of them}}'
+        return self._dns_yara_rules_uncompiled
+
+    async def dns_yara_rules(self):
+        if self._dns_yara_rules is None:
+            if self.dns_yara_rules_uncompiled is not None:
+                import yara
+
+                self._dns_yara_rules = await self.helpers.run_in_executor(
+                    yara.compile, source=self.dns_yara_rules_uncompiled
+                )
+        return self._dns_yara_rules
+
+    async def extract_in_scope_hostnames(self, s):
+        """
+        Given a string, uses yara to extract hostnames matching scan targets
+
+        Examples:
+            >>> await self.scan.extract_in_scope_hostnames("http://www.evilcorp.com")
+            ... {"www.evilcorp.com"}
+        """
+        matches = set()
+        dns_yara_rules = await self.dns_yara_rules()
+        if dns_yara_rules is not None:
+            for match in await self.helpers.run_in_executor(dns_yara_rules.match, data=s):
+                for string in match.strings:
+                    for instance in string.instances:
+                        matches.add(str(instance))
+        return matches
+
     @property
     def json(self):
         """
         A dictionary representation of the scan including its name, ID, targets, whitelist, blacklist, and modules
         """
-        j = dict()
+        j = {}
         for i in ("id", "name"):
             v = getattr(self, i, "")
             if v:
                 j.update({i: v})
         j["target"] = self.preset.target.json
         j["preset"] = self.preset.to_dict(redact_secrets=True)
+        if self.start_time is not None:
+            j["started_at"] = self.start_time.isoformat()
+        if self.end_time is not None:
+            j["finished_at"] = self.end_time.isoformat()
+        if self.duration is not None:
+            j["duration_seconds"] = self.duration_seconds
+        if self.duration_human is not None:
+            j["duration"] = self.duration_human
         return j
 
     def debug(self, *args, trace=False, **kwargs):
@@ -1181,7 +1291,7 @@ class Scanner:
             context = f"{context.__qualname__}()"
         filename, lineno, funcname = self.helpers.get_traceback_details(e)
         if self.helpers.in_exception_chain(e, (KeyboardInterrupt,)):
-            log.debug(f"Interrupted")
+            log.debug("Interrupted")
             self.stop()
         elif isinstance(e, BrokenPipeError):
             log.debug(f"BrokenPipeError in {filename}:{lineno}:{funcname}(): {e}")