diff --git a/app/image_studio.py b/app/image_studio.py index 2bf34b2..b2968c1 100644 --- a/app/image_studio.py +++ b/app/image_studio.py @@ -553,6 +553,23 @@ def mark_asset_status(asset_id, status, path=None, conn=None): return get_asset(asset_id, conn=database) +def update_asset_local_path(asset_id, local_path, status=ASSET_STATUS_AVAILABLE, path=None, conn=None): + if str(status) not in ASSET_STATUSES: + raise db.DbError("AI工场资产状态必须是 available 或 missing") + abs_local_path = os.path.abspath(local_path) if local_path else None + with _connection(conn, path) as database: + with database: + database.execute( + """ + UPDATE image_studio_assets + SET local_path = ?, status = ?, updated_at = ? + WHERE id = ? + """, + (abs_local_path, str(status), _now(), int(asset_id)), + ) + return get_asset(asset_id, conn=database) + + def create_job( project_id, *, diff --git a/app/image_studio_images.py b/app/image_studio_images.py new file mode 100644 index 0000000..4eddd07 --- /dev/null +++ b/app/image_studio_images.py @@ -0,0 +1,355 @@ +"""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, + ) diff --git a/docs/tasks/T-588.md b/docs/tasks/T-588.md index aee29ba..c9252e2 100644 --- a/docs/tasks/T-588.md +++ b/docs/tasks/T-588.md @@ -2,8 +2,8 @@ id: T-588 title: AI工场远程缩略图与蝦皮原图进入照片池 phase: 7 -deps: [T-587, T-546] -status: TODO +deps: [T-587, T-548] +status: DONE created: 2026-07-11 --- @@ -35,4 +35,10 @@ created: 2026-07-11 ## 执行记录 -（完成后记录 CDN 实探、缓存和下载验证。） +- 2026-07-11：完成 AI工场远程图片底层能力。 + - 修正任务依赖：原 `T-546` 任务文件不存在，实际复用已完成的 T-548 图片下载安全/超时口径，frontmatter 改为 `deps: [T-587, T-548]`。 + - 新增 `app/image_studio_images.py`：提供远程图片安全下载（http/https、公网地址、无系统代理、连接/读取超时、大小上限、Content-Type 与最终重定向 URL 校验、PIL 解码校验）、内存缩略图生成、线程池 `ThumbnailLoader`（缓存、取消、失败回调）和单张原图下载落 `originals/`。 + - `download_original_asset()` 只在用户单击时下载对应原图，写临时文件、解码校验后 `os.replace()` 原子落盘，并通过 `image_studio.update_asset_local_path()` 回写同一 remote asset；重复单击已有可解码本地文件时直接返回，不重复请求。 + - 新增 `tests/test_image_studio_images.py` 覆盖缩略图不落盘、缓存命中、取消、原图落盘幂等、坏图片不留下半文件/可用 asset、非图片、超限、内网 URL 和重定向到内网 URL。 + - 验证：在只套用 T-588 diff 的 clean worktree 中运行 `python -m ruff check app tests main.py`、`py -3.10 -m compileall app main.py`、`py -3.10 -m unittest discover -s tests`（355 tests）和 `git diff --check`，全部通过。 + - CDN 真实实探未执行：本轮没有指定可用的已登录测试账号/商品现场；代码已记录可观测字段 `content_type/final_url/redirected` 并覆盖回退与边界，后续接入 T-591 GUI 前建议用测试账号补一条真实 Shopee CDN 记录。 diff --git a/tests/test_image_studio_images.py b/tests/test_image_studio_images.py new file mode 100644 index 0000000..2e5d5ac --- /dev/null +++ b/tests/test_image_studio_images.py @@ -0,0 +1,251 @@ +import io +import os +import socket +import sys +import threading +import time +import unittest +from types import SimpleNamespace +from unittest import mock + +sys.path.insert(0, os.path.dirname(__file__)) + +from _helpers import TempDirMixin + +from app import db, image_studio, image_studio_images + + +class _Response: + def __init__( + self, + content=b"", + status_code=200, + headers=None, + url="https://cdn.example.com/image.png", + chunk_size=None, + ): + self.content = content + self.status_code = status_code + self.headers = dict(headers or {}) + self.url = url + self.chunk_size = chunk_size + + def iter_content(self, chunk_size=65536): + if self.chunk_size: + for index in range(0, len(self.content), self.chunk_size): + yield self.content[index:index + self.chunk_size] + return + yield self.content + + +class ImageStudioImageTests(TempDirMixin, unittest.TestCase): + def _png_bytes(self, size=(20, 16), color=(80, 120, 200, 255)): + from PIL import Image + + output = io.BytesIO() + Image.new("RGBA", size, color).save(output, format="PNG") + return output.getvalue() + + def _public_dns(self): + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))] + + def _project_and_asset(self, temp_dir, remote_url="https://cdn.example.com/source.png"): + db_path = os.path.join(temp_dir, "cmshopee.db") + db.init_db(db_path) + project = image_studio.create_or_get_project( + account_alias="alias", + account_slug="alias_slug", + item_id="51100639510", + path=db_path, + ) + asset = image_studio.sync_original_asset_urls( + project.id, + [{"index": 1, "src": remote_url}], + path=db_path, + )[0] + return db_path, project, asset + + def test_download_remote_image_rejects_private_non_image_overlimit_and_private_redirect(self): + with self.assertRaisesRegex(image_studio_images.ImageStudioImageError, "http/https|本机或内网"): + image_studio_images.download_remote_image("http://127.0.0.1/a.png") + + non_image_session = SimpleNamespace( + get=lambda *args, **kwargs: _Response( + content=b"hello", + headers={"Content-Type": "text/html"}, + ) + ) + with mock.patch("app.image_studio_images.socket.getaddrinfo", return_value=self._public_dns()): + with self.assertRaisesRegex(image_studio_images.ImageStudioImageError, "不是图片"): + image_studio_images.download_remote_image( + "https://cdn.example.com/a.png", + session=non_image_session, + ) + + oversized_session = SimpleNamespace( + get=lambda *args, **kwargs: _Response( + content=b"x", + headers={"Content-Type": "image/png", "Content-Length": "999"}, + ) + ) + with mock.patch("app.image_studio_images.socket.getaddrinfo", return_value=self._public_dns()): + with self.assertRaisesRegex(image_studio_images.ImageStudioImageError, "大小上限"): + image_studio_images.download_remote_image( + "https://cdn.example.com/a.png", + session=oversized_session, + max_bytes=10, + ) + + redirect_session = SimpleNamespace( + get=lambda *args, **kwargs: _Response( + content=self._png_bytes(), + headers={"Content-Type": "image/png"}, + url="http://127.0.0.1/private.png", + ) + ) + with mock.patch("app.image_studio_images.socket.getaddrinfo", return_value=self._public_dns()): + with self.assertRaisesRegex(image_studio_images.ImageStudioImageError, "本机或内网"): + image_studio_images.download_remote_image( + "https://cdn.example.com/a.png", + session=redirect_session, + ) + + def test_load_thumbnail_decodes_image_without_writing_file(self): + png = self._png_bytes(size=(640, 480)) + session = SimpleNamespace( + get=lambda *args, **kwargs: _Response( + content=png, + headers={"Content-Type": "image/png"}, + ) + ) + + with mock.patch("app.image_studio_images.socket.getaddrinfo", return_value=self._public_dns()): + result = image_studio_images.load_thumbnail( + "https://cdn.example.com/a.png", + key="thumb-1", + max_size=120, + session=session, + ) + + self.assertEqual("thumb-1", result.key) + self.assertLessEqual(result.width, 120) + self.assertLessEqual(result.height, 120) + self.assertTrue(result.image_bytes.startswith(b"\x89PNG")) + + def test_thumbnail_loader_uses_memory_cache_and_can_cancel(self): + calls = [] + successes = [] + errors = [] + + def fake_loader(url, key=None): + calls.append((key, url)) + return image_studio_images.ThumbnailResult( + key=str(key), + url=url, + image_bytes=b"png", + width=10, + height=10, + ) + + loader = image_studio_images.ThumbnailLoader(max_workers=1, loader=fake_loader) + try: + future = loader.submit("a", "https://cdn.example.com/a.png", successes.append, errors.append) + future.result(timeout=2) + for _ in range(20): + if successes: + break + time.sleep(0.01) + loader.submit("b", "https://cdn.example.com/a.png", successes.append, errors.append) + + self.assertEqual([("a", "https://cdn.example.com/a.png")], calls) + self.assertEqual(2, len(successes)) + self.assertFalse(successes[0].from_cache) + self.assertTrue(successes[1].from_cache) + self.assertEqual([], errors) + + release = threading.Event() + + def slow_loader(url, key=None): + release.wait(1) + return image_studio_images.ThumbnailResult( + key=str(key), + url=url, + image_bytes=b"png", + width=10, + height=10, + ) + + loader._loader = slow_loader + future = loader.submit("c", "https://cdn.example.com/c.png", successes.append, errors.append) + loader.cancel("c") + release.set() + future.result(timeout=2) + time.sleep(0.05) + self.assertEqual(2, len(successes)) + finally: + loader.close() + + def test_download_original_asset_writes_atomically_and_is_idempotent(self): + with self.make_temp_dir() as temp_dir: + db_path, _, asset = self._project_and_asset(temp_dir) + png = self._png_bytes() + session = mock.Mock() + session.get.return_value = _Response( + content=png, + headers={"Content-Type": "image/png"}, + ) + + with mock.patch("app.image_studio_images.socket.getaddrinfo", return_value=self._public_dns()): + downloaded = image_studio_images.download_original_asset( + asset.id, + path=db_path, + config={"data_dir": temp_dir, "image_dir": os.path.join(temp_dir, "images")}, + session=session, + ) + second = image_studio_images.download_original_asset( + asset.id, + path=db_path, + config={"data_dir": temp_dir, "image_dir": os.path.join(temp_dir, "images")}, + session=mock.Mock(), + ) + + self.assertEqual(downloaded.id, second.id) + self.assertEqual(downloaded.local_path, second.local_path) + self.assertTrue(os.path.isfile(downloaded.local_path)) + self.assertTrue(downloaded.local_path.endswith(".png")) + self.assertEqual(image_studio.ASSET_STATUS_AVAILABLE, downloaded.status) + session.get.assert_called_once() + originals = os.path.dirname(downloaded.local_path) + self.assertFalse([name for name in os.listdir(originals) if ".tmp-" in name]) + + self.assert_removed(temp_dir) + + def test_download_original_asset_bad_image_leaves_asset_without_local_path(self): + with self.make_temp_dir() as temp_dir: + db_path, project, asset = self._project_and_asset(temp_dir) + session = SimpleNamespace( + get=lambda *args, **kwargs: _Response( + content=b"not an image", + headers={"Content-Type": "image/png"}, + ) + ) + + with mock.patch("app.image_studio_images.socket.getaddrinfo", return_value=self._public_dns()): + with self.assertRaisesRegex(image_studio_images.ImageStudioImageError, "解码失败"): + image_studio_images.download_original_asset( + asset.id, + path=db_path, + config={"data_dir": temp_dir, "image_dir": os.path.join(temp_dir, "images")}, + session=session, + ) + + refreshed = image_studio.get_asset(asset.id, path=db_path) + self.assertIsNone(refreshed.local_path) + dirs = image_studio.project_image_dirs(os.path.join(temp_dir, "images"), project) + self.assertFalse(os.path.exists(dirs["originals"])) + + self.assert_removed(temp_dir) + + +if __name__ == "__main__": + unittest.main()