ajyy/backend/app/core/storage.py

"""文件存储：local | minio。upload_file 返回 {url, objectKey}。"""
from __future__ import annotations

import os
import uuid
from datetime import datetime, timezone
from io import BytesIO
from pathlib import Path
from typing import Any, Dict

from app.core.config import get_settings


class StorageError(Exception):
    """存储参数/写入错误（如路径穿越）。"""

    def __init__(self, message: str):
        self.message = message
        super().__init__(message)


def _unique_name(filename: str) -> str:
    ext = os.path.splitext(filename)[1] or ""
    return f"{uuid.uuid4().hex}{ext}"


def _normalize_folder(folder: str) -> str:
    """规范化 folder：拒绝空段以外的 '..' / 绝对路径痕迹。"""
    raw = (folder or "").strip().replace("\\", "/")
    if not raw:
        return ""
    # 拒绝绝对路径或盘符式前缀
    if raw.startswith("/") or (len(raw) >= 2 and raw[1] == ":"):
        raise StorageError("非法的文件夹路径")
    parts: list[str] = []
    for part in raw.split("/"):
        if part in ("", "."):
            continue
        if part == "..":
            raise StorageError("非法的文件夹路径")
        parts.append(part)
    return "/".join(parts)


def _object_key(folder: str, filename: str) -> str:
    unique = _unique_name(filename)
    date_part = datetime.now(timezone.utc).strftime("%Y%m/%d")
    folder = _normalize_folder(folder)
    if folder:
        return f"{folder}/{date_part}/{unique}"
    return f"{date_part}/{unique}"


def _ensure_within_base(base_dir: Path, dest: Path) -> None:
    base_resolved = base_dir.resolve()
    dest_resolved = dest.resolve()
    try:
        dest_resolved.relative_to(base_resolved)
    except ValueError as exc:
        raise StorageError("非法的文件夹路径") from exc


class LocalStorage:
    def upload_file(self, content: bytes, filename: str, folder: str = "") -> Dict[str, str]:
        settings = get_settings()
        base_dir = Path(settings.LOCAL_STORAGE_DIR)
        object_key = _object_key(folder, filename)
        dest = base_dir / object_key
        _ensure_within_base(base_dir, dest)
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_bytes(content)
        url = f"/uploads/{object_key}"
        return {"url": url, "objectKey": object_key}


class MinIOStorage:
    def __init__(self) -> None:
        from minio import Minio

        settings = get_settings()
        self._settings = settings
        self.client = Minio(
            settings.MINIO_ENDPOINT,
            access_key=settings.MINIO_ACCESS_KEY,
            secret_key=settings.MINIO_SECRET_KEY,
            secure=settings.MINIO_SECURE,
        )
        self.bucket_name = settings.MINIO_BUCKET
        if settings.MINIO_BASE_URL:
            self.base_url = settings.MINIO_BASE_URL.rstrip("/")
        else:
            scheme = "https" if settings.MINIO_SECURE else "http"
            self.base_url = f"{scheme}://{settings.MINIO_ENDPOINT}/{self.bucket_name}"
        self._ensure_bucket()

    def _ensure_bucket(self) -> None:
        if not self.client.bucket_exists(self.bucket_name):
            self.client.make_bucket(self.bucket_name)

    def upload_file(self, content: bytes, filename: str, folder: str = "") -> Dict[str, str]:
        object_key = _object_key(folder, filename)
        self.client.put_object(
            bucket_name=self.bucket_name,
            object_name=object_key,
            data=BytesIO(content),
            length=len(content),
        )
        url = f"{self.base_url}/{object_key}"
        return {"url": url, "objectKey": object_key}


def get_storage() -> Any:
    settings = get_settings()
    if settings.STORAGE_TYPE.lower() == "minio":
        return MinIOStorage()
    return LocalStorage()


def upload_file(content: bytes, filename: str, folder: str = "") -> Dict[str, str]:
    """上传文件，返回 {url, objectKey}。"""
    return get_storage().upload_file(content, filename, folder)