Fix bugs in error handling, URL filtering, and response validation

AGENTS.md

@@ -12,7 +12,8 @@ AI-powered spam detection tool that integrates with rspamd using AI providers (A
 │   ├── prompts.py                 # Shared SYSTEM_PROMPT constant
 │   ├── ai_spam_check.py           # Provider dispatcher (check_spam_with_ai)
 │   ├── anthropic_spam_check.py    # Anthropic provider implementation
-│   └── openai_spam_check.py       # OpenAI provider implementation
+│   ├── openai_spam_check.py       # OpenAI provider implementation
+│   ├── validation.py              # AI response validation/normalization
 │   └── process_email.py           # Email parsing and formatting
 ├── rspamd/
 │   └── external_ai_test.lua       # rspamd Lua plugin

cli_toolbox.py

@@ -57,7 +57,7 @@ def main() -> None:
                     ai_output += f"Confidence: {result.get('confidence', 'Unknown')}\n"
                     ai_output += f"Reason: {result.get('reason', 'No reason provided')}"
 
-                if not result.get("is_spam"):
+                if "is_spam" not in result or "confidence" not in result:
                     error = True
             except Exception as e:
                 output += f"\n\nError performing AI spam check: {e}\n"

email_ai_interface.py

@@ -22,7 +22,7 @@ def main() -> None:
             email_file = Path(sys.argv[1])
             if not email_file.exists():
                 raise FileNotFoundError(f"Email file not found: {email_file}")
-            with open(email_file, encoding="utf-8") as f:
+            with open(email_file, encoding="utf-8", errors="replace") as f:
                 email_content = f.read()
         else:
             # Read email from STDIN if no filename provided

email_utils/anthropic_spam_check.py

@@ -5,6 +5,7 @@
 
 from email_utils.config import load_config
 from email_utils.prompts import SYSTEM_PROMPT
+from email_utils.validation import validate_ai_response
 
 
 def check_spam_with_anthropic(email_content: str) -> dict[str, Any]:
@@ -38,7 +39,7 @@ def check_spam_with_anthropic(email_content: str) -> dict[str, Any]:
 
         response = client.messages.create(**kwargs)
 
-        return json.loads(response.content[0].text)
+        return validate_ai_response(json.loads(response.content[0].text))
 
     except Exception as e:
         return {

email_utils/openai_spam_check.py

@@ -5,6 +5,7 @@
 
 from email_utils.config import load_config
 from email_utils.prompts import SYSTEM_PROMPT
+from email_utils.validation import validate_ai_response
 
 
 def check_spam_with_openai(email_content: str) -> dict[str, Any]:
@@ -40,7 +41,7 @@ def check_spam_with_openai(email_content: str) -> dict[str, Any]:
 
         response = client.chat.completions.create(**kwargs)
 
-        return json.loads(response.choices[0].message.content)
+        return validate_ai_response(json.loads(response.choices[0].message.content))
 
     except Exception as e:
         return {

email_utils/process_email.py

@@ -146,6 +146,8 @@ def get_email_content(source: str | Path, is_string: bool = False) -> dict[str,
                 body_text += payload.decode("utf-8", errors="replace")
 
     # Prefer HTML body over plain text, converting HTML to text
+    # Parse once and reuse for text extraction and link/image collection
+    soup = None
     if html_content:
         soup = BeautifulSoup(html_content, "html.parser")
         body_text = soup.get_text()
@@ -156,16 +158,18 @@ def get_email_content(source: str | Path, is_string: bool = False) -> dict[str,
     # Extract images and links if HTML content exists (with deduplication)
     images = []
     links = []
-    if html_content:
-        soup = BeautifulSoup(html_content, "html.parser")
+    skip_schemes = {"data", "javascript", "vbscript", "mailto"}
+    if soup is not None:
         for img in soup.find_all("img"):
-            src = img.get("src")
-            if src and src not in images:
+            src = (img.get("src") or "").strip()
+            scheme = src.split(":", 1)[0].lower() if ":" in src else ""
+            if src and scheme not in skip_schemes and src not in images:
                 images.append(src)
 
         for a in soup.find_all("a"):
-            href = a.get("href")
-            if href and href not in links:
+            href = (a.get("href") or "").strip()
+            scheme = href.split(":", 1)[0].lower() if ":" in href else ""
+            if href and scheme not in skip_schemes and href not in links:
                 links.append(href)
 
     return {

email_utils/validation.py

@@ -0,0 +1,25 @@
+from typing import Any
+
+
+def validate_ai_response(result: dict[str, Any]) -> dict[str, Any]:
+    """Ensure AI response has required fields with correct types.
+
+    Returns a normalized dict with is_spam, confidence, and reason.
+    Missing or invalid fields are replaced with safe defaults.
+    """
+    is_spam = result.get("is_spam", "no")
+    if is_spam not in ("yes", "no"):
+        is_spam = "no"
+
+    confidence = result.get("confidence", 0)
+    try:
+        confidence = int(confidence)
+    except (TypeError, ValueError):
+        confidence = 0
+    confidence = max(0, min(100, confidence))
+
+    reason = result.get("reason", "No reason provided")
+    if not isinstance(reason, str):
+        reason = str(reason)
+
+    return {"is_spam": is_spam, "confidence": confidence, "reason": reason}

rspamd/external_ai_test.lua

@@ -64,7 +64,7 @@ local function check_with_external_script(task)
 
         local json_result = parser:get_object()
         local is_spam = (json_result.is_spam == "yes")
-        local confidence = tonumber(json_result.confidence) or 0
+        local confidence = tonumber(json_result.confidence or 0) or 0
         local reason = json_result.reason or "No reason provided"
 
         -- Calculate score (confidence/10, range 0-10)

tests/test_process_email.py

@@ -141,6 +141,41 @@ def test_dkim_header_removed(self):
         result = get_email_content(DKIM_SPF_EMAIL, is_string=True)
         assert "DKIM-Signature" in result["removed_headers"]
 
+    def test_data_uris_filtered(self):
+        email_str = (
+            'From: a@b.com\nContent-Type: text/html; charset="utf-8"\n\n'
+            '<html><body><img src="data:image/png;base64,abc123"/>'
+            '<img src="https://example.com/real.png"/></body></html>'
+        )
+        result = get_email_content(email_str, is_string=True)
+        assert result["images"] == ["https://example.com/real.png"]
+
+    def test_javascript_links_filtered(self):
+        email_str = (
+            'From: a@b.com\nContent-Type: text/html; charset="utf-8"\n\n'
+            '<html><body><a href="javascript:alert(1)">Click</a>'
+            '<a href="https://example.com">Real</a></body></html>'
+        )
+        result = get_email_content(email_str, is_string=True)
+        assert result["links"] == ["https://example.com"]
+
+    def test_mailto_links_filtered(self):
+        email_str = (
+            'From: a@b.com\nContent-Type: text/html; charset="utf-8"\n\n'
+            '<html><body><a href="mailto:user@example.com">Email</a>'
+            '<a href="https://example.com">Web</a></body></html>'
+        )
+        result = get_email_content(email_str, is_string=True)
+        assert result["links"] == ["https://example.com"]
+
+    def test_url_whitespace_stripped(self):
+        email_str = (
+            'From: a@b.com\nContent-Type: text/html; charset="utf-8"\n\n'
+            '<html><body><a href="  https://example.com  ">Link</a></body></html>'
+        )
+        result = get_email_content(email_str, is_string=True)
+        assert result["links"] == ["https://example.com"]
+
 
 class TestFormatEmailContent:
     def test_basic_structure(self):

tests/test_validation.py

@@ -0,0 +1,47 @@
+from email_utils.validation import validate_ai_response
+
+
+class TestValidateAiResponse:
+    def test_valid_response_passes_through(self):
+        result = {"is_spam": "yes", "confidence": 85, "reason": "Obvious spam"}
+        assert validate_ai_response(result) == result
+
+    def test_missing_is_spam_defaults_to_no(self):
+        result = {"confidence": 50, "reason": "test"}
+        assert validate_ai_response(result)["is_spam"] == "no"
+
+    def test_invalid_is_spam_defaults_to_no(self):
+        result = {"is_spam": "maybe", "confidence": 50, "reason": "test"}
+        assert validate_ai_response(result)["is_spam"] == "no"
+
+    def test_missing_confidence_defaults_to_zero(self):
+        result = {"is_spam": "yes", "reason": "test"}
+        assert validate_ai_response(result)["confidence"] == 0
+
+    def test_string_confidence_converted(self):
+        result = {"is_spam": "yes", "confidence": "75", "reason": "test"}
+        assert validate_ai_response(result)["confidence"] == 75
+
+    def test_confidence_clamped_to_100(self):
+        result = {"is_spam": "yes", "confidence": 150, "reason": "test"}
+        assert validate_ai_response(result)["confidence"] == 100
+
+    def test_confidence_clamped_to_zero(self):
+        result = {"is_spam": "yes", "confidence": -10, "reason": "test"}
+        assert validate_ai_response(result)["confidence"] == 0
+
+    def test_invalid_confidence_defaults_to_zero(self):
+        result = {"is_spam": "yes", "confidence": "high", "reason": "test"}
+        assert validate_ai_response(result)["confidence"] == 0
+
+    def test_missing_reason_gets_default(self):
+        result = {"is_spam": "no", "confidence": 10}
+        assert validate_ai_response(result)["reason"] == "No reason provided"
+
+    def test_non_string_reason_converted(self):
+        result = {"is_spam": "no", "confidence": 10, "reason": 42}
+        assert validate_ai_response(result)["reason"] == "42"
+
+    def test_empty_dict(self):
+        result = validate_ai_response({})
+        assert result == {"is_spam": "no", "confidence": 0, "reason": "No reason provided"}