dsh-plugin-image-workbench 0.0.0-stage → 1.1.6

package/vendor/chatgpt2api/source/services/image_storage_service.py

@@ -0,0 +1,405 @@
+from __future__ import annotations
+
+import hashlib
+import io
+import json
+import time
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from threading import Lock
+from urllib.parse import quote, urlparse
+
+from curl_cffi import requests
+from fastapi import HTTPException
+from PIL import Image
+
+from services.config import DATA_DIR, config
+
+IMAGE_INDEX_FILE = DATA_DIR / "image_index.json"
+IMAGE_INDEX_LOCK = Lock()
+IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
+
+
+class ImageStorageError(RuntimeError):
+    pass
+
+
+@dataclass(frozen=True)
+class StoredImage:
+    rel: str
+    url: str
+    storage: str
+    size: int
+
+
+def _clean(value: object) -> str:
+    return str(value or "").strip()
+
+
+def _now_iso() -> str:
+    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+
+def _safe_relative_path(path: str) -> str:
+    value = str(path or "").strip().replace("\\", "/").lstrip("/")
+    if not value:
+        raise HTTPException(status_code=404, detail="image not found")
+    parts = Path(value).parts
+    if any(part in {"", ".", ".."} for part in parts):
+        raise HTTPException(status_code=404, detail="image not found")
+    return Path(*parts).as_posix()
+
+
+def _image_dimensions(payload: bytes) -> tuple[int, int] | None:
+    try:
+        with Image.open(io.BytesIO(payload)) as image:
+            return image.size
+    except Exception:
+        return None
+
+
+def _is_image_rel(path: str) -> bool:
+    try:
+        safe_rel = _safe_relative_path(path)
+    except HTTPException:
+        return False
+    return Path(safe_rel).suffix.lower() in IMAGE_EXTENSIONS
+
+
+def _local_image_path(relative_path: str) -> Path:
+    rel = _safe_relative_path(relative_path)
+    root = config.images_dir.resolve()
+    path = (root / rel).resolve()
+    try:
+        path.relative_to(root)
+    except ValueError as exc:
+        raise HTTPException(status_code=404, detail="image not found") from exc
+    return path
+
+
+def _read_json_object(path: Path) -> dict[str, object]:
+    if not path.exists():
+        return {}
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+    except Exception:
+        return {}
+    return data if isinstance(data, dict) else {}
+
+
+def _write_json_object(path: Path, data: dict[str, object]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp_path = path.with_suffix(path.suffix + ".tmp")
+    tmp_path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+    tmp_path.replace(path)
+
+
+class WebDAVClient:
+    def __init__(self, settings: dict[str, object]):
+        self.url = _clean(settings.get("webdav_url")).rstrip("/")
+        self.username = _clean(settings.get("webdav_username"))
+        self.password = _clean(settings.get("webdav_password"))
+        self.root_path = _clean(settings.get("webdav_root_path")).strip("/")
+        self.session = requests.Session()
+
+    def _auth_kwargs(self) -> dict[str, object]:
+        return {"auth": (self.username, self.password)} if self.username or self.password else {}
+
+    def _request(self, method: str, url: str, **kwargs):
+        response = self.session.request(method, url, timeout=30, **self._auth_kwargs(), **kwargs)
+        if response.status_code >= 400 and not (method == "MKCOL" and response.status_code in {405}):
+            raise ImageStorageError(f"WebDAV {method} failed: HTTP {response.status_code}")
+        return response
+
+    def remote_url(self, rel: str = "") -> str:
+        parts = [part for part in [self.root_path, _safe_relative_path(rel) if rel else ""] if part]
+        encoded = "/".join(quote(part, safe="") for item in parts for part in item.split("/") if part)
+        return f"{self.url}/{encoded}" if encoded else self.url
+
+    def ensure_dirs(self, rel: str) -> None:
+        parts = [part for part in [self.root_path, Path(_safe_relative_path(rel)).parent.as_posix()] if part and part != "."]
+        current = self.url
+        for item in "/".join(parts).split("/"):
+            if not item:
+                continue
+            current = f"{current}/{quote(item, safe='')}"
+            response = self.session.request("MKCOL", current, timeout=30, **self._auth_kwargs())
+            if response.status_code in {201, 405}:
+                continue
+            if response.status_code >= 400:
+                raise ImageStorageError(f"WebDAV MKCOL failed: HTTP {response.status_code}")
+
+    def put(self, rel: str, payload: bytes, content_type: str = "image/png") -> str:
+        self.ensure_dirs(rel)
+        url = self.remote_url(rel)
+        self._request("PUT", url, data=payload, headers={"Content-Type": content_type})
+        return url
+
+    def get(self, rel: str) -> bytes:
+        response = self._request("GET", self.remote_url(rel))
+        return bytes(response.content)
+
+    def delete(self, rel: str) -> bool:
+        response = self.session.request("DELETE", self.remote_url(rel), timeout=30, **self._auth_kwargs())
+        if response.status_code in {200, 202, 204, 404}:
+            return response.status_code != 404
+        raise ImageStorageError(f"WebDAV DELETE failed: HTTP {response.status_code}")
+
+    def test(self) -> dict[str, object]:
+        if not self.url:
+            return {"ok": False, "status": 0, "error": "WebDAV URL is required"}
+        if urlparse(self.url).scheme not in {"http", "https"}:
+            return {"ok": False, "status": 0, "error": "invalid WebDAV URL"}
+        test_rel = ".chatgpt2api_webdav_test.txt"
+        try:
+            self.put(test_rel, b"chatgpt2api webdav test\n", content_type="text/plain")
+            self.delete(test_rel)
+            return {"ok": True, "status": 200, "error": None}
+        except ImageStorageError as exc:
+            return {"ok": False, "status": 0, "error": str(exc)}
+        except Exception as exc:
+            return {"ok": False, "status": 0, "error": str(exc) or exc.__class__.__name__}
+        finally:
+            self.session.close()
+
+
+class ImageStorageService:
+    def __init__(self, index_file: Path = IMAGE_INDEX_FILE):
+        self.index_file = index_file
+        self._index_lock = IMAGE_INDEX_LOCK
+
+    def settings(self) -> dict[str, object]:
+        return config.get_image_storage_settings()
+
+    def mode(self) -> str:
+        return _clean(self.settings().get("mode")) or "local"
+
+    def _load_index(self) -> dict[str, dict[str, object]]:
+        raw = _read_json_object(self.index_file)
+        items = raw.get("items")
+        if not isinstance(items, dict):
+            return {}
+        return {str(key): value for key, value in items.items() if isinstance(value, dict)}
+
+    def _load_clean_index(self) -> dict[str, dict[str, object]]:
+        items = self._load_index()
+        return {rel: item for rel, item in items.items() if _is_image_rel(rel)}
+
+    def _save_index(self, items: dict[str, dict[str, object]]) -> None:
+        _write_json_object(self.index_file, {"items": items})
+
+    def _public_url(self, rel: str, base_url: str | None = None) -> str:
+        settings = self.settings()
+        public_base_url = _clean(settings.get("public_base_url"))
+        if public_base_url:
+            return f"{public_base_url.rstrip('/')}/{_safe_relative_path(rel)}"
+        return f"{(base_url or config.base_url).rstrip('/')}/images/{_safe_relative_path(rel)}"
+
+    def make_relative_path(self, image_data: bytes) -> str:
+        file_hash = hashlib.md5(image_data).hexdigest()
+        filename = f"{int(time.time())}_{file_hash}.png"
+        relative_dir = Path(time.strftime("%Y"), time.strftime("%m"), time.strftime("%d"))
+        return f"{relative_dir.as_posix()}/{filename}"
+
+    def save(self, image_data: bytes, base_url: str | None = None) -> StoredImage:
+        config.cleanup_old_images()
+        rel = self.make_relative_path(image_data)
+        mode = self.mode()
+        if mode not in {"local", "webdav", "both"}:
+            mode = "local"
+        stored_local = False
+        stored_webdav = False
+        remote_url = ""
+
+        if mode in {"local", "both"}:
+            path = _local_image_path(rel)
+            path.parent.mkdir(parents=True, exist_ok=True)
+            path.write_bytes(image_data)
+            stored_local = True
+
+        if mode in {"webdav", "both"}:
+            remote_url = WebDAVClient(self.settings()).put(rel, image_data)
+            stored_webdav = True
+
+        dimensions = _image_dimensions(image_data)
+        item = {
+            "rel": rel,
+            "path": rel,
+            "name": Path(rel).name,
+            "date": "-".join(rel.split("/")[:3]),
+            "size": len(image_data),
+            "created_at": _now_iso(),
+            "storage": "both" if stored_local and stored_webdav else ("webdav" if stored_webdav else "local"),
+            "local": stored_local,
+            "webdav": stored_webdav,
+            "remote_url": remote_url,
+        }
+        if dimensions:
+            item["width"], item["height"] = dimensions
+        with self._index_lock:
+            items = self._load_clean_index()
+            items[rel] = item
+            self._save_index(items)
+        return StoredImage(rel=rel, url=self._public_url(rel, base_url), storage=str(item["storage"]), size=len(image_data))
+
+    def get_bytes(self, rel: str) -> bytes:
+        safe_rel = _safe_relative_path(rel)
+        if not _is_image_rel(safe_rel):
+            raise HTTPException(status_code=404, detail="image not found")
+        path = _local_image_path(safe_rel)
+        if path.is_file():
+            return path.read_bytes()
+        item = self._load_clean_index().get(safe_rel, {})
+        if item.get("webdav"):
+            return WebDAVClient(self.settings()).get(safe_rel)
+        raise HTTPException(status_code=404, detail="image not found")
+
+    def exists(self, rel: str) -> bool:
+        safe_rel = _safe_relative_path(rel)
+        if not _is_image_rel(safe_rel):
+            return False
+        if _local_image_path(safe_rel).is_file():
+            return True
+        item = self._load_clean_index().get(safe_rel, {})
+        return bool(item.get("webdav"))
+
+    def has_local(self, rel: str) -> bool:
+        safe_rel = _safe_relative_path(rel)
+        return _is_image_rel(safe_rel) and _local_image_path(safe_rel).is_file()
+
+    def list_items(self, base_url: str, start_date: str = "", end_date: str = "") -> list[dict[str, object]]:
+        with self._index_lock:
+            indexed = self._load_clean_index()
+            root = config.images_dir
+            changed = False
+            for path in root.rglob("*"):
+                if not path.is_file() or not _is_image_rel(path.name):
+                    continue
+                rel = path.relative_to(root).as_posix()
+                if rel in indexed:
+                    continue
+                dimensions = None
+                try:
+                    dimensions = _image_dimensions(path.read_bytes())
+                except Exception:
+                    dimensions = None
+                indexed[rel] = {
+                    "rel": rel,
+                    "path": rel,
+                    "name": path.name,
+                    "date": "-".join(rel.split("/")[:3]) if len(rel.split("/")) >= 4 else datetime.fromtimestamp(path.stat().st_mtime).strftime("%Y-%m-%d"),
+                    "size": path.stat().st_size,
+                    "created_at": datetime.fromtimestamp(path.stat().st_mtime).strftime("%Y-%m-%d %H:%M:%S"),
+                    "storage": "local",
+                    "local": True,
+                    "webdav": False,
+                    **({"width": dimensions[0], "height": dimensions[1]} if dimensions else {}),
+                }
+                changed = True
+
+            items: list[dict[str, object]] = []
+            for rel, item in list(indexed.items()):
+                if not _is_image_rel(rel):
+                    indexed.pop(rel, None)
+                    changed = True
+                    continue
+                local = _local_image_path(rel).is_file()
+                webdav = bool(item.get("webdav"))
+                if not local and not webdav:
+                    indexed.pop(rel, None)
+                    changed = True
+                    continue
+                storage = "both" if local and webdav else ("webdav" if webdav else "local")
+                if item.get("local") != local or item.get("storage") != storage:
+                    item = {
+                        **item,
+                        "local": local,
+                        "storage": storage,
+                    }
+                    indexed[rel] = item
+                    changed = True
+                day = str(item.get("date") or "")
+                if start_date and day < start_date:
+                    continue
+                if end_date and day > end_date:
+                    continue
+                items.append({
+                    **item,
+                    "rel": rel,
+                    "path": rel,
+                    "url": self._public_url(rel, base_url),
+                })
+            if changed:
+                self._save_index(indexed)
+        items.sort(key=lambda item: str(item.get("created_at") or ""), reverse=True)
+        return items
+
+    def delete(self, rel: str) -> bool:
+        safe_rel = _safe_relative_path(rel)
+        removed = False
+        path = _local_image_path(safe_rel)
+        if path.is_file():
+            path.unlink()
+            removed = True
+        with self._index_lock:
+            items = self._load_clean_index()
+            item = items.get(safe_rel, {})
+            if item.get("webdav"):
+                try:
+                    removed = WebDAVClient(self.settings()).delete(safe_rel) or removed
+                except ImageStorageError:
+                    if not removed:
+                        raise
+            if safe_rel in items:
+                items.pop(safe_rel, None)
+                self._save_index(items)
+        return removed
+
+    def sync_all(self) -> dict[str, int]:
+        settings = self.settings()
+        if self.mode() not in {"webdav", "both"}:
+            raise ImageStorageError("WebDAV 图片存储未启用")
+        uploaded = 0
+        skipped = 0
+        failed = 0
+        with self._index_lock:
+            items = self._load_clean_index()
+            client = WebDAVClient(settings)
+            for path in sorted(config.images_dir.rglob("*")):
+                if not path.is_file() or not _is_image_rel(path.name):
+                    continue
+                rel = path.relative_to(config.images_dir).as_posix()
+                item = items.get(rel, {})
+                if item.get("webdav"):
+                    skipped += 1
+                    continue
+                try:
+                    payload = path.read_bytes()
+                    remote_url = client.put(rel, payload)
+                    dimensions = _image_dimensions(payload)
+                    items[rel] = {
+                        **item,
+                        "rel": rel,
+                        "path": rel,
+                        "name": path.name,
+                        "date": "-".join(rel.split("/")[:3]) if len(rel.split("/")) >= 4 else datetime.fromtimestamp(path.stat().st_mtime).strftime("%Y-%m-%d"),
+                        "size": len(payload),
+                        "created_at": str(item.get("created_at") or datetime.fromtimestamp(path.stat().st_mtime).strftime("%Y-%m-%d %H:%M:%S")),
+                        "storage": "both",
+                        "local": True,
+                        "webdav": True,
+                        "remote_url": remote_url,
+                        **({"width": dimensions[0], "height": dimensions[1]} if dimensions else {}),
+                    }
+                    uploaded += 1
+                except Exception:
+                    failed += 1
+            self._save_index(items)
+        return {"uploaded": uploaded, "skipped": skipped, "failed": failed}
+
+    def test_webdav(self) -> dict[str, object]:
+        return WebDAVClient(self.settings()).test()
+
+
+image_storage_service = ImageStorageService()

package/vendor/chatgpt2api/source/services/image_tags_service.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from services.config import DATA_DIR
+
+TAGS_FILE = DATA_DIR / "image_tags.json"
+
+
+def _ensure_file() -> None:
+    TAGS_FILE.parent.mkdir(parents=True, exist_ok=True)
+    if not TAGS_FILE.exists():
+        TAGS_FILE.write_text("{}", encoding="utf-8")
+
+
+def load_tags() -> dict[str, list[str]]:
+    _ensure_file()
+    try:
+        data = json.loads(TAGS_FILE.read_text(encoding="utf-8"))
+    except Exception:
+        return {}
+    return data if isinstance(data, dict) else {}
+
+
+def save_tags(data: dict[str, list[str]]) -> None:
+    _ensure_file()
+    TAGS_FILE.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+
+
+def get_tags(image_rel: str) -> list[str]:
+    return load_tags().get(image_rel, [])
+
+
+def set_tags(image_rel: str, tags: list[str]) -> list[str]:
+    data = load_tags()
+    cleaned = list(dict.fromkeys(t.strip() for t in tags if t.strip()))
+    if cleaned:
+        data[image_rel] = cleaned
+    else:
+        data.pop(image_rel, None)
+    save_tags(data)
+    return cleaned
+
+
+def remove_tags(image_rel: str) -> None:
+    data = load_tags()
+    if data.pop(image_rel, None) is not None:
+        save_tags(data)
+
+
+def delete_tag(tag: str) -> int:
+    """从所有图片中删除指定标签，返回受影响的图片数。"""
+    data = load_tags()
+    count = 0
+    for rel in list(data):
+        if tag in data[rel]:
+            data[rel] = [t for t in data[rel] if t != tag]
+            if not data[rel]:
+                del data[rel]
+            count += 1
+    if count > 0:
+        save_tags(data)
+    return count
+
+
+def get_all_tags() -> list[str]:
+    data = load_tags()
+    seen: set[str] = set()
+    result: list[str] = []
+    for tags in data.values():
+        for t in tags:
+            if t not in seen:
+                seen.add(t)
+                result.append(t)
+    return result