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()