diff --git a/nodes/image_node.py b/nodes/image_node.py index 0ffabc5..da76bd6 100644 --- a/nodes/image_node.py +++ b/nodes/image_node.py @@ -19,6 +19,7 @@ from shared import ( update_github_cdn_urls, extract_prompts_from_workflow, send_to_discord_with_retry, + sanitize_token_from_text, tensor_to_numpy_uint8, build_filename_with_metadata, get_output_directory, @@ -771,7 +772,8 @@ class DiscordSendSaveImage: message="Discord CDN URLs for the uploaded images:" ) else: - print(f"Error sending batch to Discord: Status code {response.status_code} - {response.text}") + error_msg = sanitize_token_from_text(response.text, webhook_url) + print(f"Error sending batch to Discord: Status code {response.status_code} - {error_msg}") discord_send_success = False except Exception as e: print(f"Error sending batch to Discord: {e}") diff --git a/nodes/video_node.py b/nodes/video_node.py index 952520d..3420141 100644 --- a/nodes/video_node.py +++ b/nodes/video_node.py @@ -26,6 +26,7 @@ from shared import ( sanitize_json_for_export, update_github_cdn_urls, send_to_discord_with_retry, + sanitize_token_from_text, tensor_to_numpy_uint8, build_filename_with_metadata, get_output_directory, @@ -984,7 +985,8 @@ class DiscordSendSaveVideo: new_urls = extract_cdn_urls_from_response(response) discord_cdn_urls.extend(new_urls) else: - print(f"Discord API error: {response.status_code} - {response.text}") + error_msg = sanitize_token_from_text(response.text, webhook_url) + print(f"Discord API error: {response.status_code} - {error_msg}") discord_send_success = False except Exception as e: diff --git a/shared/__init__.py b/shared/__init__.py index 4b89ccc..682b809 100644 --- a/shared/__init__.py +++ b/shared/__init__.py @@ -15,7 +15,8 @@ from .workflow.workflow_builder import WorkflowBuilder from .discord.webhook_client import ( DiscordWebhookClient, validate_webhook_url, - send_to_discord_with_retry + send_to_discord_with_retry, + sanitize_token_from_text ) from .discord.message_builder import ( build_metadata_section, @@ -42,6 +43,7 @@ __all__ = [ 'DiscordWebhookClient', 'validate_webhook_url', 'send_to_discord_with_retry', + 'sanitize_token_from_text', # Discord message building 'build_metadata_section', 'build_prompt_section', diff --git a/shared/discord/webhook_client.py b/shared/discord/webhook_client.py index 27ac77c..59899fc 100644 --- a/shared/discord/webhook_client.py +++ b/shared/discord/webhook_client.py @@ -70,6 +70,31 @@ def sanitize_webhook_for_logging(url: str) -> str: return "[REDACTED_WEBHOOK_URL]" +def sanitize_token_from_text(text: str, webhook_url: str) -> str: + """ + Sanitize the webhook token from arbitrary text. + + Args: + text: The text to sanitize + webhook_url: The webhook URL containing the token + + Returns: + Text with the token replaced by [REDACTED] + """ + if not text or not webhook_url: + return text + + # Pattern: https://discord.com/api/webhooks/{id}/{token} + # Use case-insensitive matching to handle potential variations + match = re.search(r"/api/webhooks/\d+/([\w-]+)", webhook_url, re.IGNORECASE) + if match: + token = match.group(1) + if token in text: + return text.replace(token, "[REDACTED]") + + return text + + class DiscordWebhookClient: """ Client for sending messages and files to Discord via webhooks. @@ -243,9 +268,10 @@ class DiscordWebhookClient: # Client errors (don't retry) if 400 <= response.status_code < 500: + sanitized_details = sanitize_token_from_text(response.text[:500], self.webhook_url) return False, { "error": f"Discord API error: {response.status_code}", - "details": response.text[:500] + "details": sanitized_details } # Server errors (retry) @@ -255,14 +281,7 @@ class DiscordWebhookClient: last_error = "Request timed out" except requests.exceptions.RequestException as e: # Sanitize error message to prevent token leakage - error_msg = str(e) - # Use case-insensitive matching to handle uppercase URLs - match = re.search(r"/api/webhooks/\d+/([\w-]+)", self.webhook_url, re.IGNORECASE) - if match: - token = match.group(1) - if token in error_msg: - error_msg = error_msg.replace(token, "[REDACTED]") - + error_msg = sanitize_token_from_text(str(e), self.webhook_url) last_error = error_msg # Exponential backoff @@ -405,17 +424,12 @@ def send_to_discord_with_retry( last_exception = requests.exceptions.Timeout("Discord request timed out") except requests.exceptions.RequestException as e: # Sanitize error message to prevent token leakage - error_msg = str(e) - # Use case-insensitive matching to handle uppercase URLs - match = re.search(r"/api/webhooks/\d+/([\w-]+)", webhook_url, re.IGNORECASE) - if match: - token = match.group(1) - if token in error_msg: - error_msg = error_msg.replace(token, "[REDACTED]") - + error_msg = sanitize_token_from_text(str(e), webhook_url) logger.warning(f"Request error: {error_msg}, attempt {attempt + 1}/{max_retries}") # Store sanitized exception to avoid leaking token if raised later + # Use case-insensitive matching to handle uppercase URLs + match = re.search(r"/api/webhooks/\d+/([\w-]+)", webhook_url, re.IGNORECASE) if match and match.group(1) in str(e): # Create a new exception of the same type with sanitized message # We try to preserve the exception type, but fallback to RequestException if init fails diff --git a/tests/test_webhook_security.py b/tests/test_webhook_security.py new file mode 100644 index 0000000..e7775b6 --- /dev/null +++ b/tests/test_webhook_security.py @@ -0,0 +1,70 @@ + +import unittest +from unittest.mock import MagicMock, patch +import sys +import os + +# Mock modules that might be missing in the environment or causing issues +sys.modules["folder_paths"] = MagicMock() +sys.modules["comfy"] = MagicMock() +sys.modules["server"] = MagicMock() +sys.modules["torch"] = MagicMock() +sys.modules["numpy"] = MagicMock() +sys.modules["cv2"] = MagicMock() + +# Import the module under test +sys.path.append(os.getcwd()) + +from shared.discord.webhook_client import DiscordWebhookClient, send_to_discord_with_retry, sanitize_token_from_text + +class TestWebhookSecurity(unittest.TestCase): + def test_token_leak_in_client_error(self): + """ + Test that webhook tokens are NOT leaked in client error details. + """ + webhook_url = "https://discord.com/api/webhooks/123456789/SuperSecretToken123" + client = DiscordWebhookClient(webhook_url) + + # Mock response to simulate a 400 error that echoes the URL + mock_response = MagicMock() + mock_response.status_code = 400 + # Simulate an API that echoes the request URL in the error body + mock_response.text = f"Error processing request to {webhook_url}: Invalid payload" + mock_response.content = mock_response.text.encode('utf-8') + + with patch('requests.post', return_value=mock_response): + success, response = client.send_message("test") + + self.assertFalse(success) + error_details = response.get("details", "") + + print(f"\nDEBUG: Error details: {error_details}") + self.assertNotIn("SuperSecretToken123", error_details) + self.assertIn("[REDACTED]", error_details) + + def test_sanitize_token_from_text(self): + """ + Test the standalone sanitization function. + """ + webhook_url = "https://discord.com/api/webhooks/123456789/MySecretToken-Part2" + + # Test 1: Simple URL in text + text = f"Failed to send to {webhook_url}" + sanitized = sanitize_token_from_text(text, webhook_url) + self.assertNotIn("MySecretToken-Part2", sanitized) + self.assertIn("[REDACTED]", sanitized) + + # Test 2: Token embedded in other text + text = "Some error occurred with token MySecretToken-Part2 processing" + sanitized = sanitize_token_from_text(text, webhook_url) + self.assertNotIn("MySecretToken-Part2", sanitized) + self.assertIn("[REDACTED]", sanitized) + + # Test 3: Multiple occurrences + text = f"URL: {webhook_url}, Retry: {webhook_url}" + sanitized = sanitize_token_from_text(text, webhook_url) + self.assertNotIn("MySecretToken-Part2", sanitized) + self.assertEqual(sanitized.count("[REDACTED]"), 2) + +if __name__ == '__main__': + unittest.main()