"""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 .httpread import read_capped 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() # limit + 1: an over-limit body must stay over-limit so # ingest_bytes rejects it instead of caching it truncated return await read_capped(response, 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)