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