Merge pull request #2413 from blacklanternsecurity/fix-shodan-ratelimit

Fix shodan internetdb ratelimit

Author: TheTechromancer

bbot/defaults.yml

@@ -92,7 +92,11 @@ web:
   # These are attached to all in-scope HTTP requests
   # Note that some modules (e.g. github) may end up sending these to out-of-scope resources
   http_headers: {}
-  # HTTP retries (for Python requests; API calls, etc.)
+  # How many times to retry API requests
+  # Note that this is a separate mechanism on top of HTTP retries
+  # which will retry API requests that don't return a successful status code
+  api_retries: 2
+  # HTTP retries - try again if the raw connection fails
   http_retries: 1
   # HTTP retries (for httpx)
   httpx_retries: 1

bbot/modules/base.py

@@ -104,8 +104,6 @@ class BaseModule:
     _batch_size = 1
     batch_wait = 10
 
-    # API retries, etc.
-    _api_retries = 2
     # disable the module after this many failed attempts in a row
     _api_failure_abort_threshold = 3
 
@@ -157,6 +155,8 @@ def __init__(self, scan):
         # track number of failures (for .api_request())
         self._api_request_failures = 0
 
+        self._default_api_retries = self.scan.config.get("web", {}).get("api_retries", 2)
+
         self._tasks = []
         self._event_received = None
 
@@ -340,7 +340,7 @@ def cycle_api_key(self):
 
     @property
     def api_retries(self):
-        return max(self._api_retries + 1, len(self._api_keys))
+        return max(self._default_api_retries + 1, len(self._api_keys))
 
     @property
     def api_failure_abort_threshold(self):
@@ -1212,7 +1212,8 @@ def _prepare_api_iter_req(self, url, page, page_size, offset, **requests_kwargs)
         return url, requests_kwargs
 
     def _api_response_is_success(self, r):
-        return r.is_success
+        # 404s typically indicate no data rather than an actual error with the API, so we don't want to retry them
+        return getattr(r, "is_success", False) or getattr(r, "status_code", 0) == 404
 
     async def api_page_iter(self, url, page_size=100, _json=True, next_key=None, iter_key=None, **requests_kwargs):
         """

bbot/modules/shodan_idb.py

@@ -1,4 +1,5 @@
 from bbot.modules.base import BaseModule
+import time
 
 
 class shodan_idb(BaseModule):
@@ -46,23 +47,48 @@ class shodan_idb(BaseModule):
         "created_date": "2023-12-22",
         "author": "@TheTechromancer",
     }
+    options = {"retries": None}
+    options_desc = {
+        "retries": "How many times to retry API requests (e.g. after a 429 error). Overrides the global web.api_retries setting."
+    }
 
-    # we get lots of 404s, that's normal
+    # we typically don't want to abort this module
     _api_failure_abort_threshold = 9999999999
 
-    # there aren't any rate limits to speak of, so our outgoing queue can be pretty big
-    _qsize = 500
+    # since there are rate limits, we set a lower qsize
+    # this way when our queue is full, we can give the API a break
+    _qsize = 100
 
     base_url = "https://internetdb.shodan.io"
 
+    async def setup(self):
+        await super().setup()
+        self.last_request_time = 0
+        return True
+
     def _incoming_dedup_hash(self, event):
         return hash(self.get_ip(event))
 
+    @property
+    def api_retries(self):
+        # allow the module to override global retry setting
+        return self.config.get("retries", None) or super().api_retries
+
     async def handle_event(self, event):
         ip = self.get_ip(event)
         if ip is None:
             return
         url = f"{self.base_url}/{ip}"
+
+        # Rate limiting: ensure at least 1 second between requests
+        current_time = time.time()
+        time_since_last = current_time - self.last_request_time
+        if time_since_last < 1:
+            await self.helpers.sleep(1 - time_since_last)
+
+        # Update the last request time
+        self.last_request_time = time.time()
+
         r = await self.api_request(url)
         if r is None:
             self.debug(f"No response for {event.data}")

bbot/modules/templates/webhook.py

@@ -16,10 +16,9 @@ class WebhookOutputModule(BaseOutputModule):
     # abort module after 10 failed requests (not including retries)
     _api_failure_abort_threshold = 10
     # retry each request up to 10 times, respecting the Retry-After header
-    _api_retries = 10
+    _default_api_retries = 10
 
     async def setup(self):
-        self._api_retries = self.config.get("retries", 10)
         self.webhook_url = self.config.get("webhook_url", "")
         self.min_severity = self.config.get("min_severity", "LOW").strip().upper()
         assert self.min_severity in self.vuln_severities, (
@@ -31,6 +30,10 @@ async def setup(self):
             return False
         return await super().setup()
 
+    @property
+    def api_retries(self):
+        return self.config.get("retries", self._default_api_retries)
+
     async def handle_event(self, event):
         message = self.format_message(event)
         data = {self.content_key: message}