Strip HTTP response bodies after all modules finish processing

bbot/core/event/base.py

@@ -164,6 +164,8 @@ class BaseEvent:
         "dns_resolve_distance",
         # Host metadata (cloud providers, ASN, whois, etc.)
         "_host_metadata",
+        # Memory management
+        "_module_consumers",
         # Public attributes
         "module",
         "scan",
@@ -224,6 +226,7 @@ def __init__(
         self.dns_children = {}
         self.raw_dns_records = {}
         self._discovery_context = ""
+        self._module_consumers = 0
 
         # for creating one-off events without enforcing parent requirement
         self._dummy = _dummy
@@ -689,6 +692,25 @@ def get_parents(self, omit=False, include_self=False):
             e = parent
         return parents
 
+    def _minimize(self):
+        """
+        Called when a module is done processing this event.
+
+        Decrements the consumer count. When no modules are left waiting to
+        process this event, heavy payload data (e.g. HTTP response bodies)
+        is stripped to free memory.
+
+        The event object stays alive (for parent-chain references, etc.)
+        but large fields like HTTP response bodies and raw headers are removed.
+
+        So basically, the parent becomes dead inside for the sake of the children.
+        Just like real life.
+        """
+        self._module_consumers = max(0, self._module_consumers - 1)
+        if self._module_consumers <= 0 and isinstance(self._data, dict):
+            self._data.pop("body", None)
+            self._data.pop("raw_header", None)
+
     def clone(self):
         # Create a shallow copy of the event first
         cloned_event = copy(self)

bbot/modules/base.py

@@ -494,6 +494,9 @@ async def _handle_batch(self):
                     await self.run_task(self.handle_batch(*events), context, n=len(events))
                 except asyncio.CancelledError:
                     self.debug(f"{context} was cancelled")
+                finally:
+                    for event in events:
+                        event._minimize()
                 self.verbose(f"Finished handling batch of {len(events):,} events")
         if finish:
             context = f"{self.name}.finish()"
@@ -663,11 +666,14 @@ async def _events_waiting(self, batch_size=None):
                 if acceptable:
                     if event.type == "FINISHED":
                         finish = True
+                        event._minimize()
                     else:
                         events.append(event)
                         self.scan.stats.event_consumed(event, self)
-                elif reason:
-                    self.debug(f"Not accepting {event} because {reason}")
+                else:
+                    event._minimize()
+                    if reason:
+                        self.debug(f"Not accepting {event} because {reason}")
             except asyncio.queues.QueueEmpty:
                 break
         return events, finish
@@ -761,28 +767,31 @@ async def _worker(self):
                         except asyncio.queues.QueueEmpty:
                             continue
                         self.debug(f"Got {event} from {getattr(event, 'module', 'unknown_module')}")
-                        async with self._task_counter.count(f"event_postcheck({event})"):
-                            acceptable, reason = await self._event_postcheck(event)
-                        if acceptable:
-                            if event.type == "FINISHED":
-                                context = f"{self.name}.finish()"
-                                try:
-                                    await self.run_task(self.finish(), context)
-                                except asyncio.CancelledError:
-                                    self.debug(f"{context} was cancelled")
-                                    continue
+                        try:
+                            async with self._task_counter.count(f"event_postcheck({event})"):
+                                acceptable, reason = await self._event_postcheck(event)
+                            if acceptable:
+                                if event.type == "FINISHED":
+                                    context = f"{self.name}.finish()"
+                                    try:
+                                        await self.run_task(self.finish(), context)
+                                    except asyncio.CancelledError:
+                                        self.debug(f"{context} was cancelled")
+                                        continue
+                                else:
+                                    context = f"{self.name}.handle_event({event})"
+                                    self.scan.stats.event_consumed(event, self)
+                                    self.debug(f"Handling {event}")
+                                    try:
+                                        await self.run_task(self.handle_event(event), context)
+                                    except asyncio.CancelledError:
+                                        self.debug(f"{context} was cancelled")
+                                        continue
+                                    self.debug(f"Finished handling {event}")
                             else:
-                                context = f"{self.name}.handle_event({event})"
-                                self.scan.stats.event_consumed(event, self)
-                                self.debug(f"Handling {event}")
-                                try:
-                                    await self.run_task(self.handle_event(event), context)
-                                except asyncio.CancelledError:
-                                    self.debug(f"{context} was cancelled")
-                                    continue
-                                self.debug(f"Finished handling {event}")
-                        else:
-                            self.debug(f"Not accepting {event} because {reason}")
+                                self.debug(f"Not accepting {event} because {reason}")
+                        finally:
+                            event._minimize()
         except asyncio.CancelledError:
             # this trace was used for debugging leaked CancelledErrors from inside httpx
             # self.log.trace("Worker cancelled")
@@ -1022,6 +1031,7 @@ async def queue_event(self, event):
                 self.debug(f"Queueing {event} because {reason}")
             try:
                 self.incoming_event_queue.put_nowait(event)
+                event._module_consumers += 1
                 async with self.event_received:
                     self.event_received.notify()
                 if event.type != "FINISHED":

bbot/scanner/manager.py

@@ -271,3 +271,7 @@ async def forward_event(self, event, kwargs):
             # don't distribute events to intercept modules
             if not mod._intercept:
                 await mod.queue_event(event)
+
+        # if no module accepted this event, minimize it now
+        if event._module_consumers <= 0:
+            event._minimize()