Merge pull request #1555 from blacklanternsecurity/misc-small-bugfixes

Misc bugfixes

.flake8

@@ -0,0 +1,5 @@
+[flake8]
+select = F,E722
+ignore = F403,F405,F541
+per-file-ignores =
+    */__init__.py:F401,F403

.github/workflows/tests.yml

@@ -32,6 +32,8 @@ jobs:
     needs: lint
     runs-on: ubuntu-latest
     strategy:
+      # if one python version fails, let the others finish
+      fail-fast: false
       matrix:
         python-version: ["3.9", "3.10", "3.11", "3.12"]
     steps:

bbot/cli.py

@@ -219,7 +219,7 @@ def handle_keyboard_input(keyboard_input):
 
                 reader = asyncio.StreamReader()
                 protocol = asyncio.StreamReaderProtocol(reader)
-                await asyncio.get_event_loop().connect_read_pipe(lambda: protocol, sys.stdin)
+                await asyncio.get_running_loop().connect_read_pipe(lambda: protocol, sys.stdin)
 
                 # set stdout and stderr to blocking mode
                 # this is needed to prevent BlockingIOErrors in logging etc.
@@ -250,7 +250,7 @@ async def akeyboard_listen():
 
                 asyncio.create_task(akeyboard_listen())
 
-        await scan.async_start_without_generator()
+            await scan.async_start_without_generator()
 
         return True
 

bbot/core/event/base.py

@@ -712,7 +712,7 @@ def json(self, mode="json", siem_friendly=False):
         if self.scan:
             j["scan"] = self.scan.id
         # timestamp
-        j["timestamp"] = self.timestamp.timestamp()
+        j["timestamp"] = self.timestamp.isoformat()
         # parent event
         parent_id = self.parent_id
         if parent_id:
@@ -1604,7 +1604,7 @@ def event_from_json(j, siem_friendly=False):
         resolved_hosts = j.get("resolved_hosts", [])
         event._resolved_hosts = set(resolved_hosts)
 
-        event.timestamp = datetime.datetime.fromtimestamp(j["timestamp"])
+        event.timestamp = datetime.datetime.fromisoformat(j["timestamp"])
         event.scope_distance = j["scope_distance"]
         parent_id = j.get("parent", None)
         if parent_id is not None:

bbot/core/helpers/async_helpers.py

@@ -3,7 +3,6 @@
 import asyncio
 import logging
 from datetime import datetime
-from queue import Queue, Empty
 from cachetools import LRUCache
 from .misc import human_timedelta
 from contextlib import asynccontextmanager
@@ -91,37 +90,18 @@ def __str__(self):
             return f"{self.task_name} running for {running_for}"
 
 
-def async_to_sync_gen(async_gen):
-    # Queue to hold generated values
-    queue = Queue()
-
-    # Flag to indicate if the async generator is done
-    is_done = False
+def get_event_loop():
+    try:
+        return asyncio.get_running_loop()
+    except RuntimeError:
+        log.verbose("Starting new event loop")
+        return asyncio.new_event_loop()
 
-    # Function to run in the separate thread
-    async def runner():
-        nonlocal is_done
-        try:
-            async for value in async_gen:
-                queue.put(value)
-        finally:
-            is_done = True
 
-    def generator():
+def async_to_sync_gen(async_gen):
+    loop = get_event_loop()
+    try:
         while True:
-            # Try to get a value from the queue
-            try:
-                yield queue.get(timeout=0.1)
-            except Empty:
-                # If the queue is empty, check if the async generator is done
-                if is_done:
-                    break
-
-    from .process import BBOTThread
-
-    # Start the event loop in a separate thread
-    thread = BBOTThread(target=lambda: asyncio.run(runner()), daemon=True, custom_name="bbot async_to_sync_gen()")
-    thread.start()
-
-    # Return the generator
-    return generator()
+            yield loop.run_until_complete(async_gen.__anext__())
+    except StopAsyncIteration:
+        pass

bbot/core/helpers/dns/brute.py

@@ -50,7 +50,7 @@ async def dnsbrute(self, module, domain, subdomains, type=None):
             )
             return []
         else:
-            self.log.trace(f"{domain}: A is not in domain_wildcard_rdtypes:{domain_wildcard_rdtypes}")
+            self.log.debug(f"{domain}: A is not in domain_wildcard_rdtypes:{domain_wildcard_rdtypes}")
 
         canaries = self.gen_random_subdomains(self.num_canaries)
         canaries_list = list(canaries)

bbot/core/helpers/dns/dns.py

@@ -86,12 +86,22 @@ async def resolve_raw(self, query, **kwargs):
         return await self.run_and_return("resolve_raw", query=query, **kwargs)
 
     async def resolve_batch(self, queries, **kwargs):
-        async for _ in self.run_and_yield("resolve_batch", queries=queries, **kwargs):
-            yield _
+        agen = self.run_and_yield("resolve_batch", queries=queries, **kwargs)
+        while 1:
+            try:
+                yield await agen.__anext__()
+            except (StopAsyncIteration, GeneratorExit):
+                await agen.aclose()
+                break
 
     async def resolve_raw_batch(self, queries):
