Merge pull request #2206 from domwhewell-sage/teams_output_module

Add "retries" config to teams output module

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/discord.py

@@ -8,9 +8,10 @@ class Discord(WebhookOutputModule):
         "created_date": "2023-08-14",
         "author": "@TheTechromancer",
     }
-    options = {"webhook_url": "", "event_types": ["VULNERABILITY", "FINDING"], "min_severity": "LOW"}
+    options = {"webhook_url": "", "event_types": ["VULNERABILITY", "FINDING"], "min_severity": "LOW", "retries": 10}
     options_desc = {
         "webhook_url": "Discord webhook URL",
         "event_types": "Types of events to send",
         "min_severity": "Only allow VULNERABILITY events of this severity or higher",
+        "retries": "Number of times to retry sending the message before skipping the event",
     }

bbot/modules/output/slack.py

@@ -10,11 +10,12 @@ class Slack(WebhookOutputModule):
         "created_date": "2023-08-14",
         "author": "@TheTechromancer",
     }
-    options = {"webhook_url": "", "event_types": ["VULNERABILITY", "FINDING"], "min_severity": "LOW"}
+    options = {"webhook_url": "", "event_types": ["VULNERABILITY", "FINDING"], "min_severity": "LOW", "retries": 10}
     options_desc = {
         "webhook_url": "Discord webhook URL",
         "event_types": "Types of events to send",
         "min_severity": "Only allow VULNERABILITY events of this severity or higher",
+        "retries": "Number of times to retry sending the message before skipping the event",
     }
     content_key = "text"
 

bbot/modules/output/teams.py

@@ -8,36 +8,21 @@ class Teams(WebhookOutputModule):
         "created_date": "2023-08-14",
         "author": "@TheTechromancer",
     }
-    options = {"webhook_url": "", "event_types": ["VULNERABILITY", "FINDING"], "min_severity": "LOW"}
+    options = {"webhook_url": "", "event_types": ["VULNERABILITY", "FINDING"], "min_severity": "LOW", "retries": 10}
     options_desc = {
         "webhook_url": "Teams webhook URL",
         "event_types": "Types of events to send",
         "min_severity": "Only allow VULNERABILITY events of this severity or higher",
+        "retries": "Number of times to retry sending the message before skipping the event",
     }
-    _module_threads = 5
 
     async def handle_event(self, event):
-        while 1:
-            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_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)
+        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,7 +13,13 @@ 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._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 (
@@ -23,31 +29,16 @@ async def setup(self):
         if not self.webhook_url:
             self.warning("Must set Webhook URL")
             return False
-        return True
+        return await super().setup()
 
     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:

bbot/test/test_step_2/module_tests/test_module_teams.py

@@ -7,14 +7,16 @@ class TestTeams(DiscordBase):
     modules_overrides = ["teams", "excavate", "badsecrets", "httpx"]
 
     webhook_url = "https://evilcorp.webhook.office.com/webhookb2/deadbeef@deadbeef/IncomingWebhook/deadbeef/deadbeef"
-    config_overrides = {"modules": {"teams": {"webhook_url": webhook_url}}}
+    config_overrides = {"modules": {"teams": {"webhook_url": webhook_url, "retries": 5}}}
 
     async def setup_after_prep(self, module_test):
         self.custom_setup(module_test)
 
         def custom_response(request: httpx.Request):
             module_test.request_count += 1
             if module_test.request_count == 2:
+                return httpx.Response(status_code=429, headers={"Retry-After": "0.01"})
+            elif module_test.request_count == 3:
                 return httpx.Response(
                     status_code=400,
                     json={
@@ -28,3 +30,10 @@ def custom_response(request: httpx.Request):
                 return httpx.Response(status_code=200)
 
         module_test.httpx_mock.add_callback(custom_response, url=self.webhook_url)
+
+    def check(self, module_test, events):
+        vulns = [e for e in events if e.type == "VULNERABILITY"]
+        findings = [e for e in events if e.type == "FINDING"]
+        assert len(findings) == 1
+        assert len(vulns) == 2
+        assert module_test.request_count == 5