diff --git a/bbot/modules/crt.py b/bbot/modules/crt.py index 60ea6b6a0c..e977ddccbd 100644 --- a/bbot/modules/crt.py +++ b/bbot/modules/crt.py @@ -23,6 +23,10 @@ async def request_url(self, query): url = self.helpers.add_get_params(self.base_url, params).geturl() return await self.api_request(url, timeout=self.http_timeout + 30) + def _api_response_is_success(self, r): + # crt.sh returns 404/503 transiently; the default treats 404 as "no data" which is wrong here + return getattr(r, "is_success", False) + async def parse_results(self, r, query): results = set() j = r.json() diff --git a/bbot/modules/crt_db.py b/bbot/modules/crt_db.py index 4008394f1d..85eb97b0f7 100644 --- a/bbot/modules/crt_db.py +++ b/bbot/modules/crt_db.py @@ -26,15 +26,20 @@ async def setup(self): self.db_conn = None return await super().setup() + async def _connect(self): + return await asyncpg.connect( + host=self.db_host, + port=self.db_port, + user=self.db_user, + database=self.db_name, + statement_cache_size=0, # Disable automatic statement preparation + ) + async def request_url(self, query): - if not self.db_conn: - self.db_conn = await asyncpg.connect( - host=self.db_host, - port=self.db_port, - user=self.db_user, - database=self.db_name, - statement_cache_size=0, # Disable automatic statement preparation - ) + # asyncpg connections can drop mid-scan when the upstream Postgres restarts, + # exhausts shared memory, or idles us out. Reconnect on the next call. + if self.db_conn is None or self.db_conn.is_closed(): + self.db_conn = await self._connect() sql = """ WITH ci AS ( @@ -51,7 +56,16 @@ async def request_url(self, query): SELECT DISTINCT unnest(NAME_VALUES) as name_value FROM ci; """ start = time.time() - results = await self.db_conn.fetch(sql, query) + try: + results = await self.db_conn.fetch(sql, query) + except (asyncpg.InterfaceError, asyncpg.PostgresConnectionError, ConnectionError): + # connection died between requests; drop it so the next call reconnects + self.db_conn = None + raise + except asyncpg.OutOfMemoryError as e: + # upstream is overloaded; bail out for the rest of the scan rather than hammer it + self.set_error_state(f"crt.sh Postgres reported out-of-memory: {e}") + return [] end = time.time() self.verbose(f"SQL query executed in: {end - start} seconds with {len(results):,} results") return results diff --git a/bbot/test/test_step_2/module_tests/test_module_crt_db.py b/bbot/test/test_step_2/module_tests/test_module_crt_db.py index 995d2cd52b..d79b5f6b30 100644 --- a/bbot/test/test_step_2/module_tests/test_module_crt_db.py +++ b/bbot/test/test_step_2/module_tests/test_module_crt_db.py @@ -1,23 +1,94 @@ from .base import ModuleTestBase -class TestCRT_DB(ModuleTestBase): - async def setup_after_prep(self, module_test): - class AsyncMock: - async def fetch(self, *args, **kwargs): - return [ - {"name_value": "asdf.blacklanternsecurity.com"}, - {"name_value": "zzzz.blacklanternsecurity.com"}, - ] +class FakeAsyncpgConn: + def __init__(self): + self._closed = False + self.fetch_calls = 0 + + def is_closed(self): + return self._closed + + async def fetch(self, *args, **kwargs): + self.fetch_calls += 1 + return [ + {"name_value": "asdf.blacklanternsecurity.com"}, + {"name_value": "zzzz.blacklanternsecurity.com"}, + ] - async def close(self): - pass + async def close(self): + self._closed = True + +class TestCRT_DB(ModuleTestBase): + async def setup_after_prep(self, module_test): async def mock_connect(*args, **kwargs): - return AsyncMock() + return FakeAsyncpgConn() module_test.monkeypatch.setattr("asyncpg.connect", mock_connect) def check(self, module_test, events): assert any(e.data == "asdf.blacklanternsecurity.com" for e in events), "Failed to detect subdomain" assert any(e.data == "zzzz.blacklanternsecurity.com" for e in events), "Failed to detect subdomain" + + +class TestCRT_DB_Reconnect(ModuleTestBase): + """ + Asyncpg connections drop mid-scan when the upstream Postgres restarts or idles us out. + The module must reopen on the next call instead of staying broken for the rest of the scan. + """ + + targets = ["blacklanternsecurity.com", "evilcorp.com"] + modules_overrides = ["crt_db"] + + async def setup_after_prep(self, module_test): + async def fetch_succeed(self, *args, **kwargs): + # args[1] is the query (domain); return a subdomain matching it + query = args[1] if len(args) > 1 else "blacklanternsecurity.com" + return [{"name_value": f"reconnect.{query}"}] + + async def fetch_then_fail(self, *args, **kwargs): + self._closed = True + raise ConnectionError("connection is closed") + + self.connect_count = 0 + + async def mock_connect(*args, **kwargs): + self.connect_count += 1 + conn = FakeAsyncpgConn() + if self.connect_count == 1: + conn.fetch = fetch_then_fail.__get__(conn) + else: + conn.fetch = fetch_succeed.__get__(conn) + return conn + + module_test.monkeypatch.setattr("asyncpg.connect", mock_connect) + + def check(self, module_test, events): + assert self.connect_count >= 2, "Module did not reconnect after closed connection" + assert any(isinstance(e.data, str) and e.data.startswith("reconnect.") for e in events), ( + "Failed to detect subdomain after reconnect" + ) + + +class TestCRT_DB_OOM(ModuleTestBase): + """When crt.sh's Postgres reports out-of-memory, the module must back off (error state) rather than hammer it.""" + + module_name = "crt_db" + modules_overrides = ["crt_db"] + + async def setup_after_prep(self, module_test): + import asyncpg + + async def fetch_oom(self, *args, **kwargs): + raise asyncpg.OutOfMemoryError("out of shared memory") + + async def mock_connect(*args, **kwargs): + conn = FakeAsyncpgConn() + conn.fetch = fetch_oom.__get__(conn) + return conn + + module_test.monkeypatch.setattr("asyncpg.connect", mock_connect) + + def check(self, module_test, events): + assert module_test.module.errored is True, "Module did not enter error state on Postgres out-of-memory"