Merge pull request #1603 from blacklanternsecurity/fix-dns-parents

DEF CON Update #1

bbot/core/engine.py

@@ -63,13 +63,14 @@ def unpickle(self, binary):
 
     async def _infinite_retry(self, callback, *args, **kwargs):
         interval = kwargs.pop("_interval", 10)
+        context = kwargs.pop("_context", "")
+        if not context:
+            context = f"{callback.__name__}({args}, {kwargs})"
         while not self._shutdown_status:
             try:
                 return await asyncio.wait_for(callback(*args, **kwargs), timeout=interval)
             except (TimeoutError, asyncio.TimeoutError):
-                self.log.debug(
-                    f"{self.name}: Timeout waiting for response for {callback.__name__}({args}, {kwargs}), retrying..."
-                )
+                self.log.debug(f"{self.name}: Timeout waiting for response for {context}, retrying...")
 
 
 class EngineClient(EngineBase):
@@ -144,10 +145,10 @@ async def run_and_return(self, command, *args, **kwargs):
                 if message is error_sentinel:
                     return
                 await self._infinite_retry(socket.send, message)
-                binary = await self._infinite_retry(socket.recv)
+                binary = await self._infinite_retry(socket.recv, _context=f"waiting for return value from {fn_str}")
             except BaseException:
                 try:
-                    await self.send_cancel_message(socket)
+                    await self.send_cancel_message(socket, fn_str)
                 except Exception:
                     self.log.debug(f"{self.name}: {fn_str} failed to send cancel message after exception")
                     self.log.trace(traceback.format_exc())
@@ -176,7 +177,9 @@ async def run_and_yield(self, command, *args, **kwargs):
             await socket.send(message)
             while 1:
                 try:
-                    binary = await self._infinite_retry(socket.recv)
+                    binary = await self._infinite_retry(
+                        socket.recv, _context=f"waiting for new iteration from {fn_str}"
+                    )
                     # self.log.debug(f"{self.name}.{command}({kwargs}) got binary: {binary}")
                     message = self.unpickle(binary)
                     self.log.debug(f"{self.name} {command} got iteration: {message}")
@@ -188,21 +191,21 @@ async def run_and_yield(self, command, *args, **kwargs):
                     exc_name = e.__class__.__name__
                     self.log.debug(f"{self.name}.{command} got {exc_name}")
                     try:
-                        await self.send_cancel_message(socket)
+                        await self.send_cancel_message(socket, fn_str)
                     except Exception:
                         self.log.debug(f"{self.name}.{command} failed to send cancel message after {exc_name}")
                         self.log.trace(traceback.format_exc())
                     break
 
-    async def send_cancel_message(self, socket):
+    async def send_cancel_message(self, socket, context):
         """
         Send a cancel message and wait for confirmation from the server
         """
         # -1 == special "cancel" signal
         message = pickle.dumps({"c": -1})
         await self._infinite_retry(socket.send, message)
         while 1:
-            response = await self._infinite_retry(socket.recv)
+            response = await self._infinite_retry(socket.recv, _context=f"waiting for CANCEL_OK from {context}")
             response = pickle.loads(response)
             if isinstance(response, dict):
                 response = response.get("m", "")

bbot/core/event/base.py

