"""Media upload security primitives: MIME sniff, size validation, SSRF guard, filename safety.""" from __future__ import annotations import ipaddress import mimetypes import os import re import socket from collections.abc import Iterable from dataclasses import dataclass from urllib.parse import urlparse try: import magic as _magic # type: ignore except Exception: _magic = None class UploadError(Exception): """Structured upload error with stable code for JSON responses.""" def __init__(self, code: str, message: str, details: dict | None = None): super().__init__(message) self.code = code self.message = message self.details = details or {} def to_dict(self) -> dict: return {"error_code": self.code, "message": self.message, "details": self.details} DEFAULT_MAX_BYTES = int(os.environ.get("WP_MEDIA_MAX_MB", "10")) * 1024 * 1024 ALLOWED_MIMES: set[str] = { "image/jpeg", "image/png", "image/webp", "image/gif", "image/avif", "image/heic", "image/bmp", "image/tiff", "application/pdf", "video/mp4", "video/webm", "video/quicktime", "audio/mpeg", "audio/mp4", "audio/webm", "audio/ogg", "audio/wav", } _EXT_MAP = { "image/jpeg": ".jpg", "image/png": ".png", "image/webp": ".webp", "image/gif": ".gif", "image/avif": ".avif", "image/heic": ".heic", "application/pdf": ".pdf", "video/mp4": ".mp4", "video/webm": ".webm", "video/quicktime": ".mov", "audio/mpeg": ".mp3", "audio/mp4": ".m4a", "audio/ogg": ".ogg", "audio/wav": ".wav", } def _builtin_sniff(data: bytes) -> str | None: """Minimal magic-byte sniffer for common media types (libmagic-free fallback).""" if not data: return None if data.startswith(b"\x89PNG\r\n\x1a\n"): return "image/png" if data.startswith(b"\xff\xd8\xff"): return "image/jpeg" if data[:4] == b"GIF8": return "image/gif" if data.startswith(b"%PDF-"): return "application/pdf" if data[:4] == b"RIFF" and data[8:12] == b"WEBP": return "image/webp" if data[:4] == b"RIFF" and data[8:12] == b"WAVE": return "audio/wav" if data[4:12] in (b"ftypmp42", b"ftypisom", b"ftypMSNV", b"ftypavc1"): return "video/mp4" if data[4:8] == b"ftyp": # Generic ISO BMFF — likely mp4/heic brand = data[8:12] if brand in (b"heic", b"heix", b"mif1"): return "image/heic" return "video/mp4" if data.startswith(b"\x1aE\xdf\xa3"): return "video/webm" if data.startswith(b"ID3") or data[:2] == b"\xff\xfb": return "audio/mpeg" if data.startswith(b"OggS"): return "audio/ogg" return None def sniff_mime(data: bytes, *, hint: str | None = None) -> str: """Detect MIME from magic bytes. Uses libmagic if available, else built-in sniff.""" if _magic is not None: try: detected = _magic.from_buffer(data[:4096], mime=True) if detected and detected != "application/octet-stream": return detected except Exception: pass built_in = _builtin_sniff(data) if built_in: return built_in if hint and "/" in hint: return hint.lower() return "application/octet-stream" def validate_size(data: bytes, *, max_bytes: int = DEFAULT_MAX_BYTES) -> None: if len(data) == 0: raise UploadError("EMPTY_FILE", "Upload data is empty.") if len(data) > max_bytes: raise UploadError( "TOO_LARGE", f"File is {len(data)} bytes; limit is {max_bytes} bytes " f"(~{max_bytes // 1024 // 1024} MB). Use chunked upload for larger files.", {"size": len(data), "max": max_bytes}, ) def validate_mime(mime: str, *, allowed: Iterable[str] = ALLOWED_MIMES) -> None: if mime not in allowed: raise UploadError( "MIME_REJECTED", f"MIME type '{mime}' is not allowed. " f"Supported: {', '.join(sorted(allowed))}.", {"mime": mime}, ) _FILENAME_SAFE_RE = re.compile(r"[^A-Za-z0-9._-]+") def safe_filename(filename: str | None, *, mime: str) -> tuple[str, str | None]: """Return (ascii_filename, rfc5987_encoded_original_if_non_ascii).""" original = (filename or "").strip() or "upload" # Strip any path components original = original.rsplit("/", 1)[-1].rsplit("\\", 1)[-1] try: original.encode("ascii") ascii_name = _FILENAME_SAFE_RE.sub("_", original) encoded = None except UnicodeEncodeError: from urllib.parse import quote ascii_name = _FILENAME_SAFE_RE.sub( "_", original.encode("ascii", "ignore").decode() or "upload" ) encoded = "UTF-8''" + quote(original, safe="") # Ensure extension matches MIME ext = _EXT_MAP.get(mime) or (mimetypes.guess_extension(mime) or "") if ext and not ascii_name.lower().endswith(ext): base = ascii_name.rsplit(".", 1)[0] if "." in ascii_name else ascii_name ascii_name = f"{base}{ext}" return ascii_name[:255], encoded def content_disposition(filename_ascii: str, filename_encoded: str | None) -> str: """Build Content-Disposition header, optionally with RFC 5987 filename* for non-ASCII.""" base = f'attachment; filename="{filename_ascii}"' if filename_encoded: return f"{base}; filename*={filename_encoded}" return base # --- SSRF guard ------------------------------------------------------------ _BLOCKED_HOSTS = { "metadata.google.internal", "metadata", "169.254.169.254", } @dataclass class SSRFCheck: allowed: bool reason: str | None = None resolved_ip: str | None = None def ssrf_check(url: str, *, allow_http: bool = False) -> SSRFCheck: """Reject URLs pointing to private/loopback/link-local/metadata endpoints.""" try: parsed = urlparse(url) except Exception as e: return SSRFCheck(False, f"URL parse failed: {e}") if parsed.scheme not in ("http", "https"): return SSRFCheck(False, f"Scheme '{parsed.scheme}' not allowed; use https (or http).") if parsed.scheme == "http" and not allow_http: return SSRFCheck(False, "HTTP URLs are disabled; use HTTPS.") host = (parsed.hostname or "").lower() if not host: return SSRFCheck(False, "URL has no host.") if host in _BLOCKED_HOSTS: return SSRFCheck(False, f"Host '{host}' is on the SSRF blocklist.") # Resolve and check each IP try: infos = socket.getaddrinfo(host, None) except socket.gaierror as e: return SSRFCheck(False, f"DNS resolution failed: {e}") seen: list[str] = [] for info in infos: ip_str = info[4][0] seen.append(ip_str) try: ip = ipaddress.ip_address(ip_str) except ValueError: return SSRFCheck(False, f"Unparseable IP '{ip_str}'.") if ( ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved or ip.is_multicast or ip.is_unspecified ): return SSRFCheck(False, f"Host resolves to disallowed IP {ip}.", ip_str) return SSRFCheck(True, None, seen[0] if seen else None)