cmshoppe/app/image_studio_images.py

"""Remote image helpers for the AI image studio."""

from __future__ import annotations

import hashlib
import io
import ipaddress
import os
import socket
import threading
import urllib.parse
import uuid
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from typing import Callable, Optional

import requests

from . import appconfig, image_studio


THUMBNAIL_MAX_BYTES = 3 * 1024 * 1024
ORIGINAL_MAX_BYTES = 12 * 1024 * 1024
DEFAULT_CONNECT_TIMEOUT_SECONDS = 5
DEFAULT_READ_TIMEOUT_SECONDS = 30
DEFAULT_THUMBNAIL_SIZE = 220
DEFAULT_THUMBNAIL_WORKERS = 4


class ImageStudioImageError(RuntimeError):
    """Raised when a remote image cannot be safely loaded or saved."""


@dataclass(frozen=True)
class RemoteImage:
    url: str
    content: bytes
    content_type: str
    final_url: str
    redirected: bool


@dataclass(frozen=True)
class ThumbnailResult:
    key: str
    url: str
    image_bytes: bytes
    width: int
    height: int
    from_cache: bool = False


def _session():
    session = requests.Session()
    session.trust_env = False
    return session


def _assert_public_http_url(url):
    parts = urllib.parse.urlsplit(str(url or ""))
    if parts.scheme not in {"http", "https"}:
        raise ImageStudioImageError("远程图片地址只允许 http/https")
    host = parts.hostname
    if not host:
        raise ImageStudioImageError("远程图片地址缺少域名")
    if _is_local_hostname(host):
        raise ImageStudioImageError("远程图片地址不能指向本机或内网")
    try:
        _assert_public_ip(host)
        return
    except ValueError:
        pass
    try:
        addresses = socket.getaddrinfo(
            host,
            parts.port or (443 if parts.scheme == "https" else 80),
            type=socket.SOCK_STREAM,
        )
    except OSError as exc:
        raise ImageStudioImageError(f"远程图片地址无法解析：{exc}") from exc
    if not addresses:
        raise ImageStudioImageError("远程图片地址无法解析")
    for address in addresses:
        _assert_public_ip(address[4][0])


def _is_local_hostname(host):
    lowered = str(host or "").strip().lower().rstrip(".")
    return lowered in {"localhost"} or lowered.endswith(".localhost") or lowered.endswith(".local")


def _assert_public_ip(value):
    ip = ipaddress.ip_address(value)
    if (
        ip.is_private
        or ip.is_loopback
        or ip.is_link_local
        or ip.is_multicast
        or ip.is_reserved
        or ip.is_unspecified
    ):
        raise ImageStudioImageError("远程图片地址不能指向本机或内网")


def _safe_headers(referer=None):
    headers = {
        "User-Agent": "Mozilla/5.0",
        "Accept": "image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8",
    }
    if referer:
        headers["Referer"] = str(referer)
    return headers


def _validate_content_type(value):
    content_type = str(value or "").split(";", 1)[0].strip().lower()
    if content_type and not content_type.startswith("image/"):
        raise ImageStudioImageError("远程地址返回的不是图片内容")
    return content_type


def download_remote_image(
    url,
    *,
    referer=None,
    max_bytes=ORIGINAL_MAX_BYTES,
    timeout=(DEFAULT_CONNECT_TIMEOUT_SECONDS, DEFAULT_READ_TIMEOUT_SECONDS),
    session=None,
) -> RemoteImage:
    """Download remote image bytes with SSRF, timeout and size guards."""

    url = str(url or "").strip()
    _assert_public_http_url(url)
    client = session or _session()
    try:
        response = client.get(
            url,
            stream=True,
            timeout=timeout,
            headers=_safe_headers(referer),
        )
    except requests.exceptions.RequestException as exc:
        raise ImageStudioImageError(f"远程图片下载失败：{exc}") from exc
    status = getattr(response, "status_code", 200)
    if status >= 400:
        raise ImageStudioImageError(f"远程图片下载失败：HTTP {status}")
    final_url = str(getattr(response, "url", "") or url)
    _assert_public_http_url(final_url)
    content_type = _validate_content_type(getattr(response, "headers", {}).get("Content-Type", ""))
    content_length = getattr(response, "headers", {}).get("Content-Length")
    if content_length:
        try:
            if int(content_length) > int(max_bytes):
                raise ImageStudioImageError("远程图片超过大小上限")
        except ValueError:
            pass
    chunks = []
    total = 0
    iterator = response.iter_content(chunk_size=65536) if hasattr(response, "iter_content") else [response.content]
    for chunk in iterator:
        if not chunk:
            continue
        total += len(chunk)
        if total > int(max_bytes):
            raise ImageStudioImageError("远程图片超过大小上限")
        chunks.append(chunk)
    content = b"".join(chunks)
    _image_info(content)
    return RemoteImage(
        url=url,
        content=content,
        content_type=content_type,
        final_url=final_url,
        redirected=final_url != url,
    )


def _image_info(image_bytes):
    try:
        from PIL import Image

        with Image.open(io.BytesIO(image_bytes)) as image:
            image.verify()
        with Image.open(io.BytesIO(image_bytes)) as image:
            return {
                "format": str(image.format or "JPEG").upper(),
                "width": int(image.width),
                "height": int(image.height),
            }
    except Exception as exc:
        raise ImageStudioImageError(f"图片解码失败：{exc}") from exc