@@ -164,10 +164,12 @@ def __init__(
         self._priority = None
         self._parent_id = None
         self._host_original = None
+        self._scope_distance = None
         self._module_priority = None
         self._resolved_hosts = set()
         self.dns_children = dict()
         self._discovery_context = ""
+        self.web_spider_distance = 0
 
         # for creating one-off events without enforcing parent requirement
         self._dummy = _dummy
@@ -199,8 +201,6 @@ def __init__(
         if self.scan:
             self.scans = list(set([self.scan.id] + self.scans))
 
-        self._scope_distance = -1
-
         try:
             self.data = self._sanitize_data(data)
         except Exception as e:
@@ -214,9 +214,6 @@ def __init__(
         if (not self.parent) and (not self._dummy):
             raise ValidationError(f"Must specify event parent")
 
-        # inherit web spider distance from parent
-        self.web_spider_distance = getattr(self.parent, "web_spider_distance", 0)
-
         if tags is not None:
             for tag in tags:
                 self.add_tag(tag)
@@ -435,29 +432,29 @@ def scope_distance(self, scope_distance):
         Note:
             The method will automatically update the relevant 'distance-' tags associated with the event.
         """
-        if scope_distance >= 0:
-            new_scope_distance = None
-            # ensure scope distance does not increase (only allow setting to smaller values)
-            if self.scope_distance == -1:
-                new_scope_distance = scope_distance
+        if scope_distance < 0:
+            raise ValueError(f"Invalid scope distance: {scope_distance}")
+        # ensure scope distance does not increase (only allow setting to smaller values)
+        if self.scope_distance is None:
+            new_scope_distance = scope_distance
+        else:
+            new_scope_distance = min(self.scope_distance, scope_distance)
+        if self._scope_distance != new_scope_distance:
+            # remove old scope distance tags
+            for t in list(self.tags):
+                if t.startswith("distance-"):
+                    self.remove_tag(t)
+            if scope_distance == 0:
+                self.add_tag("in-scope")
+                self.remove_tag("affiliate")
             else:
-                new_scope_distance = min(self.scope_distance, scope_distance)
-            if self._scope_distance != new_scope_distance:
-                # remove old scope distance tags
-                for t in list(self.tags):
-                    if t.startswith("distance-"):
-                        self.remove_tag(t)
-                if scope_distance == 0:
-                    self.add_tag("in-scope")
-                    self.remove_tag("affiliate")
-                else:
-                    self.remove_tag("in-scope")
-                    self.add_tag(f"distance-{new_scope_distance}")
-                self._scope_distance = new_scope_distance
-            # apply recursively to parent events
-            parent_scope_distance = getattr(self.parent, "scope_distance", -1)
-            if parent_scope_distance >= 0 and self != self.parent:
-                self.parent.scope_distance = scope_distance + 1
+                self.remove_tag("in-scope")
+                self.add_tag(f"distance-{new_scope_distance}")
+            self._scope_distance = new_scope_distance
+        # apply recursively to parent events
+        parent_scope_distance = getattr(self.parent, "scope_distance", None)
+        if parent_scope_distance is not None and self != self.parent:
+            self.parent.scope_distance = scope_distance + 1
 
     @property
     def scope_description(self):
@@ -493,20 +490,27 @@ def parent(self, parent):
         """
         if is_event(parent):
             self._parent = parent
-            hosts_are_same = self.host and (self.host == parent.host)
-            if parent.scope_distance >= 0:
-                new_scope_distance = int(parent.scope_distance)
+            hosts_are_same = (self.host and parent.host) and (self.host == parent.host)
+            new_scope_distance = int(parent.scope_distance)
+            if self.host and parent.scope_distance is not None:
                 # only increment the scope distance if the host changes
                 if self._scope_distance_increment_same_host or not hosts_are_same:
                     new_scope_distance += 1
-                self.scope_distance = new_scope_distance
+            self.scope_distance = new_scope_distance
             # inherit certain tags
             if hosts_are_same:
+                # inherit web spider distance from parent
+                self.web_spider_distance = getattr(parent, "web_spider_distance", 0)
+                event_has_url = getattr(self, "parsed_url", None) is not None
                 for t in parent.tags:
-                    if t == "affiliate":
-                        self.add_tag("affiliate")
+                    if t in ("affiliate",):
+                        self.add_tag(t)
                     elif t.startswith("mutation-"):
                         self.add_tag(t)
+                    # only add these tags if the event has a URL
+                    if event_has_url:
+                        if t in ("spider-danger", "spider-max"):
+                            self.add_tag(t)
         elif not self._dummy:
             log.warning(f"Tried to set invalid parent on {self}: (got: {parent})")
 
@@ -539,9 +543,11 @@ def get_parent(self):
             return self.parent.get_parent()
         return self.parent
 
-    def get_parents(self, omit=False):
+    def get_parents(self, omit=False, include_self=False):
         parents = []
         e = self
+        if include_self:
+            parents.append(self)
         while 1:
             if omit:
                 parent = e.get_parent()
@@ -1098,12 +1104,13 @@ def sanitize_data(self, data):
         return data
 
     def add_tag(self, tag):
-        if tag == "spider-danger":
+        host_same_as_parent = self.parent and self.host == self.parent.host
+        if tag == "spider-danger" and host_same_as_parent and not "spider-danger" in self.tags:
             # increment the web spider distance
             if self.type == "URL_UNVERIFIED":
                 self.web_spider_distance += 1
-            if self.is_spider_max:
-                self.add_tag("spider-max")
+                if self.is_spider_max:
+                    self.add_tag("spider-max")
         super().add_tag(tag)
 
     @property

bbot/core/helpers/dns/engine.py

@@ -373,12 +373,12 @@ def new_task(query):
                 if queries:  # Start a new task for each one completed, if URLs remain
                     new_task(queries.pop(0))
 
-    async def resolve_raw_batch(self, queries, threads=10):
+    async def resolve_raw_batch(self, queries, threads=10, **kwargs):
         tasks = {}
         client_id = self.client_id_var.get()
 
         def new_task(query, rdtype):
-            task = self.new_child_task(client_id, self.resolve_raw(query, type=rdtype))
+            task = self.new_child_task(client_id, self.resolve_raw(query, type=rdtype, **kwargs))
             tasks[task] = (query, rdtype)
 
         queries = list(queries)
@@ -469,7 +469,12 @@ async def is_wildcard(self, query, ips=None, rdtype=None):
         parent = parent_domain(query)
         parents = list(domain_parents(query))
 
-        rdtypes_to_check = [rdtype] if rdtype is not None else all_rdtypes
+        if rdtype is not None:
+            if isinstance(rdtype, str):
+                rdtype = [rdtype]
+            rdtypes_to_check = rdtype
+        else:
+            rdtypes_to_check = all_rdtypes
 
         query_baseline = dict()
         # if the caller hasn't already done the work of resolving the IPs
@@ -534,6 +539,10 @@ async def is_wildcard(self, query, ips=None, rdtype=None):
         except DNSWildcardBreak:
             pass
 
+        for _rdtype, answers in query_baseline.items():
+            if answers and _rdtype not in result:
+                result[_rdtype] = (False, query)
+
         return result
 
     async def is_wildcard_domain(self, domain, log_info=False):
@@ -581,13 +590,13 @@ async def is_wildcard_domain(self, domain, log_info=False):
                 is_wildcard = False
                 wildcard_results = dict()
 
-                queries = []
+                rand_queries = []
                 for rdtype in rdtypes_to_check:
                     for _ in range(self.wildcard_tests):
                         rand_query = f"{rand_string(digits=False, length=10)}.{host}"
-                        queries.append((rand_query, rdtype))
+                        rand_queries.append((rand_query, rdtype))
 
-                async for (query, rdtype), (answers, errors) in self.resolve_raw_batch(queries):
+                async for (query, rdtype), (answers, errors) in self.resolve_raw_batch(rand_queries, use_cache=False):
                     answers = extract_targets(answers)
                     if answers:
                         is_wildcard = True

bbot/modules/base.py

@@ -1211,7 +1211,7 @@ def log_table(self, *args, **kwargs):
         return table
 
     def _is_graph_important(self, event):
-        return self.preserve_graph and getattr(event, "_graph_important", False)
+        return self.preserve_graph and getattr(event, "_graph_important", False) and not getattr(event, "_omit", False)
 
     @property
     def preserve_graph(self):
@@ -1380,7 +1380,7 @@ def error(self, *args, trace=True, **kwargs):
         if trace:
             self.trace()
 
-    def trace(self):
+    def trace(self, msg=None):
         """Logs the stack trace of the most recently caught exception.
 
         This method captures the type, value, and traceback of the most recent exception and logs it using the trace level. It is typically used for debugging purposes.
@@ -1393,9 +1393,12 @@ def trace(self):
             >>> except ZeroDivisionError:
             >>>     self.trace()
         """
-        e_type, e_val, e_traceback = exc_info()
-        if e_type is not None:
-            self.log.trace(traceback.format_exc())
+        if msg is None:
+            e_type, e_val, e_traceback = exc_info()
+            if e_type is not None:
+                self.log.trace(traceback.format_exc())
+        else:
+            self.log.trace(msg)
 
     def critical(self, *args, trace=True, **kwargs):
         """Logs a whole message in emboldened red text, and optionally the stack trace of the most recent exception.
@@ -1454,8 +1457,6 @@ async def _worker(self):
                             await self.finish()
                         continue
 
-                    self.debug(f"Got {event} from {getattr(event, 'module', 'unknown_module')}")
-
                     acceptable = True
                     async with self._task_counter.count(f"event_precheck({event})"):
                         precheck_pass, reason = self._event_precheck(event)
@@ -1482,12 +1483,11 @@ async def _worker(self):
                             with suppress(ValueError, TypeError):
                                 forward_event, forward_event_reason = forward_event
 
-                        self.debug(f"Finished intercepting {event}")
-
                         if forward_event is False:
                             self.debug(f"Not forwarding {event} because {forward_event_reason}")
                             continue
 
+                    self.debug(f"Forwarding {event}")
                     await self.forward_event(event, kwargs)
 
             except asyncio.CancelledError:

bbot/modules/code_repository.py

@@ -49,7 +49,6 @@ async def handle_event(self, event):
                         tags=platform,
                         parent=event,
                     )
-                    repo_event.scope_distance = event.scope_distance
                     await self.emit_event(
                         repo_event,
                         context=f"{{module}} detected {platform} {{event.type}} at {url}",

bbot/modules/docker_pull.py

@@ -60,7 +60,6 @@ async def handle_event(self, event):
                 parent=event,
             )
             if codebase_event:
-                codebase_event.scope_distance = event.scope_distance
                 await self.emit_event(
                     codebase_event, context=f"{{module}} downloaded Docker image to {{event.type}}: {repo_path}"
                 )

bbot/modules/filedownload.py

@@ -144,7 +144,6 @@ async def download_file(self, url, content_type=None, source_event=None):
                 file_event = self.make_event(
                     {"path": str(file_destination)}, "FILESYSTEM", tags=["filedownload", "file"], parent=source_event
                 )
-                file_event.scope_distance = source_event.scope_distance
                 await self.emit_event(file_event)
         self.urls_downloaded.add(hash(url))
 

bbot/modules/git_clone.py

@@ -40,7 +40,6 @@ async def handle_event(self, event):
         if repo_path:
             self.verbose(f"Cloned {repo_url} to {repo_path}")
             codebase_event = self.make_event({"path": str(repo_path)}, "FILESYSTEM", tags=["git"], parent=event)
-            codebase_event.scope_distance = event.scope_distance
             await self.emit_event(
                 codebase_event,
                 context=f"{{module}} downloaded git repo at {repo_url} to {{event.type}}: {repo_path}",

bbot/modules/github_codesearch.py

@@ -35,7 +35,6 @@ async def handle_event(self, event):
                 url_event = self.make_event(raw_url, "URL_UNVERIFIED", parent=repo_event, tags=["httpx-safe"])
                 if not url_event:
                     continue
-                url_event.scope_distance = repo_event.scope_distance
                 await self.emit_event(
                     url_event, context=f'file matching query "{query}" is at {{event.type}}: {raw_url}'
                 )

bbot/modules/github_org.py

@@ -62,7 +62,6 @@ async def handle_event(self, event):
                 repo_event = self.make_event({"url": repo_url}, "CODE_REPOSITORY", tags="git", parent=event)
                 if not repo_event:
                     continue
-                repo_event.scope_distance = event.scope_distance
                 await self.emit_event(
                     repo_event,
                     context=f"{{module}} listed repos for GitHub profile and discovered {{event.type}}: {repo_url}",
@@ -97,7 +96,6 @@ async def handle_event(self, event):
             event_data = {"platform": "github", "profile_name": user, "url": user_url}
             github_org_event = self.make_event(event_data, "SOCIAL", tags="github-org", parent=event)
             if github_org_event:
-                github_org_event.scope_distance = event.scope_distance
                 await self.emit_event(
                     github_org_event,
                     context=f'{{module}} tried "{user}" as GitHub profile and discovered {{event.type}}: {user_url}',

bbot/modules/github_workflows.py

@@ -58,7 +58,6 @@ async def handle_event(self, event):
                         tags=["textfile"],
                         parent=event,
                     )
-                    logfile_event.scope_distance = event.scope_distance
                     await self.emit_event(
                         logfile_event,
                         context=f"{{module}} downloaded workflow run logs from {workflow_url} to {{event.type}}: {log}",