FabledScribe/src/fabledassistant/services/images.py

"""Image cache service: fetch from web, store on disk, serve locally.

Images are stored in IMAGE_CACHE_DIR (default /data/images) and tracked in
the image_cache DB table.  The original URL is preserved for citation.
Deduplication is by SHA-256 of the original URL — the same image requested
twice shares one file and one DB row.
"""

import hashlib
import ipaddress
import logging
import socket
from pathlib import Path
from urllib.parse import urlparse

import httpx
from sqlalchemy import select
from sqlalchemy.exc import IntegrityError

from fabledassistant.config import Config
from fabledassistant.models import async_session
from fabledassistant.models.image_cache import ImageCache

logger = logging.getLogger(__name__)


def _is_safe_image_url(url: str) -> bool:
    """Return True only if the URL is a public http/https address (SSRF guard)."""
    try:
        parsed = urlparse(url)
        if parsed.scheme not in ("http", "https"):
            return False
        host = parsed.hostname
        if not host:
            return False
        if host.lower() in ("localhost", "::1"):
            return False
        addr_info = socket.getaddrinfo(host, None, proto=socket.IPPROTO_TCP)
        for entry in addr_info:
            ip = ipaddress.ip_address(entry[4][0])
            if ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved or ip.is_multicast:
                return False
        return True
    except Exception:
        return False  # Block on resolution failure

_ALLOWED_TYPES = {
    "image/jpeg",
    "image/png",
    "image/gif",
    "image/webp",
    "image/avif",
}

_EXT_MAP = {
    "image/jpeg": "jpg",
    "image/png":  "png",
    "image/gif":  "gif",
    "image/webp": "webp",
    "image/avif": "avif",
}


def _cache_dir() -> Path:
    d = Path(Config.IMAGE_CACHE_DIR)
    d.mkdir(parents=True, exist_ok=True)
    return d


def _url_hash(url: str) -> str:
    return hashlib.sha256(url.encode()).hexdigest()


async def fetch_and_store_image(
    url: str,
    title: str | None = None,
    source_domain: str | None = None,
) -> ImageCache | None:
    """Fetch an image URL, write it to disk, persist metadata to DB.

    Returns the ImageCache record (new or existing).
    Returns None if the URL is unreachable, not a valid image type, or
    exceeds IMAGE_MAX_BYTES.
    """
    if not _is_safe_image_url(url):
        logger.warning("Blocked image fetch of private/internal URL: %s", url[:80])
        return None

    url_hash = _url_hash(url)

    # Return existing record if already cached
    async with async_session() as session:
        existing = (
            await session.execute(
                select(ImageCache).where(ImageCache.url_hash == url_hash)
            )
        ).scalar_one_or_none()
        if existing:
            logger.debug("Image cache hit id=%d for %s", existing.id, url[:80])
            return existing

    # Infer source domain from URL if not supplied
    if not source_domain:
        try:
            source_domain = urlparse(url).netloc or None
        except Exception:
            source_domain = None

    # Fetch from origin
    try:
        async with httpx.AsyncClient(timeout=15.0, follow_redirects=True) as client:
            headers = {
                "User-Agent": "Mozilla/5.0 (compatible; FabledAssistant/1.0)",
            }
            # Provide a same-origin Referer to bypass basic hotlink protection
            if source_domain:
                headers["Referer"] = f"https://{source_domain}/"
            resp = await client.get(url, headers=headers)
            resp.raise_for_status()

            ct = resp.headers.get("content-type", "").split(";")[0].strip().lower()
            if ct not in _ALLOWED_TYPES:
                logger.warning("Skipping non-image content-type '%s' for %s", ct, url[:80])
                return None

            # Reject oversized images before buffering all bytes
            cl = resp.headers.get("content-length")
            if cl and int(cl) > Config.IMAGE_MAX_BYTES:
                logger.warning("Image too large (%s bytes, limit %d) at %s", cl, Config.IMAGE_MAX_BYTES, url[:80])
                return None

            data = resp.content
            if len(data) > Config.IMAGE_MAX_BYTES:
                logger.warning("Image too large (%d bytes) at %s", len(data), url[:80])
                return None
            if not data:
                return None

    except Exception:
        logger.warning("Failed to fetch image from %s", url[:80], exc_info=True)
        return None

    # Write to disk
    file_ext = _EXT_MAP.get(ct, "jpg")
    file_path = _cache_dir() / f"{url_hash}.{file_ext}"
    try:
        file_path.write_bytes(data)
    except Exception:
        logger.warning("Failed to write image to disk: %s", file_path, exc_info=True)
        return None

    # Persist metadata — handle the rare race-condition duplicate gracefully
    record = ImageCache(
        url_hash=url_hash,
        original_url=url,
        source_domain=source_domain,
        title=title,
        content_type=ct,
        file_size=len(data),
        file_ext=file_ext,
    )
    async with async_session() as session:
        try:
            session.add(record)
            await session.commit()
            await session.refresh(record)
        except IntegrityError:
            await session.rollback()
            existing = (
                await session.execute(
                    select(ImageCache).where(ImageCache.url_hash == url_hash)
                )
            ).scalar_one_or_none()
            if existing:
                return existing
            logger.warning("Failed to persist image cache record for %s", url[:80], exc_info=True)
            return None

    logger.info(
        "Cached image id=%d (%s, %d bytes) from %s",
        record.id, ct, len(data), url[:80],
    )
    return record


async def get_image_record(image_id: int) -> ImageCache | None:
    """Return the ImageCache row for a given ID, or None."""
    async with async_session() as session:
        return (
            await session.execute(
                select(ImageCache).where(ImageCache.id == image_id)
            )
        ).scalar_one_or_none()


def get_image_path(record: ImageCache) -> Path:
    """Filesystem path for a cached image record."""
    return _cache_dir() / f"{record.url_hash}.{record.file_ext}"