265 lines
10 KiB
Python
265 lines
10 KiB
Python
"""
|
|
Unit tests for ComfyUI-DiscordSend utilities.
|
|
|
|
Run tests with: python -m pytest tests/test_utils.py -v
|
|
Or without pytest: python tests/test_utils.py
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import unittest
|
|
from unittest.mock import patch, MagicMock
|
|
import requests
|
|
|
|
# Mock dependencies before importing project modules
|
|
sys.modules["torch"] = MagicMock()
|
|
sys.modules["numpy"] = MagicMock()
|
|
sys.modules["cv2"] = MagicMock()
|
|
# sys.modules["PIL"] = MagicMock() # PIL might be installed, so maybe not mock it if not needed, but safer to mock if we don't rely on it for these tests
|
|
|
|
# Add parent directory to path for imports
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from shared.workflow.sanitizer import sanitize_json_for_export
|
|
from shared.discord.webhook_client import validate_webhook_url, sanitize_webhook_for_logging, send_to_discord_with_retry, DiscordWebhookClient
|
|
from shared.github_integration import update_github_cdn_urls
|
|
|
|
|
|
class TestSanitizer(unittest.TestCase):
|
|
"""Tests for the sanitize_json_for_export function."""
|
|
|
|
def test_sanitize_none(self):
|
|
"""Should handle None input."""
|
|
self.assertIsNone(sanitize_json_for_export(None))
|
|
|
|
def test_sanitize_webhook_url_in_string(self):
|
|
"""Should remove webhook URLs from strings."""
|
|
test_url = "https://discord.com/api/webhooks/123456789/abcdefg"
|
|
result = sanitize_json_for_export(test_url)
|
|
self.assertEqual(result, "")
|
|
|
|
def test_sanitize_github_token_in_string(self):
|
|
"""Should remove GitHub tokens from strings."""
|
|
test_token = "ghp_1234567890abcdefghijklmnopqrstuvwxyz"
|
|
result = sanitize_json_for_export(test_token)
|
|
self.assertEqual(result, "")
|
|
|
|
def test_sanitize_dict_with_webhook_key(self):
|
|
"""Should remove values for webhook-related keys."""
|
|
test_data = {
|
|
"webhook_url": "https://discord.com/api/webhooks/123/abc",
|
|
"message": "Hello world"
|
|
}
|
|
result = sanitize_json_for_export(test_data)
|
|
self.assertEqual(result["webhook_url"], "")
|
|
self.assertEqual(result["message"], "Hello world")
|
|
|
|
def test_sanitize_dict_with_github_token_key(self):
|
|
"""Should remove values for github_token keys."""
|
|
test_data = {
|
|
"github_token": "ghp_secret123",
|
|
"repo": "user/repo"
|
|
}
|
|
result = sanitize_json_for_export(test_data)
|
|
self.assertEqual(result["github_token"], "")
|
|
self.assertEqual(result["repo"], "user/repo")
|
|
|
|
def test_sanitize_nested_dict(self):
|
|
"""Should sanitize nested dictionaries."""
|
|
test_data = {
|
|
"level1": {
|
|
"webhook_url": "https://discord.com/api/webhooks/123/abc",
|
|
"level2": {
|
|
"github_token": "ghp_secret"
|
|
}
|
|
}
|
|
}
|
|
result = sanitize_json_for_export(test_data)
|
|
self.assertEqual(result["level1"]["webhook_url"], "")
|
|
self.assertEqual(result["level1"]["level2"]["github_token"], "")
|
|
|
|
def test_sanitize_list(self):
|
|
"""Should sanitize lists."""
|
|
test_data = [
|
|
"normal string",
|
|
"https://discord.com/api/webhooks/123/abc",
|
|
{"webhook_url": "secret"}
|
|
]
|
|
result = sanitize_json_for_export(test_data)
|
|
self.assertEqual(result[0], "normal string")
|
|
self.assertEqual(result[1], "")
|
|
self.assertEqual(result[2]["webhook_url"], "")
|
|
|
|
def test_preserve_normal_data(self):
|
|
"""Should preserve non-sensitive data."""
|
|
test_data = {
|
|
"prompt": "a beautiful sunset",
|
|
"seed": 12345,
|
|
"steps": 20,
|
|
"cfg": 7.5
|
|
}
|
|
result = sanitize_json_for_export(test_data)
|
|
self.assertEqual(result, test_data)
|
|
|
|
|
|
class TestWebhookValidation(unittest.TestCase):
|
|
"""Tests for webhook URL validation."""
|
|
|
|
def test_valid_webhook_url(self):
|
|
"""Should accept valid Discord webhook URLs."""
|
|
valid_url = "https://discord.com/api/webhooks/1234567890/abcdefghijklmnop"
|
|
is_valid, message = validate_webhook_url(valid_url)
|
|
self.assertTrue(is_valid)
|
|
|
|
def test_valid_discordapp_url(self):
|
|
"""Should accept discordapp.com URLs."""
|
|
valid_url = "https://discordapp.com/api/webhooks/1234567890/abcdefghijklmnop"
|
|
is_valid, message = validate_webhook_url(valid_url)
|
|
self.assertTrue(is_valid)
|
|
|
|
def test_empty_url(self):
|
|
"""Should reject empty URLs."""
|
|
is_valid, message = validate_webhook_url("")
|
|
self.assertFalse(is_valid)
|
|
|
|
def test_invalid_url(self):
|
|
"""Should reject non-webhook URLs."""
|
|
is_valid, message = validate_webhook_url("https://example.com")
|
|
self.assertFalse(is_valid)
|
|
|
|
def test_bypass_attempt_url(self):
|
|
"""Should reject URLs that attempt to bypass validation."""
|
|
# This URL contains 'discord' and 'webhook' but is not hosted on discord.com
|
|
bypass_url = "http://evil-site.com/discord/webhook"
|
|
is_valid, message = validate_webhook_url(bypass_url)
|
|
self.assertFalse(is_valid)
|
|
|
|
def test_localhost_url(self):
|
|
"""Should reject localhost URLs (SSRF protection)."""
|
|
is_valid, message = validate_webhook_url("http://localhost:8080/admin")
|
|
self.assertFalse(is_valid)
|
|
|
|
def test_http_url_rejected(self):
|
|
"""Should reject HTTP URLs (must be HTTPS)."""
|
|
is_valid, message = validate_webhook_url("http://discord.com/api/webhooks/123/abc")
|
|
self.assertFalse(is_valid)
|
|
self.assertIn("must start with https://", message)
|
|
|
|
def test_ip_encoding_urls(self):
|
|
"""Should reject alternate IP encodings."""
|
|
self.assertFalse(validate_webhook_url("http://127.0.0.1")[0])
|
|
self.assertFalse(validate_webhook_url("http://0177.0.0.1")[0]) # Octal
|
|
self.assertFalse(validate_webhook_url("http://0x7f.0.0.1")[0]) # Hex
|
|
self.assertFalse(validate_webhook_url("http://[::1]")[0]) # IPv6
|
|
|
|
def test_domain_spoofing_urls(self):
|
|
"""Should reject domains that contain 'discord' but aren't Discord."""
|
|
self.assertFalse(validate_webhook_url("https://discord.com.evil.co/api/webhooks/123/abc")[0])
|
|
self.assertFalse(validate_webhook_url("https://evil-discord.com/api/webhooks/123/abc")[0])
|
|
|
|
def test_path_traversal_urls(self):
|
|
"""Should reject path traversal attempts."""
|
|
self.assertFalse(validate_webhook_url("https://discord.com/api/webhooks/123/abc/../../admin")[0])
|
|
self.assertFalse(validate_webhook_url("https://discord.com/api/webhooks/123/abc%2f..%2f..%2fadmin")[0])
|
|
|
|
|
|
class TestSSRFPrevention(unittest.TestCase):
|
|
"""Tests for SSRF prevention mechanisms."""
|
|
|
|
def test_send_to_discord_validates_url(self):
|
|
"""Should raise ValueError for invalid URLs before sending request."""
|
|
malicious_url = "http://localhost:8080/admin/delete"
|
|
|
|
# We don't need to mock requests.post because it should fail before calling it
|
|
with self.assertRaises(ValueError) as cm:
|
|
send_to_discord_with_retry(malicious_url, data={"content": "test"})
|
|
|
|
self.assertIn("Invalid webhook URL", str(cm.exception))
|
|
|
|
@patch('shared.discord.webhook_client.requests.post')
|
|
def test_send_to_discord_allows_valid_url(self, mock_post):
|
|
"""Should allow valid Discord URLs."""
|
|
valid_url = "https://discord.com/api/webhooks/123/abc"
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_post.return_value = mock_response
|
|
|
|
send_to_discord_with_retry(valid_url, data={"content": "test"})
|
|
|
|
mock_post.assert_called_once()
|
|
|
|
|
|
class TestWebhookSanitization(unittest.TestCase):
|
|
"""Tests for webhook URL sanitization for logging."""
|
|
|
|
def test_sanitize_webhook_for_logging(self):
|
|
"""Should redact the token portion of webhook URLs."""
|
|
url = "https://discord.com/api/webhooks/123456789/secrettoken"
|
|
result = sanitize_webhook_for_logging(url)
|
|
self.assertIn("[REDACTED]", result)
|
|
self.assertIn("123456789", result)
|
|
self.assertNotIn("secrettoken", result)
|
|
|
|
def test_sanitize_empty_url(self):
|
|
"""Should handle empty URLs."""
|
|
result = sanitize_webhook_for_logging("")
|
|
self.assertEqual(result, "")
|
|
|
|
|
|
class TestDiscordWebhookClient(unittest.TestCase):
|
|
"""Tests for DiscordWebhookClient security features."""
|
|
|
|
@patch('shared.discord.webhook_client.requests.post')
|
|
def test_exception_token_leakage(self, mock_post):
|
|
"""Should redact tokens from exception messages in last_error."""
|
|
token = "SUPER_SECRET_TOKEN"
|
|
url = f"https://discord.com/api/webhooks/123456/{token}"
|
|
client = DiscordWebhookClient(url)
|
|
|
|
# Configure mock to raise an exception containing the token
|
|
error_message = f"Max retries exceeded with url: /api/webhooks/123456/{token}"
|
|
mock_post.side_effect = requests.exceptions.ConnectionError(error_message)
|
|
|
|
success, result = client.send_message("Test message")
|
|
|
|
self.assertFalse(success)
|
|
self.assertIn("error", result)
|
|
self.assertNotIn(token, result["error"])
|
|
self.assertIn("[REDACTED]", result["error"])
|
|
|
|
|
|
class TestGitHubIntegration(unittest.TestCase):
|
|
"""Tests for GitHub integration security features."""
|
|
|
|
@patch('shared.github_integration.requests.put')
|
|
@patch('shared.github_integration.requests.get')
|
|
def test_github_token_redaction_in_response(self, mock_get, mock_put):
|
|
"""Should redact GitHub token from error messages including response text."""
|
|
token = "ghp_SECRET_TOKEN"
|
|
repo = "user/repo"
|
|
file_path = "cdn_urls.md"
|
|
|
|
# Mock GET to return 404 (file doesn't exist)
|
|
mock_get_response = MagicMock()
|
|
mock_get_response.status_code = 404
|
|
mock_get.return_value = mock_get_response
|
|
|
|
# Mock PUT to fail and return the token in the response text (simulating leak)
|
|
mock_put_response = MagicMock()
|
|
mock_put_response.status_code = 401
|
|
mock_put_response.text = f"Bad credentials: {token} is invalid"
|
|
mock_put.return_value = mock_put_response
|
|
|
|
success, message = update_github_cdn_urls(repo, token, file_path, [("test.png", "http://url")])
|
|
|
|
self.assertFalse(success)
|
|
self.assertNotIn(token, message)
|
|
self.assertIn("[REDACTED_TOKEN]", message)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Run tests
|
|
print("Running ComfyUI-DiscordSend utility tests...\n")
|
|
unittest.main(verbosity=2)
|