From 22823dbf327bddc28fce8ab2b688c276839606ca Mon Sep 17 00:00:00 2001 From: TheTechromancer Date: Tue, 30 Jan 2024 17:43:39 -0500 Subject: [PATCH 1/6] dns mock cleanup --- bbot/core/helpers/dns.py | 25 ------- bbot/test/bbot_fixtures.py | 65 ++++++++++++++++++ bbot/test/conftest.py | 67 ------------------- bbot/test/test_step_1/test_dns.py | 10 ++- .../test_step_1/test_manager_deduplication.py | 16 ++--- .../test_manager_scope_accuracy.py | 8 +-- bbot/test/test_step_1/test_modules_basic.py | 46 +++++++++---- bbot/test/test_step_1/test_scan.py | 9 +-- bbot/test/test_step_2/module_tests/base.py | 14 ++-- .../module_tests/test_module_affiliates.py | 9 ++- .../test_module_asset_inventory.py | 6 +- .../module_tests/test_module_baddns.py | 14 ++-- .../module_tests/test_module_baddns_zone.py | 21 +++--- .../module_tests/test_module_internetdb.py | 8 +-- .../module_tests/test_module_ipneighbor.py | 18 +---- .../module_tests/test_module_postman.py | 2 +- .../module_tests/test_module_speculate.py | 6 +- 17 files changed, 160 insertions(+), 184 deletions(-) diff --git a/bbot/core/helpers/dns.py b/bbot/core/helpers/dns.py index 9ad22116f..825536fa8 100644 --- a/bbot/core/helpers/dns.py +++ b/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 diff --git a/bbot/test/bbot_fixtures.py b/bbot/test/bbot_fixtures.py index a08ef3ece..7170e5a70 100644 --- a/bbot/test/bbot_fixtures.py +++ b/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 diff --git a/bbot/test/conftest.py b/bbot/test/conftest.py index dd59eac30..b60a0633d 100644 --- a/bbot/test/conftest.py +++ b/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 diff --git a/bbot/test/test_step_1/test_dns.py b/bbot/test/test_step_1/test_dns.py index debaed2b9..37ddf86a8 100644 --- a/bbot/test/test_step_1/test_dns.py +++ b/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( diff --git a/bbot/test/test_step_1/test_manager_deduplication.py b/bbot/test/test_step_1/test_manager_deduplication.py index e046988ae..e8ec69edc 100644 --- a/bbot/test/test_step_1/test_manager_deduplication.py +++ b/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 diff --git a/bbot/test/test_step_1/test_manager_scope_accuracy.py b/bbot/test/test_step_1/test_manager_scope_accuracy.py index e8e5da391..1175d0608 100644 --- a/bbot/test/test_step_1/test_manager_scope_accuracy.py +++ b/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="") @@ -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()] diff --git a/bbot/test/test_step_1/test_modules_basic.py b/bbot/test/test_step_1/test_modules_basic.py index 7a25bb4ea..e226eaea4 100644 --- a/bbot/test/test_step_1/test_modules_basic.py +++ b/bbot/test/test_step_1/test_modules_basic.py @@ -303,7 +303,7 @@ 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): @@ -311,23 +311,43 @@ class dummy(BaseModule): 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 diff --git a/bbot/test/test_step_1/test_scan.py b/bbot/test/test_step_1/test_scan.py index fde8dc368..27ebe3404 100644 --- a/bbot/test/test_step_1/test_scan.py +++ b/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) diff --git a/bbot/test/test_step_2/module_tests/base.py b/bbot/test/test_step_2/module_tests/base.py index e511d7256..be8c282fb 100644 --- a/bbot/test/test_step_2/module_tests/base.py +++ b/bbot/test/test_step_2/module_tests/base.py @@ -8,7 +8,7 @@ from bbot.scanner import Scanner from bbot.modules import module_loader from bbot.core.helpers.misc import rand_string -from ...bbot_fixtures import test_config +from ...bbot_fixtures import test_config, MockResolver log = logging.getLogger("bbot.test.modules") @@ -93,17 +93,17 @@ def set_expect_requests(self, expect_args={}, respond_args={}): def set_expect_requests_handler(self, expect_args=None, request_handler=None): self.httpserver.expect_request(expect_args).respond_with_handler(request_handler) + def mock_dns(self, mock_data, scan=None): + if scan is None: + scan = self.scan + scan.helpers.dns.resolver = MockResolver(mock_data) + @property def module(self): return self.scan.modules[self.name] - def mock_record(self, *args, **kwargs): - return MockRecord(*args, **kwargs) - @pytest_asyncio.fixture - async def module_test( - self, httpx_mock, bbot_httpserver, bbot_httpserver_ssl, monkeypatch, request, configure_mock_resolver - ): + async def module_test(self, httpx_mock, bbot_httpserver, bbot_httpserver_ssl, monkeypatch, request): module_test = self.ModuleTest(self, httpx_mock, bbot_httpserver, bbot_httpserver_ssl, monkeypatch, request) module_test.log.info(f"Starting {self.name} module test") await self.setup_before_prep(module_test) diff --git a/bbot/test/test_step_2/module_tests/test_module_affiliates.py b/bbot/test/test_step_2/module_tests/test_module_affiliates.py index fc9bd1c58..00135e4e4 100644 --- a/bbot/test/test_step_2/module_tests/test_module_affiliates.py +++ b/bbot/test/test_step_2/module_tests/test_module_affiliates.py @@ -6,12 +6,11 @@ class TestAffiliates(ModuleTestBase): config_overrides = {"dns_resolution": True} async def setup_before_prep(self, module_test): - module_test.scan.helpers.dns.mock_dns( + module_test.mock_dns( { - ("8.8.8.8", "PTR"): "dns.google", - ("dns.google", "A"): "8.8.8.8", - ("dns.google", "NS"): "ns1.zdns.google", - ("ns1.zdns.google", "A"): "1.2.3.4", + "8.8.8.8": {"PTR": ["dns.google"]}, + "dns.google": {"A": ["8.8.8.8"], "NS": ["ns1.zdns.google"]}, + "ns1.zdns.google": {"A": ["1.2.3.4"]}, } ) diff --git a/bbot/test/test_step_2/module_tests/test_module_asset_inventory.py b/bbot/test/test_step_2/module_tests/test_module_asset_inventory.py index cc4399266..d58808505 100644 --- a/bbot/test/test_step_2/module_tests/test_module_asset_inventory.py +++ b/bbot/test/test_step_2/module_tests/test_module_asset_inventory.py @@ -8,10 +8,10 @@ class TestAsset_Inventory(ModuleTestBase): modules_overrides = ["asset_inventory", "nmap", "sslcert"] async def setup_before_prep(self, module_test): - module_test.scan.helpers.dns.mock_dns( + module_test.mock_dns( { - ("127.0.0.1", "PTR"): "www.bbottest.notreal", - ("www.bbottest.notreal", "A"): "127.0.0.1", + "127.0.0.1": {"PTR": ["www.bbottest.notreal"]}, + "www.bbottest.notreal": {"A": ["127.0.0.1"]}, } ) diff --git a/bbot/test/test_step_2/module_tests/test_module_baddns.py b/bbot/test/test_step_2/module_tests/test_module_baddns.py index 5f1ad4d7a..57cca7d5f 100644 --- a/bbot/test/test_step_2/module_tests/test_module_baddns.py +++ b/bbot/test/test_step_2/module_tests/test_module_baddns.py @@ -24,10 +24,9 @@ async def setup_after_prep(self, module_test): from bbot.modules import baddns as baddns_module from baddns.lib.whoismanager import WhoisManager - mock_data = {"bad.dns": {"CNAME": ["baddns.azurewebsites.net."]}, "_NXDOMAIN": ["baddns.azurewebsites.net"]} - configure_mock_resolver = module_test.request_fixture.getfixturevalue("configure_mock_resolver") - mock_resolver = configure_mock_resolver(mock_data) - module_test.monkeypatch.setattr(module_test.scan.helpers.dns, "resolver", mock_resolver) + module_test.mock_dns( + {"bad.dns": {"CNAME": ["baddns.azurewebsites.net."]}, "_NXDOMAIN": ["baddns.azurewebsites.net"]} + ) module_test.monkeypatch.setattr(baddns_module.baddns, "select_modules", self.select_modules) module_test.monkeypatch.setattr(WhoisManager, "dispatchWHOIS", self.dispatchWHOIS) @@ -53,10 +52,9 @@ def set_target(self, target): respond_args = {"response_data": "

Oops! We couldn’t find that page.

", "status": 200} module_test.set_expect_requests(expect_args=expect_args, respond_args=respond_args) - mock_data = {"bad.dns": {"CNAME": ["baddns.bigcartel.com."]}, "baddns.bigcartel.com": {"A": ["127.0.0.1"]}} - configure_mock_resolver = module_test.request_fixture.getfixturevalue("configure_mock_resolver") - mock_resolver = configure_mock_resolver(mock_data) - module_test.monkeypatch.setattr(module_test.scan.helpers.dns, "resolver", mock_resolver) + module_test.mock_dns( + {"bad.dns": {"CNAME": ["baddns.bigcartel.com."]}, "baddns.bigcartel.com": {"A": ["127.0.0.1"]}} + ) module_test.monkeypatch.setattr(baddns_module.baddns, "select_modules", self.select_modules) module_test.monkeypatch.setattr(BadDNS_base, "set_target", set_target) module_test.monkeypatch.setattr(WhoisManager, "dispatchWHOIS", self.dispatchWHOIS) diff --git a/bbot/test/test_step_2/module_tests/test_module_baddns_zone.py b/bbot/test/test_step_2/module_tests/test_module_baddns_zone.py index 1a0ef9cf5..349db8ed8 100644 --- a/bbot/test/test_step_2/module_tests/test_module_baddns_zone.py +++ b/bbot/test/test_step_2/module_tests/test_module_baddns_zone.py @@ -31,12 +31,8 @@ def from_xfr(*args, **kwargs): zone = dns.zone.from_text(zone_text, origin="bad.dns.") return zone - mock_data = {"bad.dns": {"NS": ["ns1.bad.dns."]}, "ns1.bad.dns": {"A": ["127.0.0.1"]}} - configure_mock_resolver = module_test.request_fixture.getfixturevalue("configure_mock_resolver") - - mock_resolver = configure_mock_resolver(mock_data) + module_test.mock_dns({"bad.dns": {"NS": ["ns1.bad.dns."]}, "ns1.bad.dns": {"A": ["127.0.0.1"]}}) module_test.monkeypatch.setattr("dns.zone.from_xfr", from_xfr) - module_test.monkeypatch.setattr(module_test.scan.helpers.dns, "resolver", mock_resolver) module_test.monkeypatch.setattr(WhoisManager, "dispatchWHOIS", self.dispatchWHOIS) def check(self, module_test, events): @@ -50,14 +46,13 @@ class TestBaddns_zone_nsec(BaseTestBaddns_zone): async def setup_after_prep(self, module_test): from baddns.lib.whoismanager import WhoisManager - mock_data = { - "bad.dns": {"NSEC": ["asdf.bad.dns"]}, - "asdf.bad.dns": {"NSEC": ["zzzz.bad.dns"]}, - "zzzz.bad.dns": {"NSEC": ["xyz.bad.dns"]}, - } - configure_mock_resolver = module_test.request_fixture.getfixturevalue("configure_mock_resolver") - mock_resolver = configure_mock_resolver(mock_data) - module_test.monkeypatch.setattr(module_test.scan.helpers.dns, "resolver", mock_resolver) + module_test.mock_dns( + { + "bad.dns": {"NSEC": ["asdf.bad.dns"]}, + "asdf.bad.dns": {"NSEC": ["zzzz.bad.dns"]}, + "zzzz.bad.dns": {"NSEC": ["xyz.bad.dns"]}, + } + ) module_test.monkeypatch.setattr(WhoisManager, "dispatchWHOIS", self.dispatchWHOIS) def check(self, module_test, events): diff --git a/bbot/test/test_step_2/module_tests/test_module_internetdb.py b/bbot/test/test_step_2/module_tests/test_module_internetdb.py index 5b232f679..d24cdebc0 100644 --- a/bbot/test/test_step_2/module_tests/test_module_internetdb.py +++ b/bbot/test/test_step_2/module_tests/test_module_internetdb.py @@ -5,11 +5,11 @@ class TestInternetDB(ModuleTestBase): config_overrides = {"dns_resolution": True} async def setup_before_prep(self, module_test): - module_test.scan.helpers.mock_dns( + module_test.mock_dns( { - ("blacklanternsecurity.com", "A"): "1.2.3.4", - ("autodiscover.blacklanternsecurity.com", "A"): "2.3.4.5", - ("mail.blacklanternsecurity.com", "A"): "3.4.5.6", + "blacklanternsecurity.com": {"A": ["1.2.3.4"]}, + "autodiscover.blacklanternsecurity.com": {"A": ["2.3.4.5"]}, + "mail.blacklanternsecurity.com": {"A": ["3.4.5.6"]}, } ) diff --git a/bbot/test/test_step_2/module_tests/test_module_ipneighbor.py b/bbot/test/test_step_2/module_tests/test_module_ipneighbor.py index fbccc69a7..b8ba8331a 100644 --- a/bbot/test/test_step_2/module_tests/test_module_ipneighbor.py +++ b/bbot/test/test_step_2/module_tests/test_module_ipneighbor.py @@ -6,21 +6,9 @@ class TestIPNeighbor(ModuleTestBase): config_overrides = {"scope_report_distance": 1, "dns_resolution": True, "scope_dns_search_distance": 2} async def setup_after_prep(self, module_test): - old_resolve_ip = module_test.scan.helpers.dns._resolve_ip - old_resolve_hostname = module_test.scan.helpers.dns._resolve_hostname - - async def _resolve_ip(query, **kwargs): - if query == "127.0.0.3": - return [module_test.mock_record("asdf.www.bls.notreal", "PTR")], [] - return await old_resolve_ip(query, **kwargs) - - async def _resolve_hostname(query, **kwargs): - if query == "asdf.www.bls.notreal" and kwargs.get("rdtype", "") == "A": - return [module_test.mock_record("127.0.0.3", "A")], [] - return await old_resolve_hostname(query, **kwargs) - - module_test.monkeypatch.setattr(module_test.scan.helpers.dns, "_resolve_ip", _resolve_ip) - module_test.monkeypatch.setattr(module_test.scan.helpers.dns, "_resolve_hostname", _resolve_hostname) + module_test.mock_dns( + {"3.0.0.127.in-addr.arpa": {"PTR": ["asdf.www.bls.notreal"]}, "asdf.www.bls.notreal": {"A": ["127.0.0.3"]}} + ) def check(self, module_test, events): assert any(e.data == "127.0.0.3" for e in events) diff --git a/bbot/test/test_step_2/module_tests/test_module_postman.py b/bbot/test/test_step_2/module_tests/test_module_postman.py index e4c8d8ce6..21f464054 100644 --- a/bbot/test/test_step_2/module_tests/test_module_postman.py +++ b/bbot/test/test_step_2/module_tests/test_module_postman.py @@ -235,7 +235,7 @@ async def new_emit_event(event_data, event_type, **kwargs): await old_emit_event(event_data, event_type, **kwargs) module_test.monkeypatch.setattr(module_test.module, "emit_event", new_emit_event) - module_test.scan.helpers.dns.mock_dns({("asdf.blacklanternsecurity.com", "A"): "127.0.0.1"}) + module_test.mock_dns({"asdf.blacklanternsecurity.com": {"A": ["127.0.0.1"]}}) request_args = dict(uri="/_api/request/28129865-987c8ac8-bfa9-4bab-ade9-88ccf0597862") respond_args = dict(response_data="https://asdf.blacklanternsecurity.com") diff --git a/bbot/test/test_step_2/module_tests/test_module_speculate.py b/bbot/test/test_step_2/module_tests/test_module_speculate.py index db36db97b..0036451e4 100644 --- a/bbot/test/test_step_2/module_tests/test_module_speculate.py +++ b/bbot/test/test_step_2/module_tests/test_module_speculate.py @@ -28,10 +28,10 @@ class TestSpeculate_OpenPorts(ModuleTestBase): config_overrides = {"speculate": True} async def setup_before_prep(self, module_test): - module_test.scan.helpers.dns.mock_dns( + module_test.mock_dns( { - ("evilcorp.com", "A"): "127.0.254.1", - ("asdf.evilcorp.com", "A"): "127.0.254.2", + "evilcorp.com": {"A": "127.0.254.1"}, + "asdf.evilcorp.com": {"A": "127.0.254.2"}, } ) From 5a63b73801c27cd1f888c1b5f59668b44ea6248e Mon Sep 17 00:00:00 2001 From: TheTechromancer Date: Wed, 31 Jan 2024 18:02:43 -0500 Subject: [PATCH 2/6] fix dnscommonsrv tests --- .../test_manager_scope_accuracy.py | 26 +++++++++---------- .../module_tests/test_module_dnscommonsrv.py | 21 +++++++-------- 2 files changed, 23 insertions(+), 24 deletions(-) diff --git a/bbot/test/test_step_1/test_manager_scope_accuracy.py b/bbot/test/test_step_1/test_manager_scope_accuracy.py index 1175d0608..453625a06 100644 --- a/bbot/test/test_step_1/test_manager_scope_accuracy.py +++ b/bbot/test/test_step_1/test_manager_scope_accuracy.py @@ -31,7 +31,7 @@ def bbot_other_httpservers(): @pytest.mark.asyncio -async def test_manager_scope_accuracy(bbot_config, bbot_scanner, bbot_httpserver, bbot_other_httpservers, bbot_httpserver_ssl): +async def test_manager_scope_accuracy(bbot_config, bbot_scanner, bbot_httpserver, bbot_other_httpservers, bbot_httpserver_ssl, mock_dns): """ This test ensures that BBOT correctly handles different scope distance settings. It performs these tests for normal modules, output modules, and their graph variants, @@ -102,7 +102,7 @@ async def do_scan(*args, _config={}, _dns_mock={}, scan_callback=None, **kwargs) scan.modules["dummy_graph_output_module"] = dummy_graph_output_module scan.modules["dummy_graph_batch_output_module"] = dummy_graph_batch_output_module 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 ( @@ -114,12 +114,12 @@ async def do_scan(*args, _config={}, _dns_mock={}, scan_callback=None, **kwargs) ) dns_mock_chain = { - ("test.notreal", "A"): "127.0.0.66", - ("127.0.0.66", "PTR"): "test.notrealzies", - ("test.notrealzies", "CNAME"): "www.test.notreal", - ("www.test.notreal", "A"): "127.0.0.77", - ("127.0.0.77", "PTR"): "test2.notrealzies", - ("test2.notrealzies", "A"): "127.0.0.88", + "test.notreal": {"A": ["127.0.0.66"]}, + "66.0.0.127.in-addr.arpa": {"PTR": ["test.notrealzies"]}, + "test.notrealzies": {"CNAME": ["www.test.notreal"]}, + "www.test.notreal": {"A": ["127.0.0.77"]}, + "77.0.0.127.in-addr.arpa": {"PTR": ["test2.notrealzies"]}, + "test2.notrealzies": {"A": ["127.0.0.88"]}, } # dns search distance = 1, report distance = 0 @@ -240,9 +240,9 @@ async def do_scan(*args, _config={}, _dns_mock={}, scan_callback=None, **kwargs) assert 0 == len([e for e in _graph_output_events if e.type == "IP_ADDRESS" and e.data == "127.0.0.88"]) dns_mock_chain = { - ("test.notreal", "A"): "127.0.0.66", - ("127.0.0.66", "PTR"): "test.notrealzies", - ("test.notrealzies", "A"): "127.0.0.77", + "test.notreal": {"A": ["127.0.0.66"]}, + "66.0.0.127.in-addr.arpa": {"PTR": ["test.notrealzies"]}, + "test.notrealzies": {"A": ["127.0.0.77"]}, } class DummyVulnModule(BaseModule): @@ -667,7 +667,7 @@ def custom_setup(scan): "127.0.0.0/31", modules=["speculate", "sslcert"], _config={"dns_resolution": False, "scope_report_distance": 0, "internal_modules": {"speculate": {"ports": "9999"}}}, - _dns_mock={("www.bbottest.notreal", "A"): "127.0.1.0", ("test.notreal", "A"): "127.0.0.1"}, + _dns_mock={"www.bbottest.notreal": {"A": ["127.0.1.0"]}, "test.notreal": {"A": ["127.0.0.1"]}}, ) assert len(events) == 6 @@ -727,7 +727,7 @@ def custom_setup(scan): modules=["speculate", "sslcert"], whitelist=["127.0.1.0"], _config={"dns_resolution": False, "scope_report_distance": 0, "internal_modules": {"speculate": {"ports": "9999"}}}, - _dns_mock={("www.bbottest.notreal", "A"): "127.0.0.1", ("test.notreal", "A"): "127.0.1.0"}, + _dns_mock={"www.bbottest.notreal": {"A": ["127.0.0.1"]}, "test.notreal": {"A": ["127.0.1.0"]}}, ) assert len(events) == 3 diff --git a/bbot/test/test_step_2/module_tests/test_module_dnscommonsrv.py b/bbot/test/test_step_2/module_tests/test_module_dnscommonsrv.py index bdb6f1207..aaf26664c 100644 --- a/bbot/test/test_step_2/module_tests/test_module_dnscommonsrv.py +++ b/bbot/test/test_step_2/module_tests/test_module_dnscommonsrv.py @@ -3,22 +3,21 @@ class TestDNSCommonSRV(ModuleTestBase): targets = ["blacklanternsecurity.notreal"] + config_overrides = {"dns_resolution": True} async def setup_after_prep(self, module_test): - old_resolve_fn = module_test.scan.helpers.dns.resolve - - async def resolve(query, **kwargs): - if ( - query == "_ldap._tcp.gc._msdcs.blacklanternsecurity.notreal" - and kwargs.get("type", "").upper() == "SRV" - ): - return {"asdf.blacklanternsecurity.notreal"} - return await old_resolve_fn(query, **kwargs) - - module_test.monkeypatch.setattr(module_test.scan.helpers.dns, "resolve", resolve) + module_test.mock_dns( + { + "_ldap._tcp.gc._msdcs.blacklanternsecurity.notreal": { + "SRV": ["0 100 3268 asdf.blacklanternsecurity.notreal"] + }, + "asdf.blacklanternsecurity.notreal": {"A": "1.2.3.4"}, + } + ) def check(self, module_test, events): assert any( e.data == "_ldap._tcp.gc._msdcs.blacklanternsecurity.notreal" for e in events ), "Failed to detect subdomain" + assert any(e.data == "asdf.blacklanternsecurity.notreal" for e in events), "Failed to detect subdomain" assert not any(e.data == "_ldap._tcp.dc._msdcs.blacklanternsecurity.notreal" for e in events), "False positive" From ae7e41fa012f8cd66ca39216c3d0071bfa96bea6 Mon Sep 17 00:00:00 2001 From: TheTechromancer Date: Thu, 1 Feb 2024 13:01:28 -0500 Subject: [PATCH 3/6] fix tests --- bbot/test/test_step_1/test_modules_basic.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/bbot/test/test_step_1/test_modules_basic.py b/bbot/test/test_step_1/test_modules_basic.py index e226eaea4..8784923ed 100644 --- a/bbot/test/test_step_1/test_modules_basic.py +++ b/bbot/test/test_step_1/test_modules_basic.py @@ -337,14 +337,21 @@ async def handle_event(self, event): events = [e async for e in scan.async_start()] assert len(events) == 6 + assert 1 == len([e for e in events if e.type == "SCAN"]) + assert 2 == len([e for e in events if e.type == "DNS_NAME"]) + assert 1 == len([e for e in events if e.type == "DNS_NAME" and e.data == "evilcorp.com"]) + # the reason we don't have a DNS_NAME for www.evilcorp.com is because FINDING.quick_emit = True + assert 0 == len([e for e in events if e.type == "DNS_NAME" and e.data == "www.evilcorp.com"]) + assert 1 == len([e for e in events if e.type == "DNS_NAME" and e.data == "asdf.evilcorp.com"]) + assert 1 == len([e for e in events if e.type == "ORG_STUB" and e.data == "evilcorp"]) + assert 1 == len([e for e in events if e.type == "FINDING"]) assert scan.stats.events_emitted_by_type == { "SCAN": 1, - "DNS_NAME": 3, + "DNS_NAME": 2, "URL": 1, "FINDING": 1, "ORG_STUB": 1, - "OPEN_TCP_PORT": 1, } assert set(scan.stats.module_stats) == {"host", "speculate", "python", "dummy", "TARGET"} From 652b233df1e472b780972491cc344f34de7cee34 Mon Sep 17 00:00:00 2001 From: TheTechromancer Date: Thu, 1 Feb 2024 13:56:56 -0500 Subject: [PATCH 4/6] fix PTR tests --- bbot/test/test_step_1/test_scan.py | 2 +- bbot/test/test_step_2/module_tests/test_module_affiliates.py | 2 +- .../test_step_2/module_tests/test_module_asset_inventory.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/bbot/test/test_step_1/test_scan.py b/bbot/test/test_step_1/test_scan.py index 27ebe3404..9aa89ffe2 100644 --- a/bbot/test/test_step_1/test_scan.py +++ b/bbot/test/test_step_1/test_scan.py @@ -56,7 +56,7 @@ async def test_scan( assert not scan2.in_scope("1.0.0.1") dns_table = { - "1.1.1.1": {"PTR": ["one.one.one.one"]}, + "1.1.1.1.in-addr.arpa": {"PTR": ["one.one.one.one"]}, "one.one.one.one": {"A": ["1.1.1.1"]}, } diff --git a/bbot/test/test_step_2/module_tests/test_module_affiliates.py b/bbot/test/test_step_2/module_tests/test_module_affiliates.py index 00135e4e4..4afd4cd29 100644 --- a/bbot/test/test_step_2/module_tests/test_module_affiliates.py +++ b/bbot/test/test_step_2/module_tests/test_module_affiliates.py @@ -8,7 +8,7 @@ class TestAffiliates(ModuleTestBase): async def setup_before_prep(self, module_test): module_test.mock_dns( { - "8.8.8.8": {"PTR": ["dns.google"]}, + "8.8.8.8.in-addr.arpa": {"PTR": ["dns.google"]}, "dns.google": {"A": ["8.8.8.8"], "NS": ["ns1.zdns.google"]}, "ns1.zdns.google": {"A": ["1.2.3.4"]}, } diff --git a/bbot/test/test_step_2/module_tests/test_module_asset_inventory.py b/bbot/test/test_step_2/module_tests/test_module_asset_inventory.py index d58808505..b63de31ec 100644 --- a/bbot/test/test_step_2/module_tests/test_module_asset_inventory.py +++ b/bbot/test/test_step_2/module_tests/test_module_asset_inventory.py @@ -10,7 +10,7 @@ class TestAsset_Inventory(ModuleTestBase): async def setup_before_prep(self, module_test): module_test.mock_dns( { - "127.0.0.1": {"PTR": ["www.bbottest.notreal"]}, + "1.0.0.127.in-addr.arpa": {"PTR": ["www.bbottest.notreal"]}, "www.bbottest.notreal": {"A": ["127.0.0.1"]}, } ) From e5d88dcf9d6a02b1ef3a8725da9e3c83fb52c75d Mon Sep 17 00:00:00 2001 From: TheTechromancer Date: Thu, 1 Feb 2024 14:18:55 -0500 Subject: [PATCH 5/6] fix aggregate tests --- bbot/test/test_step_2/module_tests/test_module_aggregate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bbot/test/test_step_2/module_tests/test_module_aggregate.py b/bbot/test/test_step_2/module_tests/test_module_aggregate.py index 59f5370b3..7a41fe022 100644 --- a/bbot/test/test_step_2/module_tests/test_module_aggregate.py +++ b/bbot/test/test_step_2/module_tests/test_module_aggregate.py @@ -5,7 +5,7 @@ class TestAggregate(ModuleTestBase): config_overrides = {"dns_resolution": True, "scope_report_distance": 1} async def setup_before_prep(self, module_test): - module_test.scan.helpers.dns.mock_dns({("blacklanternsecurity.com", "A"): "1.2.3.4"}) + module_test.mock_dns({"blacklanternsecurity.com": {"A": ["1.2.3.4"]}}) def check(self, module_test, events): filename = next(module_test.scan.home.glob("scan-stats-table*.txt")) From 5f8d266d8a46847563adc87be82f2a71e970574d Mon Sep 17 00:00:00 2001 From: TheTechromancer Date: Thu, 1 Feb 2024 17:59:50 -0500 Subject: [PATCH 6/6] fix speculate tests --- bbot/test/test_step_2/module_tests/test_module_speculate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bbot/test/test_step_2/module_tests/test_module_speculate.py b/bbot/test/test_step_2/module_tests/test_module_speculate.py index 0036451e4..2dcafaddc 100644 --- a/bbot/test/test_step_2/module_tests/test_module_speculate.py +++ b/bbot/test/test_step_2/module_tests/test_module_speculate.py @@ -30,8 +30,8 @@ class TestSpeculate_OpenPorts(ModuleTestBase): async def setup_before_prep(self, module_test): module_test.mock_dns( { - "evilcorp.com": {"A": "127.0.254.1"}, - "asdf.evilcorp.com": {"A": "127.0.254.2"}, + "evilcorp.com": {"A": ["127.0.254.1"]}, + "asdf.evilcorp.com": {"A": ["127.0.254.2"]}, } )