Use builtin retry mechanism

bbot/modules/base.py

@@ -1169,11 +1169,13 @@ async def api_request(self, *args, **kwargs):
                     )
                 else:
                     # sleep for a bit if we're being rate limited
-                    if status_code == 429:
+                    retry_after = self._get_retry_after(r)
+                    if retry_after or status_code == 429:
+                        sleep_interval = int(retry_after) if retry_after is not None else self._429_sleep_interval
                         self.verbose(
-                            f"Sleeping for {self._429_sleep_interval:,} seconds due to rate limit (HTTP status: 429)"
+                            f"Sleeping for {sleep_interval:,} seconds due to rate limit (HTTP status: {status_code})"
                         )
-                        await asyncio.sleep(self._429_sleep_interval)
+                        await asyncio.sleep(sleep_interval)
                     elif self._api_keys:
                         # if request failed, cycle API keys and try again
                         self.cycle_api_key()
@@ -1182,6 +1184,19 @@ async def api_request(self, *args, **kwargs):
 
         return r
 
+    def _get_retry_after(self, r):
+        # try to get retry_after from headers first
+        headers = getattr(r, "headers", {})
+        retry_after = headers.get("Retry-After", None)
+        if retry_after is None:
+            # then look in body json
+            with suppress(Exception):
+                body_json = r.json()
+                if isinstance(body_json, dict):
+                    retry_after = body_json.get("retry_after", None)
+        if retry_after is not None:
+            return float(retry_after)
+
     def _prepare_api_iter_req(self, url, page, page_size, offset, **requests_kwargs):
         """
         Default function for preparing an API request for iterating through paginated data.

bbot/modules/output/teams.py

@@ -15,34 +15,18 @@ class Teams(WebhookOutputModule):
         "min_severity": "Only allow VULNERABILITY events of this severity or higher",
         "retries": "Number of times to retry sending the message before skipping the event (Default: 10)",
     }
-    _module_threads = 5
 
     async def setup(self):
-        self.retries = self.config.get("retries", 10)
+        self._api_retries = self.config.get("retries", 10)
         return await super().setup()
 
     async def handle_event(self, event):
-        for _ in range(self.retries):
-            data = self.format_message(event)
-
-            response = await self.helpers.request(
-                url=self.webhook_url,
-                method="POST",
-                json=data,
-            )
-            status_code = getattr(response, "status_code", 0)
-            if self.evaluate_response(response):
-                break
-            else:
-                response_headers = response.headers
-                try:
-                    retry_after = float(response_headers.get("Retry-After", 1))
-                except Exception:
-                    retry_after = 1
-                self.verbose(
-                    f"Error sending {event}: status code {status_code}, response headers: {response_headers}, retrying in {retry_after} seconds"
-                )
-                await self.helpers.sleep(retry_after)
+        data = self.format_message(event)
+        await self.api_request(
+            url=self.webhook_url,
+            method="POST",
+            json=data,
+        )
 
     def trim_message(self, message):
         if len(message) > self.message_size_limit:

bbot/modules/templates/webhook.py

@@ -13,6 +13,11 @@ class WebhookOutputModule(BaseOutputModule):
     content_key = "content"
     vuln_severities = ["UNKNOWN", "LOW", "MEDIUM", "HIGH", "CRITICAL"]
 
+    # 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
+
     async def setup(self):
         self.webhook_url = self.config.get("webhook_url", "")
         self.min_severity = self.config.get("min_severity", "LOW").strip().upper()
@@ -26,28 +31,13 @@ async def setup(self):
         return True
 
     async def handle_event(self, event):
-        while 1:
-            message = self.format_message(event)
-            data = {self.content_key: message}
-
-            response = await self.helpers.request(
-                url=self.webhook_url,
-                method="POST",
-                json=data,
-            )
-            status_code = getattr(response, "status_code", 0)
-            if self.evaluate_response(response):
-                break
-            else:
-                response_data = getattr(response, "text", "")
-                try:
-                    retry_after = response.json().get("retry_after", 1)
-                except Exception:
-                    retry_after = 1
-                self.verbose(
-                    f"Error sending {event}: status code {status_code}, response: {response_data}, retrying in {retry_after} seconds"
-                )
-                await self.helpers.sleep(retry_after)
+        message = self.format_message(event)
+        data = {self.content_key: message}
+        await self.api_request(
+            url=self.webhook_url,
+            method="POST",
+            json=data,
+        )
 
     def get_watched_events(self):
         if self._watched_events is None: