dns mock cleanup

Author: TheTechromancer

bbot/core/helpers/dns.py

@@ -1040,28 +1040,3 @@ def _get_dummy_module(self, name):
             dummy_module = self.parent_helper._make_dummy_module(name=name, _type="DNS")
             self._dummy_modules[name] = dummy_module
         return dummy_module
-
-    def mock_dns(self, dns_dict):
-        if self._orig_resolve_raw is None:
-            self._orig_resolve_raw = self.resolve_raw
-
-        async def mock_resolve_raw(query, **kwargs):
-            results = []
-            errors = []
-            types = self._parse_rdtype(kwargs.get("type", ["A", "AAAA"]))
-            for t in types:
-                with suppress(KeyError):
-                    results += self._mock_table[(query, t)]
-            return results, errors
-
-        for (query, rdtype), answers in dns_dict.items():
-            if isinstance(answers, str):
-                answers = [answers]
-            for answer in answers:
-                rdata = dns.rdata.from_text("IN", rdtype, answer)
-                try:
-                    self._mock_table[(query, rdtype)].append((rdtype, rdata))
-                except KeyError:
-                    self._mock_table[(query, rdtype)] = [(rdtype, [rdata])]
-
-        self.resolve_raw = mock_resolve_raw

bbot/test/bbot_fixtures.py

@@ -1,4 +1,5 @@
 import os
+import dns
 import sys
 import pytest
 import asyncio  # noqa
@@ -252,3 +253,67 @@ def install_all_python_deps():
     for module in module_loader.preloaded().values():
         deps_pip.update(set(module.get("deps", {}).get("pip", [])))
     subprocess.run([sys.executable, "-m", "pip", "install"] + list(deps_pip))
+
+
+class MockResolver:
+    import dns
+
+    def __init__(self, mock_data=None):
+        self.mock_data = mock_data if mock_data else {}
+        self.nameservers = ["127.0.0.1"]
+
+    async def resolve_address(self, ipaddr, *args, **kwargs):
+        modified_kwargs = {}
+        modified_kwargs.update(kwargs)
+        modified_kwargs["rdtype"] = "PTR"
+        return await self.resolve(str(dns.reversename.from_address(ipaddr)), *args, **modified_kwargs)
+
+    def create_dns_response(self, query_name, rdtype):
+        query_name = query_name.strip(".")
+        answers = self.mock_data.get(query_name, {}).get(rdtype, [])
+        if not answers:
+            raise self.dns.resolver.NXDOMAIN(f"No answer found for {query_name} {rdtype}")
+
+        message_text = f"""id 1234
+opcode QUERY
+rcode NOERROR
+flags QR AA RD
+;QUESTION
+{query_name}. IN {rdtype}
+;ANSWER"""
+        for answer in answers:
+            message_text += f"\n{query_name}. 1 IN {rdtype} {answer}"
+
+        message_text += "\n;AUTHORITY\n;ADDITIONAL\n"
+        message = self.dns.message.from_text(message_text)
+        return message
+
+    async def resolve(self, query_name, rdtype=None):
+        if rdtype is None:
+            rdtype = "A"
+        elif isinstance(rdtype, str):
+            rdtype = rdtype.upper()
+        else:
+            rdtype = str(rdtype.name).upper()
+
+        domain_name = self.dns.name.from_text(query_name)
+        rdtype_obj = self.dns.rdatatype.from_text(rdtype)
+
+        if "_NXDOMAIN" in self.mock_data and query_name in self.mock_data["_NXDOMAIN"]:
+            # Simulate the NXDOMAIN exception
+            raise self.dns.resolver.NXDOMAIN
+
+        try:
+            response = self.create_dns_response(query_name, rdtype)
+            answer = self.dns.resolver.Answer(domain_name, rdtype_obj, self.dns.rdataclass.IN, response)
+            return answer
+        except self.dns.resolver.NXDOMAIN:
+            return []
+
+
+@pytest.fixture()
+def mock_dns():
+    def _mock_dns(scan, mock_data):
+        scan.helpers.dns.resolver = MockResolver(mock_data)
+
+    return _mock_dns

bbot/test/conftest.py

@@ -179,70 +179,3 @@ def proxy_server():
     # Stop the server.
     server.shutdown()
     server_thread.join()
