blashttp request_batch_stream

Author: TheTechromancer

bbot/core/helpers/web/web.py

@@ -20,6 +20,23 @@
 log = logging.getLogger("bbot.core.helpers.web")
 
 
+async def iter_batch_results(stream):
+    """
+    Yield individual ``BatchResult`` objects from a ``request_batch_stream`` iterator.
+
+    The native blasthttp 0.4.0 iterator yields lists of ``BatchResult`` (drained in
+    chunks of 1000 / 200ms to amortize the Python↔Rust boundary). A future Python
+    wrapper is expected to unwrap these into individual items. This helper handles
+    both shapes so callers can write a single ``async for`` loop.
+    """
+    async for item in stream:
+        if isinstance(item, list):
+            for r in item:
+                yield r
+        else:
+            yield item
+
+
 class WebHelper:
     """
     Main utility class for managing HTTP operations in BBOT. Uses blasthttp (Rust) as the
@@ -297,23 +314,26 @@ async def request(self, *args, **kwargs):
                 log.trace(traceback.format_exc())
             raise
 
-    async def request_batch(self, urls, threads=10, **kwargs):
+    async def request_batch_stream(self, urls, threads=10, **kwargs):
         """
-        Request multiple URLs in parallel via blasthttp's native Rust batch engine.
+        Request multiple URLs in parallel via blasthttp's native Rust batch engine,
+        yielding each response as soon as it completes (completion order, not input
+        order).
 
         Applies the same header/cookie/proxy/timeout logic as ``request()`` — each
-        entry is translated into a ``blasthttp.BatchConfig`` and sent to Rust in one
-        shot.  Results are returned as a list (not streamed).
+        entry is translated into a ``blasthttp.BatchConfig`` and dispatched through
+        ``blasthttp.request_batch_stream``. A slow request no longer blocks faster
+        peers behind it, and Python work overlaps with in-flight HTTP I/O.
 
         Each entry in ``urls`` can be:
             - A plain URL string (uses shared ``**kwargs`` for all requests)
             - A ``(url, per_request_kwargs)`` tuple for per-request options
             - A ``(url, per_request_kwargs, tracker)`` tuple to attach arbitrary
-              tracking data that is returned alongside the response
+              tracking data that is yielded alongside the response
 
-        Returns:
-            When entries are plain strings: ``list[(url, response)]``
-            When any entry includes a tracker: ``list[(url, response, tracker)]``
+        Yields:
+            When entries are plain strings: ``(url, response)``
+            When any entry includes a tracker: ``(url, response, tracker)``
 
         Args:
             urls: URLs to visit — strings or ``(url, kwargs[, tracker])`` tuples.
@@ -324,15 +344,13 @@ async def request_batch(self, urls, threads=10, **kwargs):
         Examples:
             Simple (shared kwargs)::
 
-                results = await self.helpers.request_batch(urls, headers={"X-Test": "Test"})
-                for url, response in results:
+                async for url, response in self.helpers.request_batch_stream(urls, headers={"X-Test": "Test"}):
                     ...
 
             Per-request kwargs with tracker::
 
                 reqs = [("http://example.com", {"method": "POST"}, "my-tracker")]
-                results = await self.helpers.request_batch(reqs)
-                for url, response, tracker in results:
+                async for url, response, tracker in self.helpers.request_batch_stream(reqs):
                     ...
         """
         import blasthttp
@@ -354,33 +372,33 @@ async def request_batch(self, urls, threads=10, **kwargs):
                 entries.append((str(entry), kwargs, None))
 
         if not entries:
-            return []
+            return
+
+        # Build BatchConfig objects using the same logic as request().
+        # Map each config URL back to a queue of trackers so we can correlate
+        # completion-order results to original entries even when multiple entries
+        # share a URL.
+        from collections import deque
 
-        # Build BatchConfig objects using the same logic as request()
         configs = []
-        trackers = []
+        trackers_by_url = {}
         for url, req_kwargs, tracker in entries:
             url, method, blast_kwargs = self._build_blasthttp_kwargs(url, **req_kwargs)
             config = blasthttp.BatchConfig(url, **blast_kwargs)
             configs.append(config)
-            trackers.append(tracker)
-
-        # Send to Rust — all I/O happens here
-        batch_results = await self.client.request_batch(configs, concurrency=threads)
+            trackers_by_url.setdefault(config.url, deque()).append(tracker)
 
-        # Convert to (url, response[, tracker]) tuples
-        # Results are returned in the same order as configs
-        results = []
-        for i, br in enumerate(batch_results):
+        async for br in iter_batch_results(self.client.request_batch_stream(configs, concurrency=threads)):
             if br.response is not None:
                 response = BlasthttpResponse(br.response, request_url=br.url, method="GET")
             else:
                 response = None
             if has_tracker:
-                results.append((br.url, response, trackers[i]))
+                queue = trackers_by_url.get(br.url)
+                tracker = queue.popleft() if queue else None
+                yield br.url, response, tracker
             else:
-                results.append((br.url, response))
-        return results
+                yield br.url, response
 
     async def download(self, url, **kwargs):
         """

