"""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 MAX_ORIGINAL_ASSETS = 16 class ImageStudioImageError(RuntimeError): """Raised when a remote image cannot be safely loaded or saved.""" class ImageStudioImageCancelled(ImageStudioImageError): """Raised when an original image download is cancelled safely.""" @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, should_stop=None, ) -> RemoteImage: """Download remote image bytes with SSRF, timeout and size guards.""" url = str(url or "").strip() should_stop = should_stop or (lambda: False) if should_stop(): raise ImageStudioImageCancelled("用户停止下载商品原图") _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 try: if should_stop(): raise ImageStudioImageCancelled("用户停止下载商品原图") 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 should_stop(): raise ImageStudioImageCancelled("用户停止下载商品原图") if not chunk: continue total += len(chunk) if total > int(max_bytes): raise ImageStudioImageError("远程图片超过大小上限") chunks.append(chunk) if should_stop(): raise ImageStudioImageCancelled("用户停止下载商品原图") 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, ) finally: close = getattr(response, "close", None) if callable(close): close() 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 import_original_files( project_id, file_paths, *, path=None, config=None, image_root=None, max_assets=MAX_ORIGINAL_ASSETS, ): """Validate and atomically copy local product images into one studio project.""" imported = [] errors = [] for file_path in file_paths or []: try: with open(os.path.abspath(str(file_path)), "rb") as fh: content = fh.read(int(ORIGINAL_MAX_BYTES) + 1) imported.append( import_original_bytes( project_id, content, filename_hint=os.path.basename(str(file_path)), path=path, config=config, image_root=image_root, max_assets=max_assets, ) ) except Exception as exc: errors.append({"path": os.path.abspath(str(file_path)), "error": str(exc)}) return {"assets": imported, "errors": errors, "limit": int(max_assets)} def import_original_bytes( project_id, content, *, filename_hint="clipboard.png", path=None, config=None, image_root=None, max_assets=MAX_ORIGINAL_ASSETS, ): cfg = appconfig.load_config() if config is None else config database_path = path or appconfig.db_path(cfg) project = image_studio.get_project(project_id, path=database_path) if project is None: raise ImageStudioImageError("商品套图任务不存在") image_bytes = bytes(content or b"") if not image_bytes: raise ImageStudioImageError("商品原图内容为空") if len(image_bytes) > int(ORIGINAL_MAX_BYTES): raise ImageStudioImageError("商品原图超过大小上限") info = _image_info(image_bytes) originals = image_studio.list_assets( project.id, kind=image_studio.ASSET_KIND_ORIGINAL, path=database_path, ) active_originals = [ asset for asset in originals if asset.status != image_studio.ASSET_STATUS_MISSING ] digest = hashlib.sha256(image_bytes).hexdigest()[:16] duplicate = next( ( asset for asset in originals if digest in os.path.basename(str(asset.local_path or "")) and _existing_local_asset(asset) is not None ), None, ) if duplicate is not None: return duplicate if len(active_originals) >= int(max_assets): raise ImageStudioImageError("商品原图最多只能添加%d张" % int(max_assets)) directory = image_studio.project_image_dirs( image_root or appconfig.image_dir(cfg), project, )["originals"] os.makedirs(directory, exist_ok=True) order = max([int(asset.source_order or 0) for asset in originals] + [0]) + 1 extension = _extension_for_format(info["format"]) safe_hint = os.path.splitext(os.path.basename(str(filename_hint or "image")))[0] safe_hint = "".join(ch if ch.isalnum() or ch in "_-" else "_" for ch in safe_hint) safe_hint = safe_hint.strip("_")[:32] or "image" final_path = os.path.join( directory, "original_%02d_%s_%s%s" % (order, safe_hint, digest, extension), ) temp_path = final_path + ".tmp-" + uuid.uuid4().hex created = False try: with open(temp_path, "wb") as fh: fh.write(image_bytes) with open(temp_path, "rb") as fh: _image_info(fh.read()) os.replace(temp_path, final_path) created = True return image_studio.add_asset( project.id, image_studio.ASSET_KIND_ORIGINAL, local_path=final_path, aspect_ratio="%d:%d" % (int(info["width"]), int(info["height"])), source_order=order, path=database_path, ) except Exception as exc: for candidate in (temp_path, final_path if created else None): if candidate and os.path.exists(candidate): try: os.remove(candidate) except OSError: pass if isinstance(exc, ImageStudioImageError): raise raise ImageStudioImageError(f"保存商品原图失败：{exc}") from exc def trash_generated_asset(asset_id, *, path=None, config=None, image_root=None): """Move one generated image to the app-managed trash without losing DB history.""" 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 or not str(asset.kind or "").startswith("generated_"): raise ImageStudioImageError("只能删除商品套图生成结果") source_path = os.path.abspath(str(asset.local_path or "")) if not source_path or not os.path.isfile(source_path): raise ImageStudioImageError("生成图片文件不存在") project = image_studio.get_project(asset.project_id, path=database_path) if project is None: raise ImageStudioImageError("商品套图任务不存在") root = image_studio.project_image_dirs( image_root or appconfig.image_dir(cfg), project, )["root"] trash_dir = os.path.join(root, ".trash") os.makedirs(trash_dir, exist_ok=True) trash_path = os.path.join( trash_dir, "%s_%s" % (uuid.uuid4().hex, os.path.basename(source_path)), ) try: os.replace(source_path, trash_path) image_studio.update_asset_local_path( asset.id, trash_path, status=image_studio.ASSET_STATUS_MISSING, path=database_path, ) except Exception as exc: if os.path.isfile(trash_path) and not os.path.exists(source_path): try: os.replace(trash_path, source_path) except OSError: pass raise ImageStudioImageError(f"删除生成图片失败：{exc}") from exc return { "asset_id": asset.id, "original_path": source_path, "trash_path": trash_path, } def restore_trashed_asset(record, *, path=None, config=None): cfg = appconfig.load_config() if config is None else config database_path = path or appconfig.db_path(cfg) payload = dict(record or {}) asset = image_studio.get_asset(payload.get("asset_id"), path=database_path) if asset is None: raise ImageStudioImageError("待撤销的生成图片记录不存在") trash_path = os.path.abspath(str(payload.get("trash_path") or "")) original_path = os.path.abspath(str(payload.get("original_path") or "")) if not os.path.isfile(trash_path): raise ImageStudioImageError("废纸篓中的生成图片不存在") if os.path.exists(original_path): stem, extension = os.path.splitext(original_path) original_path = "%s_restored_%s%s" % (stem, uuid.uuid4().hex[:8], extension) os.makedirs(os.path.dirname(original_path), exist_ok=True) os.replace(trash_path, original_path) try: return image_studio.update_asset_local_path( asset.id, original_path, status=image_studio.ASSET_STATUS_AVAILABLE, path=database_path, ) except Exception as exc: try: os.replace(original_path, trash_path) except OSError: pass raise ImageStudioImageError(f"撤销删除生成图片失败：{exc}") from exc def download_original_asset( asset_id, *, path=None, config=None, image_root=None, session=None, should_stop=None, ): """Download one Shopee original image into originals/ and mark its asset available.""" cfg = appconfig.load_config() if config is None else config should_stop = should_stop or (lambda: False) if should_stop(): raise ImageStudioImageCancelled("用户停止下载商品原图") 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, should_stop=should_stop, ) if should_stop(): raise ImageStudioImageCancelled("用户停止下载商品原图") 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: if should_stop(): raise ImageStudioImageCancelled("用户停止下载商品原图") with open(temp_path, "wb") as fh: fh.write(remote.content) if should_stop(): raise ImageStudioImageCancelled("用户停止下载商品原图") 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, )