From 60d1af1f0b91a2b34fed45b99650131e81a86345 Mon Sep 17 00:00:00 2001
From: TheTechromancer <thetechromancer@protonmail.com>
Date: Wed, 20 Sep 2023 12:05:24 -0400
Subject: [PATCH] utilize target for source event correlation in massdns

---
 bbot/core/helpers/helper.py   |  2 +-
 bbot/modules/deadly/nuclei.py |  2 +-
 bbot/modules/massdns.py       | 15 +++------------
 bbot/modules/nmap.py          |  2 +-
 4 files changed, 6 insertions(+), 15 deletions(-)

diff --git a/bbot/core/helpers/helper.py b/bbot/core/helpers/helper.py
index b052a63fb..241ec15d8 100644
--- a/bbot/core/helpers/helper.py
+++ b/bbot/core/helpers/helper.py
@@ -70,7 +70,7 @@ def clean_old_scans(self):
         _filter = lambda x: x.is_dir() and self.regexes.scan_name_regex.match(x.name)
         self.clean_old(self.scans_dir, keep=self.keep_old_scans, filter=_filter)
 
-    def make_target(self, events):
+    def make_target(self, *events):
         return Target(self.scan, *events)
 
     @property
diff --git a/bbot/modules/deadly/nuclei.py b/bbot/modules/deadly/nuclei.py
index 0ea880b31..35eb6aabf 100644
--- a/bbot/modules/deadly/nuclei.py
+++ b/bbot/modules/deadly/nuclei.py
@@ -128,7 +128,7 @@ async def setup(self):
         return True
 
     async def handle_batch(self, *events):
-        temp_target = self.helpers.make_target(events)
+        temp_target = self.helpers.make_target(*events)
         nuclei_input = [str(e.data) for e in events]
         async for severity, template, host, url, name, extracted_results in self.execute_nuclei(nuclei_input):
             # this is necessary because sometimes nuclei is inconsistent about the data returned in the host field
diff --git a/bbot/modules/massdns.py b/bbot/modules/massdns.py
index 2ec1031af..ed54fc655 100644
--- a/bbot/modules/massdns.py
+++ b/bbot/modules/massdns.py
@@ -65,7 +65,7 @@ class massdns(crobat):
     async def setup(self):
         self.found = dict()
         self.mutations_tried = set()
-        self.source_events = dict()
+        self.source_events = self.helpers.make_target()
         self.subdomain_file = await self.helpers.wordlist(self.config.get("wordlist"))
         self.max_resolvers = self.config.get("max_resolvers", 1000)
         self.max_mutations = self.config.get("max_mutations", 500)
@@ -94,9 +94,7 @@ async def filter_event(self, event):
 
     async def handle_event(self, event):
         query = self.make_query(event)
-        h = hash(query)
-        if not h in self.source_events:
-            self.source_events[h] = event
+        self.source_events.add_target(event)
 
         self.info(f"Brute-forcing subdomains for {query} (source: {event.data})")
         for hostname in await self.massdns(query, self.helpers.read_file(self.subdomain_file)):
@@ -354,7 +352,7 @@ def add_mutation(_domain_hash, m):
                         self.info(f"Trying {len(mutations):,} mutations against {domain} ({i+1}/{len(found)})")
                         results = list(await self.massdns(query, mutations))
                         for hostname in results:
-                            source_event = self.get_source_event(hostname)
+                            source_event = self.source_events.get(hostname)
                             if source_event is None:
                                 self.warning(f"Could not correlate source event from: {hostname}")
                                 source_event = self.scan.root_event
@@ -395,10 +393,3 @@ def gen_random_subdomains(self, n=50):
             yield subdomain
         for _ in range(5):
             yield self.helpers.rand_string(length=8, digits=False)
-
-    def get_source_event(self, hostname):
-        for p in self.helpers.domain_parents(hostname):
-            try:
-                return self.source_events[hash(p)]
-            except KeyError:
-                continue
diff --git a/bbot/modules/nmap.py b/bbot/modules/nmap.py
index d0671d16b..e2900ab1b 100644
--- a/bbot/modules/nmap.py
+++ b/bbot/modules/nmap.py
@@ -35,7 +35,7 @@ async def setup(self):
         return True
 
     async def handle_batch(self, *events):
-        target = self.helpers.make_target(events)
+        target = self.helpers.make_target(*events)
         targets = list(set(str(e.data) for e in events))
         command, output_file = self.construct_command(targets)
         try: