ajyy/backend/app/core/storage.py

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

import json
import mimetypes
import os
import uuid
from datetime import datetime, timezone
from io import BytesIO
from pathlib import Path
from typing import Any, BinaryIO, Dict, Tuple, Union
from urllib.parse import urlparse

from app.core.config import get_settings
from app.core.upload_validate import IMAGE_EXTENSIONS, file_extension

StreamLike = Union[BinaryIO, Any]


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


def _safe_content_type(filename: str) -> str:
    guessed = mimetypes.guess_type(filename)[0] or "application/octet-stream"
    banned = {
        "application/x-msdownload",
        "application/x-msdos-program",
        "application/x-executable",
        "application/x-sh",
    }
    if guessed in banned:
        return "application/octet-stream"
    return guessed


def _key_has_traversal(key: str) -> bool:
    return any(part == ".." for part in key.split("/"))


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)
        os.chmod(dest, 0o644)
        url = f"/uploads/{object_key}"
        return {"url": url, "objectKey": object_key}

    def open_file(self, object_key: str) -> Tuple[StreamLike, str, str, bool]:
        settings = get_settings()
        base_dir = Path(settings.LOCAL_STORAGE_DIR)
        dest = base_dir / object_key
        _ensure_within_base(base_dir, dest)
        if not dest.is_file():
            raise StorageError("文件不存在")
        filename = dest.name
        content_type = _safe_content_type(filename)
        inline = file_extension(filename) in IMAGE_EXTENSIONS
        return dest.open("rb"), content_type, filename, inline


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)
        # 仅二维码公示图匿名可读；反馈/建言附件私有
        self._ensure_public_read_policy()

    def _ensure_public_read_policy(self) -> None:
        policy = {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Principal": {"AWS": ["*"]},
                    "Action": ["s3:GetObject"],
                    "Resource": [f"arn:aws:s3:::{self.bucket_name}/qrcode/*"],
                }
            ],
        }
        self.client.set_bucket_policy(self.bucket_name, json.dumps(policy))

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

    def open_file(self, object_key: str) -> Tuple[StreamLike, str, str, bool]:
        filename = object_key.rsplit("/", 1)[-1] or "file"
        content_type = _safe_content_type(filename)
        inline = file_extension(filename) in IMAGE_EXTENSIONS
        try:
            resp = self.client.get_object(self.bucket_name, object_key)
        except Exception as exc:
            raise StorageError("文件不存在或无法读取") from exc
        return resp, content_type, filename, inline


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)


def resolve_object_key(url: str) -> str:
    """从库内 url 解析 objectKey；非法则 StorageError。"""
    raw = (url or "").strip()
    if not raw:
        raise StorageError("无效的文件路径")
    if "?" in raw:
        raw = raw.split("?", 1)[0]

    if raw.startswith("/uploads/"):
        key = raw[len("/uploads/") :].lstrip("/")
        if not key or _key_has_traversal(key):
            raise StorageError("非法的文件夹路径")
        return key

    if raw.startswith("http://") or raw.startswith("https://"):
        settings = get_settings()
        parsed = urlparse(raw)
        path = (parsed.path or "").lstrip("/")
        bucket = settings.MINIO_BUCKET
        marker = f"{bucket}/"
        idx = path.find(marker)
        if idx < 0:
            raise StorageError("无效的文件路径")
        key = path[idx + len(marker) :].lstrip("/")
        if not key or _key_has_traversal(key):
            raise StorageError("非法的文件夹路径")

        # 校验 host：须匹配 MINIO_BASE_URL 或 MINIO_ENDPOINT
        allowed_hosts: set[str] = set()
        if settings.MINIO_BASE_URL:
            base = urlparse(
                settings.MINIO_BASE_URL
                if "://" in settings.MINIO_BASE_URL
                else f"http://{settings.MINIO_BASE_URL}"
            )
            if base.hostname:
                allowed_hosts.add(base.hostname.lower())
        endpoint_host = settings.MINIO_ENDPOINT.split(":")[0].lower()
        if endpoint_host:
            allowed_hosts.add(endpoint_host)
        host = (parsed.hostname or "").lower()
        if host not in allowed_hosts:
            raise StorageError("无效的文件路径")
        return key

    # 允许直接传 objectKey（无前缀）
    key = raw.lstrip("/")
    if not key or _key_has_traversal(key) or key.startswith(".."):
        raise StorageError("非法的文件夹路径")
    if "://" in key:
        raise StorageError("无效的文件路径")
    return key


def open_file(url: str) -> Tuple[StreamLike, str, str, bool]:
    """返回 (stream, content_type, download_name, inline)。"""
    key = resolve_object_key(url)
    return get_storage().open_file(key)