image input pipeline: content-hash cache, data-url vision, picture_edit real

fjerkroa_bot/images.py

@@ -0,0 +1,124 @@
+"""Content-hash image cache (SPEC-004, FDB-010).
+
+Attachments are downloaded once, sniffed, stored under their sha256
+and served to vision as data: URLs — Discord's expiring CDN links
+never travel further (IMG-10/11). LRU + TTL keep the cache bounded
+(IMG-12); deletions and !forgetme propagate here (IMG-14).
+"""
+
+import base64
+import hashlib
+import logging
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional
+
+import aiohttp
+
+from .persistence import PersistentStore
+
+DEFAULT_CACHE_MB = 500
+DEFAULT_TTL_DAYS = 90
+DEFAULT_MAX_BYTES = 8 * 1024 * 1024
+DOWNLOAD_TIMEOUT_S = 20
+
+MAGIC = [
+    (b"\x89PNG", "png"),
+    (b"\xff\xd8\xff", "jpg"),
+    (b"GIF87a", "gif"),
+    (b"GIF89a", "gif"),
+]
+
+
+def sniff_ext(data: bytes) -> Optional[str]:
+    """Extension from magic bytes only — names and headers lie (IMG-10)."""
+    for magic, ext in MAGIC:
+        if data.startswith(magic):
+            return ext
+    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
+        return "webp"
+    return None
+
+
+class ImageCache:
+    def __init__(self, store: PersistentStore, root: Path, config_getter: Callable[[], Dict[str, Any]]) -> None:
+        self.store = store
+        self.root = Path(root)
+        self._config = config_getter
+        self.root.mkdir(parents=True, exist_ok=True)
+
+    def _path(self, sha256: str, ext: str) -> Path:
+        return self.root / f"{sha256}.{ext}"
+
+    def ingest_bytes(self, data: bytes, channel: str, user: str, message_id: Optional[str]) -> Optional[str]:
+        ext = sniff_ext(data)
+        if ext is None:
+            logging.warning(f"image cache: rejected non-image bytes from {user} (IMG-10)")
+            return None
+        if len(data) > int(self._config().get("image-max-bytes", DEFAULT_MAX_BYTES)):
+            logging.warning(f"image cache: rejected oversized upload from {user} ({len(data)} bytes)")
+            return None
+        sha256 = hashlib.sha256(data).hexdigest()
+        path = self._path(sha256, ext)
+        if not path.exists():
+            path.write_bytes(data)
+        self.store.image_add(sha256, channel, user, message_id, ext, len(data))
+        self.evict()
+        return sha256
+
+    async def ingest_url(self, url: str, channel: str, user: str, message_id: Optional[str]) -> Optional[str]:
+        try:
+            data = await self._download(url)
+        except Exception as err:
+            logging.warning(f"image cache: download failed for {user}: {repr(err)}")
+            return None
+        return self.ingest_bytes(data, channel, user, message_id)
+
+    async def _download(self, url: str) -> bytes:
+        limit = int(self._config().get("image-max-bytes", DEFAULT_MAX_BYTES))
+        timeout = aiohttp.ClientTimeout(total=DOWNLOAD_TIMEOUT_S)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            async with session.get(url) as response:
+                response.raise_for_status()
+                return await response.content.read(limit + 1)
+
+    def data_url(self, sha256: str, ext: str) -> Optional[str]:
+        path = self._path(sha256, ext)
+        if not path.exists():
+            return None
+        mime = "jpeg" if ext == "jpg" else ext
+        return f"data:image/{mime};base64," + base64.b64encode(path.read_bytes()).decode()
+
+    def recent(self, channel: str, count: int) -> List[Dict[str, Any]]:
+        return self.store.images_recent(channel, count)
+
+    def recent_paths(self, channel: str, count: int) -> List[Path]:
+        paths = [self._path(row["sha256"], row["ext"]) for row in self.recent(channel, count)]
+        return [path for path in paths if path.exists()]
+
+    def _remove(self, sha256: str, ext: str) -> None:
+        self._path(sha256, ext).unlink(missing_ok=True)
+        self.store.images_delete(sha256)
+
+    def evict(self) -> None:
+        """TTL first, then LRU down to the byte cap (IMG-12)."""
+        config = self._config()
+        for row in self.store.images_expired(int(config.get("image-cache-ttl-days", DEFAULT_TTL_DAYS))):
+            self._remove(row["sha256"], row["ext"])
+        cap = int(config.get("image-cache-mb", DEFAULT_CACHE_MB)) * 1024 * 1024
+        while self.store.images_total_bytes() > cap:
+            victims = self.store.images_oldest(1)
+            if not victims:
+                break
+            self._remove(victims[0]["sha256"], victims[0]["ext"])
+
+    def purge_user(self, user: str) -> int:
+        rows = self.store.images_for_user(user)
+        for row in rows:
+            self._remove(row["sha256"], row["ext"])
+        return len(rows)
+
+    def purge_message(self, message_id: str) -> int:
+        rows = self.store.images_for_message(message_id)
+        for row in rows:
+            self._remove(row["sha256"], row["ext"])
+        return len(rows)