-        async for _ in self.run_and_yield("resolve_raw_batch", queries=queries):
-            yield _
+        agen = self.run_and_yield("resolve_raw_batch", queries=queries)
+        while 1:
+            try:
+                yield await agen.__anext__()
+            except (StopAsyncIteration, GeneratorExit):
+                await agen.aclose()
+                break
 
     @property
     def brute(self):

bbot/core/helpers/dns/engine.py

@@ -348,9 +348,10 @@ async def resolve_batch(self, queries, threads=10, **kwargs):
             ('evilcorp.com', {'2.2.2.2'})
         """
         tasks = {}
+        client_id = self.client_id_var.get()
 
         def new_task(query):
-            task = asyncio.create_task(self.resolve(query, **kwargs))
+            task = self.new_child_task(client_id, self.resolve(query, **kwargs))
             tasks[task] = query
 
         queries = list(queries)
@@ -360,9 +361,9 @@ def new_task(query):
 
         while tasks:  # While there are tasks pending
             # Wait for the first task to complete
-            done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+            finished = await self.finished_tasks(client_id)
 
-            for task in done:
+            for task in finished:
                 results = task.result()
                 query = tasks.pop(task)
 
@@ -374,9 +375,10 @@ def new_task(query):
 
     async def resolve_raw_batch(self, queries, threads=10):
         tasks = {}
+        client_id = self.client_id_var.get()
 
         def new_task(query, rdtype):
-            task = asyncio.create_task(self.resolve_raw(query, type=rdtype))
+            task = self.new_child_task(client_id, self.resolve_raw(query, type=rdtype))
             tasks[task] = (query, rdtype)
 
         queries = list(queries)
@@ -386,9 +388,9 @@ def new_task(query, rdtype):
 
         while tasks:  # While there are tasks pending
             # Wait for the first task to complete
-            done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+            finished = await self.finished_tasks(client_id)
 
-            for task in done:
+            for task in finished:
                 answers, errors = task.result()
                 query, rdtype = tasks.pop(task)
                 for answer in answers:

bbot/core/helpers/helper.py

@@ -1,5 +1,4 @@
 import os
-import asyncio
 import logging
 from pathlib import Path
 import multiprocessing as mp
@@ -15,6 +14,7 @@
 from .interactsh import Interactsh
 from ...scanner.target import Target
 from .depsinstaller import DepsInstaller
+from .async_helpers import get_event_loop
 
 log = logging.getLogger("bbot.core.helpers")
 
@@ -84,13 +84,19 @@ def __init__(self, preset):
         self._cloud = None
 
         self.re = RegexHelper(self)
-        self.dns = DNSHelper(self)
+        self._dns = None
         self._web = None
         self.config_aware_validators = self.validators.Validators(self)
         self.depsinstaller = DepsInstaller(self)
         self.word_cloud = WordCloud(self)
         self.dummy_modules = {}
 
+    @property
+    def dns(self):
+        if self._dns is None:
+            self._dns = DNSHelper(self)
+        return self._dns
+
     @property
     def web(self):
         if self._web is None:
@@ -170,7 +176,7 @@ def loop(self):
         Get the current event loop
         """
         if self._loop is None:
-            self._loop = asyncio.get_running_loop()
+            self._loop = get_event_loop()
         return self._loop
 
     def run_in_executor(self, callback, *args, **kwargs):

bbot/core/helpers/process.py

@@ -16,6 +16,8 @@ class BBOTThread(threading.Thread):
 
     def __init__(self, *args, **kwargs):
         self.custom_name = kwargs.pop("custom_name", self.default_name)
+        if "daemon" not in kwargs:
+            kwargs["daemon"] = True
         super().__init__(*args, **kwargs)
 
     def run(self):

bbot/core/helpers/web/engine.py

@@ -32,7 +32,7 @@ def __init__(self, socket_path, target, config={}):
         self.target = target
         self.config = config
         self.web_config = self.config.get("web", {})
-        self.http_debug = self.web_config.get("http_debug", False)
+        self.http_debug = self.web_config.get("debug", False)
         self._ssl_context_noverify = None
         self.web_client = self.AsyncClient(persist_cookies=False)
 
@@ -85,36 +85,38 @@ async def request(self, *args, **kwargs):
 
     async def request_batch(self, urls, *args, threads=10, **kwargs):
         tasks = {}
+        client_id = self.client_id_var.get()
 
         urls = list(urls)
 
         def new_task():
             if urls:
                 url = urls.pop(0)
-                task = asyncio.create_task(self.request(url, *args, **kwargs))
+                task = self.new_child_task(client_id, self.request(url, *args, **kwargs))
                 tasks[task] = url
 
         for _ in range(threads):  # Start initial batch of tasks
             new_task()
 
         while tasks:  # While there are tasks pending
             # Wait for the first task to complete
