discord_bot/fjerkroa_bot/images.py

"""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)