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

rf4_research/media_cli.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 import argparse
+from datetime import datetime, timezone
 import json
 from pathlib import Path
+import urllib.error
 import urllib.request
 
 from .community_cli import USER_AGENT, _StrictRedirectHandler, _validate_url_before_io, check_and_reserve, fetch_html, fetch_site_key
@@ -22,17 +24,25 @@ def _download_one(root: Path, state_file: Path) -> str:
     url = queued["asset_url"]
     _validate_url_before_io(url)
     check_and_reserve(fetch_site_key(url), state_file=state_file)
-    request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT, "Accept": "image/*"})
-    opener = urllib.request.build_opener(_StrictRedirectHandler())
-    with opener.open(request, timeout=30) as response:
-        content_type = response.headers.get_content_type()
-        if not content_type.startswith("image/"):
-            raise ValueError(f"expected image, got {content_type}")
-        body = response.read(MAX_ASSET_BYTES + 1)
-        if len(body) > MAX_ASSET_BYTES:
-            raise ValueError("asset exceeded 15MB limit")
-    digest, relative = store_asset(root, body, content_type=content_type, source_url=url)
-    queued.update({"status": "stored", "sha256": digest, "local_path": relative, "content_type": content_type, "bytes": len(body)})
+    attempted_at = datetime.now(timezone.utc).isoformat()
+    try:
+        request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT, "Accept": "image/*"})
+        opener = urllib.request.build_opener(_StrictRedirectHandler())
+        with opener.open(request, timeout=30) as response:
+            content_type = response.headers.get_content_type()
+            if not content_type.startswith("image/"):
+                raise ValueError(f"expected image, got {content_type}")
+            body = response.read(MAX_ASSET_BYTES + 1)
+            if len(body) > MAX_ASSET_BYTES:
+                raise ValueError("asset exceeded 15MB limit")
+        digest, relative, width, height, detected_mime = store_asset(root, body, content_type=content_type, source_url=url)
+    except Exception as exc:
+        code = exc.code if isinstance(exc, urllib.error.HTTPError) else None
+        status = "missing" if code in {404, 410} else "blocked" if code in {401, 403} else "invalid" if isinstance(exc, ValueError) else "queued"
+        queued.update({"status": status, "last_attempt_at": attempted_at, "last_error": str(exc)[:500]})
+        manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+        raise
+    queued.update({"status": "stored", "sha256": digest, "local_path": relative, "content_type": detected_mime, "bytes": len(body), "width": width, "height": height, "fetched_at": attempted_at})
     manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
     return f"stored {url} as {relative}"