-            done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+            finished = await self.finished_tasks(client_id)
 
-            for task in done:
+            for task in finished:
                 response = task.result()
                 url = tasks.pop(task)
                 yield (url, response)
                 new_task()
 
     async def request_custom_batch(self, urls_and_kwargs, threads=10):
         tasks = {}
+        client_id = self.client_id_var.get()
         urls_and_kwargs = list(urls_and_kwargs)
 
         def new_task():
             if urls_and_kwargs:  # Ensure there are args to process
                 url, kwargs, custom_tracker = urls_and_kwargs.pop(0)
-                task = asyncio.create_task(self.request(url, **kwargs))
+                task = self.new_child_task(client_id, self.request(url, **kwargs))
                 tasks[task] = (url, kwargs, custom_tracker)
 
         for _ in range(threads):  # Start initial batch of tasks
@@ -135,7 +137,10 @@ async def download(self, url, **kwargs):
         filename = kwargs.pop("filename")
         raise_error = kwargs.get("raise_error", False)
         try:
-            content, response = await self.stream_request(url, **kwargs)
+            result = await self.stream_request(url, **kwargs)
+            if result is None:
+                raise httpx.HTTPError(f"No response from {url}")
+            content, response = result
             log.debug(f"Download result: HTTP {response.status_code}")
             response.raise_for_status()
             with open(filename, "wb") as f:

bbot/core/helpers/web/web.py

@@ -124,8 +124,13 @@ async def request_batch(self, urls, *args, **kwargs):
             >>>     if response is not None and response.status_code == 200:
             >>>         self.hugesuccess(response)
         """
-        async for _ in self.run_and_yield("request_batch", urls, *args, **kwargs):
-            yield _
+        agen = self.run_and_yield("request_batch", urls, *args, **kwargs)
+        while 1:
+            try:
+                yield await agen.__anext__()
+            except (StopAsyncIteration, GeneratorExit):
+                await agen.aclose()
+                break
 
     async def request_custom_batch(self, urls_and_kwargs):
         """
@@ -149,8 +154,13 @@ async def request_custom_batch(self, urls_and_kwargs):
             >>>     if response is not None and response.status_code == 200:
             >>>         self.hugesuccess(response)
         """
-        async for _ in self.run_and_yield("request_custom_batch", urls_and_kwargs):
-            yield _
+        agen = self.run_and_yield("request_custom_batch", urls_and_kwargs)
+        while 1:
+            try:
+                yield await agen.__anext__()
+            except (StopAsyncIteration, GeneratorExit):
+                await agen.aclose()
+                break
 
     async def download(self, url, **kwargs):
         """
@@ -291,6 +301,9 @@ async def api_page_iter(self, url, page_size=100, json=True, next_key=None, **re
             else:
                 new_url = url.format(page=page, page_size=page_size, offset=offset)
             result = await self.request(new_url, **requests_kwargs)
+            if result is None:
+                log.verbose(f"api_page_iter() got no response for {url}")
+                break
             try:
                 if json:
                     result = result.json()

bbot/modules/azure_tenant.py

@@ -93,7 +93,7 @@ async def query(self, domain):
         self.debug(f"Retrieving tenant domains at {url}")
 
         autodiscover_task = self.helpers.create_task(
-            self.helpers.request(url, method="POST", headers=headers, data=data)
+            self.helpers.request(url, method="POST", headers=headers, content=data)
         )
         openid_url = f"https://login.windows.net/{domain}/.well-known/openid-configuration"
         openid_task = self.helpers.create_task(self.helpers.request(openid_url))

bbot/modules/base.py

@@ -530,7 +530,8 @@ def num_incoming_events(self):
 
     def start(self):
         self._tasks = [
-            asyncio.create_task(self._worker(), name=f"{self.name}._worker()") for _ in range(self.module_threads)
+            asyncio.create_task(self._worker(), name=f"{self.scan.name}.{self.name}._worker()")
+            for _ in range(self.module_threads)
         ]
 
     async def _setup(self):

bbot/modules/dnsbrute_mutations.py

@@ -29,7 +29,7 @@ async def setup(self):
         return True
 
     async def handle_event(self, event):
-        # here we don't brute-force, we just add the subdomain to our end-of-scan TODO
+        # here we don't brute-force, we just add the subdomain to our end-of-scan
         self.add_found(event)
 
     def add_found(self, event):
@@ -103,6 +103,13 @@ def add_mutation(m):
                     ):
                         add_mutation(subdomain)
 
+                    # skip if there's hardly any mutations
+                    if len(mutations) < 10:
+                        self.debug(
+                            f"Skipping {len(mutations):,} mutations against {domain} because there are less than 10"
+                        )
+                        break
+
                     if mutations:
                         self.info(f"Trying {len(mutations):,} mutations against {domain} ({i+1}/{len(trimmed_found)})")
                         results = await self.helpers.dns.brute(self, query, mutations)