-
-
-class MockResolver:
-    import dns
-
-    def __init__(self, mock_data=None):
-        self.mock_data = mock_data if mock_data else {}
-        self.nameservers = ["127.0.0.1"]
-
-    async def resolve_address(self, host):
-        try:
-            from dns.asyncresolver import resolve_address
-
-            result = await resolve_address(host)
-            return result
-        except ImportError:
-            raise ImportError("dns.asyncresolver.Resolver.resolve_address not found")
-
-    def create_dns_response(self, query_name, rdtype):
-        answers = self.mock_data.get(query_name, {}).get(rdtype, [])
-        if not answers:
-            raise self.dns.resolver.NXDOMAIN(f"No answer found for {query_name} {rdtype}")
-
-        message_text = f"""id 1234
-opcode QUERY
-rcode NOERROR
-flags QR AA RD
-;QUESTION
-{query_name}. IN {rdtype}
-;ANSWER"""
-        for answer in answers:
-            message_text += f"\n{query_name}. 1 IN {rdtype} {answer}"
-
-        message_text += "\n;AUTHORITY\n;ADDITIONAL\n"
-        message = self.dns.message.from_text(message_text)
-        return message
-
-    async def resolve(self, query_name, rdtype=None):
-        if rdtype is None:
-            rdtype = "A"
-        elif isinstance(rdtype, str):
-            rdtype = rdtype.upper()
-        else:
-            rdtype = str(rdtype.name).upper()
-
-        domain_name = self.dns.name.from_text(query_name)
-        rdtype_obj = self.dns.rdatatype.from_text(rdtype)
-
-        if "_NXDOMAIN" in self.mock_data and query_name in self.mock_data["_NXDOMAIN"]:
-            # Simulate the NXDOMAIN exception
-            raise self.dns.resolver.NXDOMAIN
-
-        try:
-            response = self.create_dns_response(query_name, rdtype)
-            answer = self.dns.resolver.Answer(domain_name, rdtype_obj, self.dns.rdataclass.IN, response)
-            return answer
-        except self.dns.resolver.NXDOMAIN:
-            return []
-
-
-@pytest.fixture()
-def configure_mock_resolver(monkeypatch):
-    def _configure(mock_data):
-        mock_resolver = MockResolver(mock_data)
-        return mock_resolver
-
-    return _configure

bbot/test/test_step_1/test_dns.py

@@ -2,7 +2,7 @@
 
 
 @pytest.mark.asyncio
-async def test_dns(bbot_scanner, bbot_config):
+async def test_dns(bbot_scanner, bbot_config, mock_dns):
     scan = bbot_scanner("1.1.1.1", config=bbot_config)
     helpers = scan.helpers
 
@@ -85,8 +85,12 @@ async def test_dns(bbot_scanner, bbot_config):
     dns_config = OmegaConf.create({"dns_resolution": True})
     dns_config = OmegaConf.merge(bbot_config, dns_config)
     scan2 = bbot_scanner("evilcorp.com", config=dns_config)
