Spaces:
Runtime error
Runtime error
File size: 3,186 Bytes
0d599d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """
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
|