Fix non-functional DNS brute-force canary check

bbot/core/helpers/dns/brute.py

@@ -49,16 +49,16 @@ async def dnsbrute(self, module, domain, subdomains, type=None):
 
         wildcard_domains = await self.parent_helper.dns.is_wildcard_domain(domain, (type, "CNAME"))
         wildcard_rdtypes = set()
-        for domain, rdtypes in wildcard_domains.items():
+        for wildcard_domain, rdtypes in wildcard_domains.items():
             wildcard_rdtypes.update(rdtypes)
         if wildcard_domains:
             self.log.hugewarning(
                 f"Aborting massdns on {domain} because it's a wildcard domain ({','.join(sorted(wildcard_rdtypes))})"
             )
             return []
 
-        canaries = self.gen_random_subdomains(self.num_canaries)
-        canaries_list = list(canaries)
+        canaries_list = list(self.gen_random_subdomains(self.num_canaries))
+        canary_set = set(canaries_list)
         canaries_pre = canaries_list[: int(self.num_canaries / 2)]
         canaries_post = canaries_list[int(self.num_canaries / 2) :]
         # sandwich subdomains between canaries
@@ -67,8 +67,8 @@ async def dnsbrute(self, module, domain, subdomains, type=None):
         results = []
         canaries_triggered = []
         async for hostname, ip, rdtype in self._massdns(module, domain, subdomains, rdtype=type):
-            sub = hostname.split(domain)[0]
-            if sub in canaries:
+            sub = hostname.split(domain)[0].rstrip(".")
+            if sub in canary_set:
                 canaries_triggered.append(sub)
             else:
                 results.append(hostname)

bbot/test/test_step_2/module_tests/test_module_dnsbrute.py

@@ -95,3 +95,52 @@ def check(self, module_test, events):
         assert 1 == len(
             [e for e in events if e.data == "asdf.blacklanternsecurity.com" and str(e.module) == "dnsbrute"]
         )
+
+
+class TestDnsbruteCanaryCheck(ModuleTestBase):
+    """Test that the canary check correctly aborts brute-forcing on wildcard domains.
+
+    Simulates a wildcard domain by making massdns return a result for every input
+    subdomain, including the canary subdomains. The canary check should detect
+    this and return no results.
+    """
+
+    module_name = "dnsbrute"
+    subdomain_wordlist = tempwordlist(["www", "mail"])
+    config_overrides = {"modules": {"dnsbrute": {"wordlist": str(subdomain_wordlist)}}}
+
+    async def setup_after_prep(self, module_test):
+        old_run_live = module_test.scan.helpers.run_live
+
+        async def new_run_live(*command, check=False, text=True, **kwargs):
+            if "massdns" in command[:2]:
+                # simulate a wildcard domain: return a result for EVERY input subdomain
+                _input = [l async for l in kwargs["input"]]
+                for subdomain in _input:
+                    hostname = subdomain.strip()
+                    if hostname:
+                        yield (
+                            '{"name": "'
+                            + hostname
+                            + '.", "type": "A", "class": "IN", "status": "NOERROR", "data": {"answers": [{"ttl": 86400, "type": "A", "class": "IN", "name": "'
+                            + hostname
+                            + '.", "data": "1.2.3.4"}]}, "resolver": "195.226.187.130:53", "proto": "UDP"}'
+                        )
+            else:
+                async for _ in old_run_live(*command, check=False, text=True, **kwargs):
+                    yield _
+
+        module_test.monkeypatch.setattr(module_test.scan.helpers, "run_live", new_run_live)
+
+        await module_test.mock_dns(
+            {
+                "blacklanternsecurity.com": {"A": ["4.3.2.1"]},
+            }
+        )
+
+    def check(self, module_test, events):
+        # canary check should have aborted, so no DNS_NAME events from dnsbrute
+        dnsbrute_events = [e for e in events if e.type == "DNS_NAME" and str(e.module) == "dnsbrute"]
+        assert len(dnsbrute_events) == 0, (
+            f"Expected no results from dnsbrute (canary check should abort), but got {len(dnsbrute_events)}: {[e.data for e in dnsbrute_events]}"
+        )