-    scan2.helpers.dns.mock_dns(
-        {("evilcorp.com", "TXT"): '"v=spf1 include:cloudprovider.com ~all"', ("cloudprovider.com", "A"): "1.2.3.4"}
+    mock_dns(
+        scan2,
+        {
+            "evilcorp.com": {"TXT": ['"v=spf1 include:cloudprovider.com ~all"']},
+            "cloudprovider.com": {"A": ["1.2.3.4"]},
+        },
     )
     events = [e async for e in scan2.async_start()]
     assert 1 == len(

bbot/test/test_step_1/test_manager_deduplication.py

@@ -3,7 +3,7 @@
 
 
 @pytest.mark.asyncio
-async def test_manager_deduplication(bbot_config, bbot_scanner):
+async def test_manager_deduplication(bbot_config, bbot_scanner, mock_dns):
 
     class DefaultModule(BaseModule):
         _name = "default_module"
@@ -62,7 +62,7 @@ async def do_scan(*args, _config={}, _dns_mock={}, scan_callback=None, **kwargs)
         scan.modules["per_hostport_only"] = per_hostport_only
         scan.modules["per_domain_only"] = per_domain_only
         if _dns_mock:
-            scan.helpers.dns.mock_dns(_dns_mock)
+            mock_dns(scan, _dns_mock)
         if scan_callback is not None:
             scan_callback(scan)
         return (
@@ -76,12 +76,12 @@ async def do_scan(*args, _config={}, _dns_mock={}, scan_callback=None, **kwargs)
         )
 
     dns_mock_chain = {
-        ("default_module.test.notreal", "A"): "127.0.0.3",
-        ("everything_module.test.notreal", "A"): "127.0.0.4",
-        ("no_suppress_dupes.test.notreal", "A"): "127.0.0.5",
-        ("accept_dupes.test.notreal", "A"): "127.0.0.6",
-        ("per_hostport_only.test.notreal", "A"): "127.0.0.7",
-        ("per_domain_only.test.notreal", "A"): "127.0.0.8",
+        "default_module.test.notreal": {"A": ["127.0.0.3"]},
+        "everything_module.test.notreal": {"A": ["127.0.0.4"]},
+        "no_suppress_dupes.test.notreal": {"A": ["127.0.0.5"]},
+        "accept_dupes.test.notreal": {"A": ["127.0.0.6"]},
+        "per_hostport_only.test.notreal": {"A": ["127.0.0.7"]},
+        "per_domain_only.test.notreal": {"A": ["127.0.0.8"]},
     }
 
     # dns search distance = 1, report distance = 0

bbot/test/test_step_1/test_manager_scope_accuracy.py

@@ -777,7 +777,7 @@ def custom_setup(scan):
 
 
 @pytest.mark.asyncio
-async def test_manager_blacklist(bbot_config, bbot_scanner, bbot_httpserver, caplog):
+async def test_manager_blacklist(bbot_config, bbot_scanner, bbot_httpserver, caplog, mock_dns):
 
     bbot_httpserver.expect_request(uri="/").respond_with_data(response_data="<a href='http://www-prod.test.notreal:8888'/><a href='http://www-dev.test.notreal:8888'/>")
 
@@ -791,9 +791,9 @@ async def test_manager_blacklist(bbot_config, bbot_scanner, bbot_httpserver, cap
         whitelist=["127.0.0.0/29", "test.notreal"],
         blacklist=["127.0.0.64/29"],
     )
-    scan.helpers.dns.mock_dns({
-        ("www-prod.test.notreal", "A"): "127.0.0.66",
-        ("www-dev.test.notreal", "A"): "127.0.0.22",
+    mock_dns(scan, {
+        "www-prod.test.notreal": {"A": ["127.0.0.66"]},
+        "www-dev.test.notreal": {"A": ["127.0.0.22"]},
     })
 
     events = [e async for e in scan.async_start()]

bbot/test/test_step_1/test_modules_basic.py

@@ -303,31 +303,51 @@ async def test_modules_basic_perdomainonly(scan, helpers, events, bbot_config, b
 
 
 @pytest.mark.asyncio
-async def test_modules_basic_stats(helpers, events, bbot_config, bbot_scanner, httpx_mock, monkeypatch):
+async def test_modules_basic_stats(helpers, events, bbot_config, bbot_scanner, httpx_mock, monkeypatch, mock_dns):
     from bbot.modules.base import BaseModule
 
     class dummy(BaseModule):
         _name = "dummy"
         watched_events = ["*"]
 
         async def handle_event(self, event):
+            # quick emit events like FINDINGS behave differently than normal ones
+            # hosts are not speculated from them
             await self.emit_event(
                 {"host": "www.evilcorp.com", "url": "http://www.evilcorp.com", "description": "asdf"}, "FINDING", event
             )
+            await self.emit_event("https://asdf.evilcorp.com", "URL", event, tags=["status-200"])
 
     scan = bbot_scanner(
         "evilcorp.com",
+        modules=["speculate"],
         config=bbot_config,
         force_start=True,
     )
-    scan.helpers.dns.mock_dns({("evilcorp.com", "A"): "127.0.254.1", ("www.evilcorp.com", "A"): "127.0.254.2"})
+    mock_dns(
+        scan,
+        {
+            "evilcorp.com": {"A": ["127.0.254.1"]},
+            "www.evilcorp.com": {"A": ["127.0.254.2"]},
+            "asdf.evilcorp.com": {"A": ["127.0.254.3"]},
+        },
+    )
 
     scan.modules["dummy"] = dummy(scan)
     events = [e async for e in scan.async_start()]
 
-    assert len(events) == 3
+    assert len(events) == 6
+
+    assert scan.stats.events_emitted_by_type == {
+        "SCAN": 1,
+        "DNS_NAME": 3,
+        "URL": 1,
+        "FINDING": 1,
+        "ORG_STUB": 1,
+        "OPEN_TCP_PORT": 1,
+    }
 
-    assert set(scan.stats.module_stats) == {"dummy", "python", "TARGET"}
+    assert set(scan.stats.module_stats) == {"host", "speculate", "python", "dummy", "TARGET"}
 
     target_stats = scan.stats.module_stats["TARGET"]
     assert target_stats.emitted == {"SCAN": 1, "DNS_NAME": 1}
@@ -338,19 +358,17 @@ async def handle_event(self, event):
     assert target_stats.consumed_total == 0
 
     dummy_stats = scan.stats.module_stats["dummy"]
-    assert dummy_stats.emitted == {"FINDING": 1}
-    assert dummy_stats.emitted_total == 1
-    assert dummy_stats.produced == {"FINDING": 1}
-    assert dummy_stats.produced_total == 1
-    assert dummy_stats.consumed == {"SCAN": 1, "DNS_NAME": 1}
-    assert dummy_stats.consumed_total == 2
+    assert dummy_stats.emitted == {"FINDING": 1, "URL": 1}
+    assert dummy_stats.emitted_total == 2
+    assert dummy_stats.produced == {"FINDING": 1, "URL": 1}
+    assert dummy_stats.produced_total == 2
+    assert dummy_stats.consumed == {"DNS_NAME": 2, "OPEN_TCP_PORT": 1, "SCAN": 1, "URL": 1}
+    assert dummy_stats.consumed_total == 5
 
     python_stats = scan.stats.module_stats["python"]
     assert python_stats.emitted == {}
     assert python_stats.emitted_total == 0
     assert python_stats.produced == {}
     assert python_stats.produced_total == 0
-    assert python_stats.consumed == {"SCAN": 1, "FINDING": 1, "DNS_NAME": 1}
-    assert python_stats.consumed_total == 3
-
-    assert scan.stats.events_emitted_by_type == {"SCAN": 1, "FINDING": 1, "DNS_NAME": 1}
+    assert python_stats.consumed == {"DNS_NAME": 2, "FINDING": 1, "ORG_STUB": 1, "SCAN": 1, "URL": 1}
+    assert python_stats.consumed_total == 6

bbot/test/test_step_1/test_scan.py

@@ -9,6 +9,7 @@ async def test_scan(
     neograph,
     monkeypatch,
     bbot_scanner,
+    mock_dns,
 ):
     scan0 = bbot_scanner(
         "1.1.1.1/31",
@@ -55,15 +56,15 @@ async def test_scan(
     assert not scan2.in_scope("1.0.0.1")
 
     dns_table = {
-        ("1.1.1.1", "PTR"): "one.one.one.one",
-        ("one.one.one.one", "A"): "1.1.1.1",
+        "1.1.1.1": {"PTR": ["one.one.one.one"]},
+        "one.one.one.one": {"A": ["1.1.1.1"]},
     }
 
     # make sure DNS resolution works
     dns_config = OmegaConf.create({"dns_resolution": True})
     dns_config = OmegaConf.merge(bbot_config, dns_config)
     scan4 = bbot_scanner("1.1.1.1", config=dns_config)
-    scan4.helpers.dns.mock_dns(dns_table)
+    mock_dns(scan4, dns_table)
     events = []
     async for event in scan4.async_start():
         events.append(event)
@@ -74,7 +75,7 @@ async def test_scan(
     no_dns_config = OmegaConf.create({"dns_resolution": False})
     no_dns_config = OmegaConf.merge(bbot_config, no_dns_config)
     scan5 = bbot_scanner("1.1.1.1", config=no_dns_config)
-    scan5.helpers.dns.mock_dns(dns_table)
+    mock_dns(scan5, dns_table)
     events = []
     async for event in scan5.async_start():
         events.append(event)