bbot/modules/git.py

@@ -24,7 +24,7 @@ async def handle_event(self, event):
             self.helpers.urljoin(base_url, ".git/config"),
             self.helpers.urljoin(f"{base_url}/", ".git/config"),
         }
-        for url, response in await self.helpers.request_batch(urls):
+        async for url, response in self.helpers.request_batch_stream(urls):
             text = getattr(response, "text", "")
             if not text:
                 text = ""

bbot/modules/http.py

@@ -4,6 +4,7 @@
 
 import blasthttp
 
+from bbot.core.helpers.web.web import iter_batch_results
 from bbot.modules.base import BaseModule
 
 
@@ -212,8 +213,12 @@ async def handle_batch(self, *events):
             )
             configs.append(config)
 
-        # blasthttp batch returns a native coroutine via pyo3-async-runtimes
-        results = await self.client.request_batch(configs, self.threads)
+        # Drain the streaming batch into a list — we need every result in hand
+        # before we can decide http/https suppression for paired OPEN_TCP_PORT probes.
+        # Python conversion still overlaps with in-flight HTTP I/O via the stream.
+        results = []
+        async for r in iter_batch_results(self.client.request_batch_stream(configs, concurrency=self.threads)):
+            results.append(r)
 
         # For OPEN_TCP_PORT probes, suppress redundant https when http already succeeded.
         # When probing an unknown port, we try both http:// and https://. If http works,

bbot/modules/iis_shortnames.py

@@ -142,7 +142,7 @@ async def solve_valid_chars(self, method, target, affirmative_status_code):
                 url = f"{target}{payload}{suffix}"
                 urls_and_kwargs.append((url, kwargs, (c, file_part)))
 
-        for url, response, (c, file_part) in await self.helpers.request_batch(urls_and_kwargs):
+        async for url, response, (c, file_part) in self.helpers.request_batch_stream(urls_and_kwargs):
             if response is not None:
                 if response.status_code == affirmative_status_code:
                     if file_part == "stem":
@@ -183,7 +183,7 @@ async def solve_shortname_recursive(
             kwargs = {"method": method}
             urls_and_kwargs.append((url, kwargs, c))
 
-        for url, response, c in await self.helpers.request_batch(urls_and_kwargs):
+        async for url, response, c in self.helpers.request_batch_stream(urls_and_kwargs):
             if response is not None:
                 if response.status_code == affirmative_status_code:
                     found_results = True

bbot/modules/ntlm.py

@@ -96,10 +96,9 @@ async def handle_event(self, event):
                 urls.add(f"{event.parsed_url.scheme}://{event.parsed_url.netloc}/{endpoint}")
 
         num_urls = len(urls)
-        results = await self.helpers.request_batch(
+        async for url, response in self.helpers.request_batch_stream(
             urls, headers=NTLM_test_header, allow_redirects=False, timeout=self.http_timeout
-        )
-        for url, response in results:
+        ):
             ntlm_resp = response.headers.get("WWW-Authenticate", "")
             if not ntlm_resp:
                 continue

bbot/modules/pgp.py

@@ -37,7 +37,7 @@ async def query(self, query):
         results = set()
         urls = self.config.get("search_urls", [])
         urls = [url.replace("<query>", self.helpers.quote(query)) for url in urls]
-        for url, response in await self.helpers.request_batch(urls):
+        async for url, response in self.helpers.request_batch_stream(urls):
             keyserver = self.helpers.urlparse(url).netloc
             if response is not None:
                 for email in await self.helpers.re.extract_emails(response.text):

bbot/modules/telerik.py

@@ -299,9 +299,8 @@ async def handle_event(self, event):
                 url = self.create_url(base_url, f"{dh}?dp=1")
                 urls[url] = dh
 
-            results = await self.helpers.request_batch(list(urls))
             fail_count = 0
-            for url, response in results:
+            async for url, response in self.helpers.request_batch_stream(list(urls)):
                 # cancel if we run into timeouts etc.
                 if response is None:
                     fail_count += 1

bbot/modules/templates/bucket.py

@@ -131,7 +131,9 @@ async def brute_buckets(self, buckets, permutations=False, omit_base=False):
                 for bucket_name in new_buckets:
                     url, kwargs = self.build_bucket_request(bucket_name, base_domain, region)
                     bucket_urls_kwargs.append((url, kwargs, (bucket_name, base_domain, region)))
-        for url, response, (bucket_name, base_domain, region) in await self.helpers.request_batch(bucket_urls_kwargs):
+        async for url, response, (bucket_name, base_domain, region) in self.helpers.request_batch_stream(
+            bucket_urls_kwargs
+        ):
             existent_bucket, tags = self._check_bucket_exists(bucket_name, response)
             if existent_bucket:
                 yield bucket_name, url, tags, num_buckets

bbot/modules/web_brute.py

@@ -3,6 +3,7 @@
 
 import blasthttp
 
+from bbot.core.helpers.web.web import iter_batch_results
 from bbot.modules.base import BaseModule
 
 
@@ -182,8 +183,9 @@ async def baseline_fuzz(self, url, exts=None, prefix="", suffix=""):
 
             canary_results = []
             canary_waf_count = 0
-            results = await self.blast_client.request_batch(canary_configs, 4, rate_limit=self.rate)
-            for result in results:
+            async for result in iter_batch_results(
+                self.blast_client.request_batch_stream(canary_configs, 4, rate_limit=self.rate)
+            ):
                 if result.success:
                     canary_results.append(self._batch_response_metrics(result.response))
                     if await self.helpers.yara.match(self.waf_yara_rules, result.response.body):
@@ -313,12 +315,13 @@ async def execute_fuzz(
 
             self.debug(f"Fuzzing {len(configs)} URLs for ext [{ext}]")
 
-            # Fire all requests via native blasthttp batch (Rust concurrency)
-            results = await self.blast_client.request_batch(configs, self.concurrency, rate_limit=self.rate)
-
-            # Index results by URL for ordered processing
+            # Fire all requests via native blasthttp batch (Rust concurrency).
+            # Stream results into a URL-keyed dict so we can re-process them in
+            # wordlist order (canary appended last) below.
             results_by_url = {}
-            for result in results:
+            async for result in iter_batch_results(
+                self.blast_client.request_batch_stream(configs, self.concurrency, rate_limit=self.rate)
+            ):
                 results_by_url[result.url] = result
 
             # Process in wordlist order so canary (appended last) is checked last
@@ -389,9 +392,13 @@ async def execute_fuzz(
                         proxy=proxy,
                     )
                 ]
-                canary_batch = await self.blast_client.request_batch(canary_configs, 1, rate_limit=self.rate)
-                if canary_batch and canary_batch[0].success:
-                    canary_metrics = self._batch_response_metrics(canary_batch[0].response)
+                canary_result = None
+                async for r in iter_batch_results(
+                    self.blast_client.request_batch_stream(canary_configs, 1, rate_limit=self.rate)
+                ):
+                    canary_result = r
+                if canary_result is not None and canary_result.success:
+                    canary_metrics = self._batch_response_metrics(canary_result.response)
                     if not self._is_baseline_match(canary_metrics, ext_filter):
                         self.verbose(
                             f"Would have reported {len(hits)} hit(s), but mid-scan baseline check failed. "

bbot/test/conftest.py

@@ -132,10 +132,11 @@ async def patched_request(self, *args, **kwargs):
             return result
         return await original_request(self, *args, **kwargs)
 
-    original_request_batch = WebHelper.request_batch
+    original_request_batch_stream = WebHelper.request_batch_stream
 
-    async def patched_request_batch(self, urls, threads=10, **kwargs):
+    async def patched_request_batch_stream(self, urls, threads=10, **kwargs):
         import blasthttp
+        from collections import deque
 
         # Run the real entry-parsing and config-building logic unmodified
         entries = []
@@ -154,38 +155,35 @@ async def patched_request_batch(self, urls, threads=10, **kwargs):
                 entries.append((str(entry), kwargs, None))
 
         if not entries:
-            return []
+            return
 
         configs = []
-        trackers = []
+        trackers_by_url = {}
         for url, req_kwargs, tracker in entries:
             url, method, blast_kwargs = self._build_blasthttp_kwargs(url, **req_kwargs)
             config = blasthttp.BatchConfig(url, **blast_kwargs)
             configs.append(config)
-            trackers.append(tracker)
-
-        # Route through mock's batch handler instead of Rust client directly
-        batch_results = await mock.handle_batch(self.client, configs, concurrency=threads)
+            trackers_by_url.setdefault(config.url, deque()).append(tracker)
 
         from bbot.core.helpers.web.blast_response import BlasthttpResponse
 
-        results = []
-        for i, br in enumerate(batch_results):
+        async for br in mock.handle_batch_stream(self.client, configs, concurrency=threads):
             if br.response is not None:
                 response = BlasthttpResponse(br.response, request_url=br.url, method="GET")
             else:
                 response = None
             if has_tracker:
-                results.append((br.url, response, trackers[i]))
+                queue = trackers_by_url.get(br.url)
+                tracker = queue.popleft() if queue else None
+                yield br.url, response, tracker
             else:
-                results.append((br.url, response))
-        return results
+                yield br.url, response
 
     WebHelper.request = patched_request
-    WebHelper.request_batch = patched_request_batch
+    WebHelper.request_batch_stream = patched_request_batch_stream
     yield mock
     WebHelper.request = original_request
-    WebHelper.request_batch = original_request_batch
+    WebHelper.request_batch_stream = original_request_batch_stream
 
 
 @pytest.fixture