attachment_guard.py

"""
attachment_guard.py — Protection against Indirect Prompt Injection via files.

Handles extraction and security screening of uploaded attachments.
V2: Fixed memory exhaustion DoS by checking size BEFORE base64 decoding.
"""

import os
import base64
from typing import Any, Dict


class AttachmentGuard:
    """
    Utility for validating and extracting text from uploaded attachments.
    """

    ALLOWED_EXTENSIONS = {'.txt', '.csv', '.md', '.json', '.py', '.js', '.html', '.css'}
    MAX_FILE_SIZE = 1 * 1024 * 1024  # 1MB limit for extraction

    @staticmethod
    def extract_text(filename: str, content_b64: str) -> Dict[str, Any]:
        """
        Extract text from a base64 encoded file.

        Returns:
            { "text": str, "error": str|None, "extension": str }
        """
        ext = os.path.splitext(filename)[1].lower()

        if ext not in AttachmentGuard.ALLOWED_EXTENSIONS:
            return {
                "text": "",
                "error": f"Unsupported file type: {ext}. Only text-based files allowed.",
                "extension": ext
            }

        # ── DoS fix: check estimated decoded size BEFORE decoding ─────────
        # Base64 encoding inflates size by ~33%, so decoded ≈ len(b64) * 3/4
        estimated_size = len(content_b64) * 3 // 4
        if estimated_size > AttachmentGuard.MAX_FILE_SIZE:
            return {
                "text": "",
                "error": f"File too large (estimated {estimated_size // 1024}KB, max 1MB).",
                "extension": ext
            }
        # ─────────────────────────────────────────────────────────────────

        try:
            file_bytes = base64.b64decode(content_b64)

            if len(file_bytes) > AttachmentGuard.MAX_FILE_SIZE:
                return {
                    "text": "",
                    "error": "File too large (max 1MB).",
                    "extension": ext
                }

            # Try to decode as utf-8
            text = file_bytes.decode('utf-8')

            return {
                "text": text,
                "error": None,
                "extension": ext
            }
        except Exception as e:
            return {
                "text": "",
                "error": f"Failed to extract text: {str(e)}",
                "extension": ext
            }

    @staticmethod
    def screen_with_guard(guard: Any, filename: str, text: str) -> Dict[str, Any]:
        """
        Run the PromptGuardTextGuard screening on extracted text.
        The guard's screen() method internally uses chunked scanning,
        so documents of arbitrary length are handled correctly.
        """
        if not text.strip():
            return {"blocked": False, "reason": "Empty attachment", "threat_score": 0, "flags": []}

        result = guard.screen(text)

        if result["blocked"]:
            result["reason"] = f"MALICIOUS_ATTACHMENT ({filename}): {result['reason']}"
            result["flags"].append("INDIRECT_PROMPT_INJECTION")

        return result