feat: validate and unblock media downloads

rf4_research/media_assets.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import hashlib
+import io
 import json
 import mimetypes
 import re
@@ -10,6 +11,9 @@ from pathlib import Path
 from urllib.parse import urljoin, urlsplit
 
 from bs4 import BeautifulSoup, Tag
+from PIL import Image, UnidentifiedImageError
+
+MAX_IMAGE_PIXELS = 40_000_000
 
 
 @dataclass(frozen=True, slots=True)
@@ -91,12 +95,32 @@ def reclassify_manifest(path: Path) -> dict:
     return manifest
 
 
-def store_asset(root: Path, body: bytes, *, content_type: str, source_url: str) -> tuple[str, str]:
+def inspect_image(body: bytes) -> tuple[int, int, str]:
+    try:
+        with Image.open(io.BytesIO(body)) as image:
+            width, height = image.size
+            image_format = image.format
+            image.verify()
+    except (UnidentifiedImageError, OSError, ValueError) as exc:
+        raise ValueError("asset is not a valid raster image") from exc
+    if width < 1 or height < 1 or width * height > MAX_IMAGE_PIXELS:
+        raise ValueError("asset dimensions are outside safe limits")
+    mime = Image.MIME.get(image_format or "")
+    if mime not in {"image/jpeg", "image/png", "image/webp", "image/gif"}:
+        raise ValueError(f"unsupported image format: {image_format}")
+    return width, height, mime
+
+
+def store_asset(root: Path, body: bytes, *, content_type: str, source_url: str) -> tuple[str, str, int, int, str]:
+    width, height, detected_mime = inspect_image(body)
+    declared_mime = content_type.split(";", 1)[0]
+    if declared_mime != detected_mime:
+        raise ValueError(f"image MIME mismatch: declared {declared_mime}, detected {detected_mime}")
     digest = hashlib.sha256(body).hexdigest()
-    extension = mimetypes.guess_extension(content_type.split(";", 1)[0]) or Path(urlsplit(source_url).path).suffix
-    extension = extension if extension in {".jpg", ".jpeg", ".png", ".webp", ".gif", ".svg"} else ".bin"
+    extension = mimetypes.guess_extension(detected_mime) or Path(urlsplit(source_url).path).suffix
+    extension = ".jpg" if extension == ".jpe" else extension
     target = root / "files" / digest[:2] / f"{digest}{extension}"
     target.parent.mkdir(parents=True, exist_ok=True)
     if not target.exists():
         target.write_bytes(body)
-    return digest, str(target.relative_to(root))
+    return digest, str(target.relative_to(root)), width, height, detected_mime