crt/crt_db error handling

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()

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

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"