"""Attachment storage layer — disk-based file storage with lightweight references. Files are stored on disk under `{opc_home}/projects/{project_id}/attachments/{id}/`. Only lightweight AttachmentRef objects (no binary content) flow through the engine pipeline, databases, and WebSocket messages. Base64 encoding is performed lazily at the two endpoints: ingestion (decode → disk) and LLM call (disk → encode). """ from __future__ import annotations import base64 import mimetypes import os import shutil import uuid from dataclasses import dataclass from pathlib import Path from typing import Any ALLOWED_MIME_PREFIXES = ( "image/", "text/", "application/pdf", "application/json", "application/x-yaml", "application/yaml", ) ALLOWED_EXTENSIONS = { ".png", ".jpg", ".jpeg", ".gif", ".webp", ".svg", ".bmp", ".tiff", ".txt", ".md", ".pdf", ".csv", ".json", ".yaml", ".yml", ".py", ".js", ".ts", ".tsx", ".jsx", ".html", ".css", ".java", ".c", ".cpp", ".h", ".go", ".rs", ".rb", ".sh", ".xml", ".toml", ".ini", ".cfg", ".log", ".docx", ".xlsx", ".pptx", ".mp4", ".mpeg", ".mpg", ".mov", ".webm", } MAX_FILE_SIZE = 10 * 1024 * 1024 # 10 MB per file MAX_TOTAL_SIZE = 20 * 1024 * 1024 # 20 MB per message ALLOWED_MIME_TYPES = { "video/mp4", "video/mpeg", "video/quicktime", "video/webm", } _MIME_EXTENSION_OVERRIDES = { "image/jpeg": ".jpg", "image/png": ".png", "image/gif": ".gif", "image/webp": ".webp", "image/svg+xml": ".svg", "image/bmp": ".bmp", "image/tiff": ".tiff", "application/pdf": ".pdf", "text/plain": ".txt", "text/markdown": ".md", "application/json": ".json", "application/x-yaml": ".yaml", "application/yaml": ".yaml", "video/mp4": ".mp4", "video/mpeg": ".mpeg", "video/quicktime": ".mov", "video/webm": ".webm", } @dataclass class AttachmentRef: """Lightweight reference to a stored attachment — safe for JSON / metadata.""" attachment_id: str filename: str mime_type: str size_bytes: int disk_path: str # relative to opc_home @property def is_image(self) -> bool: return self.mime_type.startswith("image/") def to_dict(self) -> dict[str, Any]: return { "attachment_id": self.attachment_id, "filename": self.filename, "mime_type": self.mime_type, "size_bytes": self.size_bytes, "disk_path": self.disk_path, } @classmethod def from_dict(cls, d: dict[str, Any]) -> AttachmentRef: return cls( attachment_id=d["attachment_id"], filename=d["filename"], mime_type=d.get("mime_type", "application/octet-stream"), size_bytes=d.get("size_bytes", 0), disk_path=d.get("disk_path", ""), ) def _sanitize_filename(name: str) -> str: """Strip path components and dangerous characters from a filename.""" name = os.path.basename(name) name = name.replace("..", "").replace("/", "").replace("\\", "") return name or "upload" def _check_mime(filename: str, mime: str) -> bool: ext = os.path.splitext(filename)[1].lower() if ext in ALLOWED_EXTENSIONS: return True if mime in ALLOWED_MIME_TYPES: return True for prefix in ALLOWED_MIME_PREFIXES: if mime.startswith(prefix): return True return False def _normalize_mime(mime_type: str | None) -> str: return str(mime_type or "").strip().lower() def _split_data_url(payload: str) -> tuple[str, str | None]: raw = str(payload or "").strip() if not raw.startswith("data:"): return raw, None header, sep, encoded = raw.partition(",") if not sep: return raw, None mime = header[5:].split(";", 1)[0].strip().lower() return encoded, mime or None def _extension_for_mime(mime_type: str) -> str: if not mime_type: return "" override = _MIME_EXTENSION_OVERRIDES.get(mime_type) if override: return override guessed = mimetypes.guess_extension(mime_type, strict=False) or "" return guessed.lower() def _ensure_filename_extension(filename: str, mime_type: str) -> str: if os.path.splitext(filename)[1]: return filename extension = _extension_for_mime(mime_type) if extension: return f"{filename}{extension}" return filename class AttachmentStore: """Manages attachment lifecycle on disk.""" def __init__(self, opc_home: Path, project_id: str) -> None: self.opc_home = opc_home self.project_id = project_id self.base_dir = opc_home / "projects" / project_id / "attachments" def _ensure_dir(self, attachment_id: str) -> Path: d = self.base_dir / attachment_id d.mkdir(parents=True, exist_ok=True) return d async def save_from_base64( self, filename: str, b64_data: str, mime_type: str | None = None, ) -> AttachmentRef: """Decode base64 data and persist to disk. Returns a lightweight ref.""" filename = _sanitize_filename(filename) b64_payload, inferred_mime = _split_data_url(b64_data) mime = _normalize_mime(mime_type) or inferred_mime or "" if not mime: guessed_mime, _ = mimetypes.guess_type(filename) mime = guessed_mime or "application/octet-stream" filename = _ensure_filename_extension(filename, mime) if mime == "application/octet-stream": guessed_mime, _ = mimetypes.guess_type(filename) mime = guessed_mime or mime try: raw = base64.b64decode(b64_payload) except Exception as exc: raise ValueError(f"Invalid base64 data for {filename}: {exc}") from exc size = len(raw) if size > MAX_FILE_SIZE: raise ValueError(f"File too large: {size} bytes (limit {MAX_FILE_SIZE})") if not _check_mime(filename, mime): raise ValueError(f"Unsupported file type: {filename} ({mime})") aid = uuid.uuid4().hex[:16] dest_dir = self._ensure_dir(aid) dest = dest_dir / filename dest.write_bytes(raw) rel = dest.relative_to(self.opc_home) return AttachmentRef( attachment_id=aid, filename=filename, mime_type=mime, size_bytes=size, disk_path=str(rel), ) async def save_from_path(self, file_path: Path) -> AttachmentRef: """Copy a local file into the attachment store. Used by CLI.""" file_path = file_path.expanduser().resolve() if not file_path.is_file(): raise FileNotFoundError(f"File not found: {file_path}") size = file_path.stat().st_size if size > MAX_FILE_SIZE: raise ValueError(f"File too large: {size} bytes (limit {MAX_FILE_SIZE})") filename = _sanitize_filename(file_path.name) mime, _ = mimetypes.guess_type(filename) mime = mime or "application/octet-stream" if not _check_mime(filename, mime): raise ValueError(f"Unsupported file type: {filename} ({mime})") aid = uuid.uuid4().hex[:16] dest_dir = self._ensure_dir(aid) dest = dest_dir / filename shutil.copy2(str(file_path), str(dest)) rel = dest.relative_to(self.opc_home) return AttachmentRef( attachment_id=aid, filename=filename, mime_type=mime, size_bytes=size, disk_path=str(rel), ) def resolve_abs_path(self, ref: AttachmentRef) -> Path: """Return the absolute path on disk for reading.""" resolved = (self.opc_home / ref.disk_path).resolve() if not str(resolved).startswith(str(self.base_dir.resolve())): raise ValueError(f"Path traversal detected: {ref.disk_path}") return resolved def read_bytes(self, ref: AttachmentRef) -> bytes: """Read raw file content from disk.""" return self.resolve_abs_path(ref).read_bytes() def read_base64(self, ref: AttachmentRef) -> str: """Read file and return as base64 string — called only at LLM call time.""" return base64.b64encode(self.read_bytes(ref)).decode("ascii") def resolve_http_path(self, ref: AttachmentRef) -> str: """Return the HTTP-accessible path for frontend display.""" return f"/api/attachments/{ref.attachment_id}/{ref.filename}"