def _thumbnail_bytes(image_bytes, max_size=DEFAULT_THUMBNAIL_SIZE):
    try:
        from PIL import Image

        with Image.open(io.BytesIO(image_bytes)) as image:
            image.thumbnail((int(max_size), int(max_size)))
            output = io.BytesIO()
            image.convert("RGBA").save(output, format="PNG")
            return output.getvalue(), image.width, image.height
    except Exception as exc:
        raise ImageStudioImageError(f"缩略图解码失败：{exc}") from exc


def load_thumbnail(url, *, key=None, max_size=DEFAULT_THUMBNAIL_SIZE, referer=None, session=None):
    remote = download_remote_image(
        url,
        referer=referer,
        max_bytes=THUMBNAIL_MAX_BYTES,
        session=session,
    )
    image_bytes, width, height = _thumbnail_bytes(remote.content, max_size=max_size)
    return ThumbnailResult(
        key=str(key or url),
        url=remote.url,
        image_bytes=image_bytes,
        width=width,
        height=height,
        from_cache=False,
    )


class ThumbnailLoader:
    """Small thread-pool thumbnail loader with memory cache and cooperative cancel."""

    def __init__(self, max_workers=DEFAULT_THUMBNAIL_WORKERS, loader: Optional[Callable] = None):
        self._executor = ThreadPoolExecutor(max_workers=max(1, int(max_workers or 1)))
        self._loader = loader or load_thumbnail
        self._cache = {}
        self._cancelled = set()
        self._closed = False
        self._lock = threading.Lock()

    def submit(self, key, url, on_success=None, on_error=None):
        key = str(key)
        url = str(url or "").strip()
        with self._lock:
            if self._closed:
                raise ImageStudioImageError("缩略图加载器已关闭")
            cached = self._cache.get(url)
            if cached is not None:
                result = ThumbnailResult(key, url, cached.image_bytes, cached.width, cached.height, from_cache=True)
                if on_success is not None:
                    on_success(result)
                return None
            self._cancelled.discard(key)
            future = self._executor.submit(self._loader, url, key=key)

        def done_callback(done):
            try:
                result = done.result()
                with self._lock:
                    if key in self._cancelled or self._closed:
                        return
                    self._cache[url] = result
                if on_success is not None:
                    on_success(result)
            except Exception as exc:
                with self._lock:
                    if key in self._cancelled or self._closed:
                        return
                if on_error is not None:
                    on_error(key, exc)

        future.add_done_callback(done_callback)
        return future

    def cancel(self, key=None):
        with self._lock:
            if key is None:
                self._cancelled.add("*")
                self._closed = True
            else:
                self._cancelled.add(str(key))

    def close(self):
        with self._lock:
            self._closed = True
        self._executor.shutdown(wait=False, cancel_futures=True)


def _extension_for_format(format_name):
    normalized = str(format_name or "").strip().upper()
    if normalized in {"JPEG", "JPG"}:
        return ".jpg"
    if normalized == "PNG":
        return ".png"
    if normalized == "WEBP":
        return ".webp"
    return ".img"


def _original_file_path(project, asset, image_root):
    dirs = image_studio.project_image_dirs(image_root, project)
    digest = hashlib.sha1(str(asset.remote_url or asset.id).encode("utf-8")).hexdigest()[:10]
    stem = "original_%02d_%s" % (max(1, int(asset.source_order or 1)), digest)
    return dirs["originals"], stem


def _existing_local_asset(asset):
    local_path = str(asset.local_path or "").strip()
    if not local_path or not os.path.isfile(local_path):
        return None
    try:
        with open(local_path, "rb") as fh:
            _image_info(fh.read())
        return asset
    except ImageStudioImageError:
        return None


def download_original_asset(asset_id, *, path=None, config=None, image_root=None, session=None):
    """Download one Shopee original image into originals/ and mark its asset available."""

    cfg = appconfig.load_config() if config is None else config
    database_path = path or appconfig.db_path(cfg)
    asset = image_studio.get_asset(asset_id, path=database_path)
    if asset is None:
        raise ImageStudioImageError("原图资产不存在")
    if not asset.remote_url:
        raise ImageStudioImageError("原图资产缺少远程 URL")
    existing = _existing_local_asset(asset)
    if existing is not None:
        return existing
    project = image_studio.get_project(asset.project_id, path=database_path)
    if project is None:
        raise ImageStudioImageError("原图资产所属项目不存在")
    remote = download_remote_image(asset.remote_url, max_bytes=ORIGINAL_MAX_BYTES, session=session)
    info = _image_info(remote.content)
    directory, stem = _original_file_path(project, asset, image_root or appconfig.image_dir(cfg))
    os.makedirs(directory, exist_ok=True)
    final_path = os.path.join(directory, stem + _extension_for_format(info["format"]))
    temp_path = final_path + ".tmp-" + uuid.uuid4().hex
    try:
        with open(temp_path, "wb") as fh:
            fh.write(remote.content)
        with open(temp_path, "rb") as fh:
            _image_info(fh.read())
        os.replace(temp_path, final_path)
    except Exception as exc:
        try:
            if os.path.exists(temp_path):
                os.remove(temp_path)
        finally:
            if isinstance(exc, ImageStudioImageError):
                raise
            raise ImageStudioImageError(f"保存原图失败：{exc}") from exc
    return image_studio.update_asset_local_path(
        asset.id,
        final_path,
        status=image_studio.ASSET_STATUS_AVAILABLE,
        path=database_path,
    )