From a1b2edff2bc18837a5af71032f4b5011965dd9cf Mon Sep 17 00:00:00 2001 From: Dom Whewell Date: Thu, 23 Jan 2025 18:25:56 +0000 Subject: [PATCH 1/8] Change the teams output module to only retry sending requests a set number of times --- bbot/modules/output/teams.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/bbot/modules/output/teams.py b/bbot/modules/output/teams.py index 5703ad80f..670baee90 100644 --- a/bbot/modules/output/teams.py +++ b/bbot/modules/output/teams.py @@ -8,16 +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 (Default: 10)", } _module_threads = 5 + async def setup(self): + self.retries = self.config.get("retries", 10) + return await super().setup() + async def handle_event(self, event): - while 1: + for _ in range(self.retries): data = self.format_message(event) response = await self.helpers.request( From d32b69928318b91875a77bfb41909e0ab8c59758 Mon Sep 17 00:00:00 2001 From: Dom Whewell Date: Thu, 23 Jan 2025 18:31:14 +0000 Subject: [PATCH 2/8] Get Retry-After from the webhook response headers --- bbot/modules/output/teams.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/bbot/modules/output/teams.py b/bbot/modules/output/teams.py index 670baee90..9b8f7a1ae 100644 --- a/bbot/modules/output/teams.py +++ b/bbot/modules/output/teams.py @@ -34,13 +34,13 @@ async def handle_event(self, event): if self.evaluate_response(response): break else: - response_data = getattr(response, "text", "") + response_headers = response.headers try: - retry_after = response.json().get("retry_after", 1) + retry_after = int(response_headers.get("Retry-After", 123)) except Exception: - retry_after = 1 + retry_after = 123 self.verbose( - f"Error sending {event}: status code {status_code}, response: {response_data}, retrying in {retry_after} seconds" + f"Error sending {event}: status code {status_code}, response headers: {response_headers}, retrying in {retry_after} seconds" ) await self.helpers.sleep(retry_after) From 1d28fcc2b80ab4c64d97045a7a5fa9fddc434082 Mon Sep 17 00:00:00 2001 From: Dom Whewell Date: Thu, 23 Jan 2025 18:57:04 +0000 Subject: [PATCH 3/8] Reduce the teams retry and change to float --- bbot/modules/output/teams.py | 4 ++-- .../test_step_2/module_tests/test_module_teams.py | 11 ++++++++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/bbot/modules/output/teams.py b/bbot/modules/output/teams.py index 9b8f7a1ae..bfebf73b5 100644 --- a/bbot/modules/output/teams.py +++ b/bbot/modules/output/teams.py @@ -36,9 +36,9 @@ async def handle_event(self, event): else: response_headers = response.headers try: - retry_after = int(response_headers.get("Retry-After", 123)) + retry_after = float(response_headers.get("Retry-After", 1)) except Exception: - retry_after = 123 + retry_after = 1 self.verbose( f"Error sending {event}: status code {status_code}, response headers: {response_headers}, retrying in {retry_after} seconds" ) diff --git a/bbot/test/test_step_2/module_tests/test_module_teams.py b/bbot/test/test_step_2/module_tests/test_module_teams.py index bd00af650..3f573dc21 100644 --- a/bbot/test/test_step_2/module_tests/test_module_teams.py +++ b/bbot/test/test_step_2/module_tests/test_module_teams.py @@ -7,7 +7,7 @@ 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) @@ -15,6 +15,8 @@ async def setup_after_prep(self, 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 From c3ad0109e8e28577020fa132e9a551baac36a206 Mon Sep 17 00:00:00 2001 From: github-actions Date: Thu, 23 Jan 2025 20:17:22 -0500 Subject: [PATCH 4/8] use builtin retry mechanism --- bbot/modules/base.py | 21 ++++++++++++++++--- bbot/modules/output/teams.py | 35 +++++++------------------------ bbot/modules/templates/webhook.py | 34 +++++++++++------------------- 3 files changed, 37 insertions(+), 53 deletions(-) diff --git a/bbot/modules/base.py b/bbot/modules/base.py index ccef6e1e7..37bfca899 100644 --- a/bbot/modules/base.py +++ b/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. diff --git a/bbot/modules/output/teams.py b/bbot/modules/output/teams.py index bfebf73b5..5916ab9b7 100644 --- a/bbot/modules/output/teams.py +++ b/bbot/modules/output/teams.py @@ -8,41 +8,20 @@ class Teams(WebhookOutputModule): "created_date": "2023-08-14", "author": "@TheTechromancer", } - options = {"webhook_url": "", "event_types": ["VULNERABILITY", "FINDING"], "min_severity": "LOW", "retries": 10} + options = {"webhook_url": "", "event_types": ["VULNERABILITY", "FINDING"], "min_severity": "LOW"} 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 (Default: 10)", } - _module_threads = 5 - - async def setup(self): - self.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: diff --git a/bbot/modules/templates/webhook.py b/bbot/modules/templates/webhook.py index 32f10eef3..c4c91df0c 100644 --- a/bbot/modules/templates/webhook.py +++ b/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_failure_retry_limit = 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: From fe580ef8a39911ad8b88ba8bbd8444e912638c91 Mon Sep 17 00:00:00 2001 From: github-actions Date: Thu, 23 Jan 2025 20:22:11 -0500 Subject: [PATCH 5/8] reintroduce retries option --- bbot/modules/output/teams.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/bbot/modules/output/teams.py b/bbot/modules/output/teams.py index 5916ab9b7..e6f3a20d9 100644 --- a/bbot/modules/output/teams.py +++ b/bbot/modules/output/teams.py @@ -8,13 +8,18 @@ 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 (Default: 10)", } + async def setup(self): + self._api_retries = self.config.get("retries", 10) + return await super().setup() + async def handle_event(self, event): data = self.format_message(event) await self.api_request( From c2cb3e117eaf55fa0a0d45e31d28a236562d71a9 Mon Sep 17 00:00:00 2001 From: github-actions Date: Thu, 23 Jan 2025 20:23:03 -0500 Subject: [PATCH 6/8] fix api retries --- bbot/modules/templates/webhook.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bbot/modules/templates/webhook.py b/bbot/modules/templates/webhook.py index c4c91df0c..19d11e28f 100644 --- a/bbot/modules/templates/webhook.py +++ b/bbot/modules/templates/webhook.py @@ -16,7 +16,7 @@ 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_failure_retry_limit = 10 + _api_retries = 10 async def setup(self): self.webhook_url = self.config.get("webhook_url", "") From 7f93b098f7ccf89ee1728dba77ad45ad9c939e77 Mon Sep 17 00:00:00 2001 From: github-actions Date: Fri, 24 Jan 2025 10:03:13 -0500 Subject: [PATCH 7/8] support retries option on all webhook modules --- bbot/modules/output/discord.py | 3 ++- bbot/modules/output/slack.py | 3 ++- bbot/modules/output/teams.py | 6 +----- bbot/modules/templates/webhook.py | 4 ++++ 4 files changed, 9 insertions(+), 7 deletions(-) diff --git a/bbot/modules/output/discord.py b/bbot/modules/output/discord.py index 3a921a900..2aa4d21f8 100644 --- a/bbot/modules/output/discord.py +++ b/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", } diff --git a/bbot/modules/output/slack.py b/bbot/modules/output/slack.py index 5d4769554..d65c816b3 100644 --- a/bbot/modules/output/slack.py +++ b/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" diff --git a/bbot/modules/output/teams.py b/bbot/modules/output/teams.py index e6f3a20d9..c9a7cf182 100644 --- a/bbot/modules/output/teams.py +++ b/bbot/modules/output/teams.py @@ -13,13 +13,9 @@ class Teams(WebhookOutputModule): "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 (Default: 10)", + "retries": "Number of times to retry sending the message before skipping the event", } - async def setup(self): - self._api_retries = self.config.get("retries", 10) - return await super().setup() - async def handle_event(self, event): data = self.format_message(event) await self.api_request( diff --git a/bbot/modules/templates/webhook.py b/bbot/modules/templates/webhook.py index 19d11e28f..caac575dc 100644 --- a/bbot/modules/templates/webhook.py +++ b/bbot/modules/templates/webhook.py @@ -18,6 +18,10 @@ class WebhookOutputModule(BaseOutputModule): # 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) + return await super().setup() + async def setup(self): self.webhook_url = self.config.get("webhook_url", "") self.min_severity = self.config.get("min_severity", "LOW").strip().upper() From d8ad85ab766d2f1b65433e9e2576554f4dda4146 Mon Sep 17 00:00:00 2001 From: github-actions Date: Fri, 24 Jan 2025 10:06:47 -0500 Subject: [PATCH 8/8] fix setup issue --- bbot/modules/templates/webhook.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/bbot/modules/templates/webhook.py b/bbot/modules/templates/webhook.py index caac575dc..12df037ae 100644 --- a/bbot/modules/templates/webhook.py +++ b/bbot/modules/templates/webhook.py @@ -20,9 +20,6 @@ class WebhookOutputModule(BaseOutputModule): async def setup(self): self._api_retries = self.config.get("retries", 10) - return await super().setup() - - async def setup(self): self.webhook_url = self.config.get("webhook_url", "") self.min_severity = self.config.get("min_severity", "LOW").strip().upper() assert ( @@ -32,7 +29,7 @@ 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): message = self.format_message(event)