HeXi/hexi/plugins/nonebot_plugin_video_analysis/storage/s3.py

"""统一 S3 存储模块 — 合并本地局域网 S3 和公网 MinIO"""

from pathlib import Path
from time import strftime, localtime

import boto3
from botocore.config import Config
from nonebot import logger

from hexi.web_config import get_effective_value

# 该插件模块名，用于读取统一配置值库(Web 修改后生效)
_PLUGIN_ID = "hexi.plugins.nonebot_plugin_video_analysis"


def _ev(name: str, default):
    """读取统一配置生效值(用户值库 > getter > 默认)，来源无关。"""
    return get_effective_value(_PLUGIN_ID, name, default)


# ============================= 本地局域网 PLANA_S3 配置 =============================
LOCAL_S3_PLANA_ENDPOINT = "http://192.168.2.15:5246"
LOCAL_S3_PLANA_ACCESS_KEY = "vS8Lf7UoiS5cxL+mz+CH"
LOCAL_S3_PLANA_SECRET_KEY = "Fi3phPW10CsZWuiQCerK4gxaIIJp41w5SQNlnDRb"
LOCAL_S3_PLANA_BUCKET = "PLANA"

# ============================= 本地局域网 PLANB_S3 配置 =============================
LOCAL_S3_PLANB_ENDPOINT = "http://192.168.2.15:5246"
LOCAL_S3_PLANB_ACCESS_KEY = "vS8Lf7UoiS5cxL+mz+CH"
LOCAL_S3_PLANB_SECRET_KEY = "Fi3phPW10CsZWuiQCerK4gxaIIJp41w5SQNlnDRb"
LOCAL_S3_PLANB_BUCKET = "PLANB"

# ============================= 本地局域网 S3 配置 =============================
LOCAL_S3_PLANC_ENDPOINT = "http://192.168.2.15:5246"
LOCAL_S3_PLANC_ACCESS_KEY = "vS8Lf7UoiS5cxL+mz+CH"
LOCAL_S3_PLANC_SECRET_KEY = "Fi3phPW10CsZWuiQCerK4gxaIIJp41w5SQNlnDRb"
LOCAL_S3_PLANC_BUCKET = "PLANC"

# ============================= 公网 MinIO 配置 =============================
PUBLIC_S3_ENDPOINT = "s3.sansenhoshi.top"
PUBLIC_S3_ACCESS_KEY = "JDMynACSjPaN8JRwriwS"
PUBLIC_S3_SECRET_KEY = "JRl4bIGxeiqwqvfeBTlAWQbUKMpNNHSZPm6ne93j"
PUBLIC_S3_BUCKET = "s-file-trans"
PUBLIC_S3_REGION = "bot"
PUBLIC_S3_SECURE = True
PUBLIC_S3_DOMAIN = f"https://{PUBLIC_S3_ENDPOINT}/{PUBLIC_S3_BUCKET}"

# ============================= 客户端实例（懒加载） =============================
_local_s3_plana = None
_local_s3_planb = None
_local_s3_planc = None
_public_s3 = None


class S3Client:
    """通用 S3 客户端"""

    def __init__(
        self,
        endpoint_url: str,
        access_key: str,
        secret_key: str,
        bucket: str,
        region: str = "us-east-1",
        secure: bool = False,
    ):
        self.bucket = bucket
        self.client = boto3.client(
            "s3",
            endpoint_url=endpoint_url,
            aws_access_key_id=access_key,
            aws_secret_access_key=secret_key,
            region_name=region,
            verify=secure,
            config=Config(
                signature_version="s3v4",
                s3={"addressing_style": "path"},
                retries={"max_attempts": 3},
                connect_timeout=10,
                read_timeout=60,
            ),
        )

    def upload_and_get_url(
        self, local_path: str, key: str, expire: int = 3600
    ) -> str:
        self.client.upload_file(local_path, self.bucket, key)
        url = self.client.generate_presigned_url(
            "get_object",
            Params={"Bucket": self.bucket, "Key": key},
            ExpiresIn=expire,
        )
        return url

    def upload_public(self, local_path: str, key: str) -> str:
        """上传并设为公开读"""
        self.client.upload_file(
            local_path,
            self.bucket,
            key,
            ExtraArgs={"ACL": "public-read"},
        )
        logger.info(f"文件 {key} 已成功上传至云端S3")
        return f"{_ev('PUBLIC_S3_DOMAIN', PUBLIC_S3_DOMAIN)}/{key}"

    def delete(self, key: str) -> bool:
        try:
            self.client.delete_object(Bucket=self.bucket, Key=key)
            logger.info(f"成功删除 S3 文件: {key}")
            return True
        except Exception:
            logger.exception(f"删除 S3 文件失败: {key}")
            return False


def _get_local_s3_plana() -> S3Client:
    global _local_s3_plana
    if _local_s3_plana is None:
        _local_s3_plana = S3Client(
            _ev("LOCAL_S3_PLANA_ENDPOINT", LOCAL_S3_PLANA_ENDPOINT),
            _ev("LOCAL_S3_PLANA_ACCESS_KEY", LOCAL_S3_PLANA_ACCESS_KEY),
            _ev("LOCAL_S3_PLANA_SECRET_KEY", LOCAL_S3_PLANA_SECRET_KEY),
            _ev("LOCAL_S3_PLANA_BUCKET", LOCAL_S3_PLANA_BUCKET),
        )
    return _local_s3_plana


def _get_local_s3_planb() -> S3Client:
    global _local_s3_planb
    if _local_s3_planb is None:
        _local_s3_planb = S3Client(
            _ev("LOCAL_S3_PLANB_ENDPOINT", LOCAL_S3_PLANB_ENDPOINT),
            _ev("LOCAL_S3_PLANB_ACCESS_KEY", LOCAL_S3_PLANB_ACCESS_KEY),
            _ev("LOCAL_S3_PLANB_SECRET_KEY", LOCAL_S3_PLANB_SECRET_KEY),
            _ev("LOCAL_S3_PLANB_BUCKET", LOCAL_S3_PLANB_BUCKET),
        )
    return _local_s3_planb


def _get_local_s3_planc() -> S3Client:
    global _local_s3_planc
    if _local_s3_planc is None:
        _local_s3_planc = S3Client(
            _ev("LOCAL_S3_PLANC_ENDPOINT", LOCAL_S3_PLANC_ENDPOINT),
            _ev("LOCAL_S3_PLANC_ACCESS_KEY", LOCAL_S3_PLANC_ACCESS_KEY),
            _ev("LOCAL_S3_PLANC_SECRET_KEY", LOCAL_S3_PLANC_SECRET_KEY),
            _ev("LOCAL_S3_PLANC_BUCKET", LOCAL_S3_PLANC_BUCKET),
        )
    return _local_s3_planc


def _get_public_s3() -> S3Client:
    global _public_s3
    if _public_s3 is None:
        secure = _ev("PUBLIC_S3_SECURE", PUBLIC_S3_SECURE)
        endpoint = _ev("PUBLIC_S3_ENDPOINT", PUBLIC_S3_ENDPOINT)
        _public_s3 = S3Client(
            f"{'https' if secure else 'http'}://{endpoint}",
            _ev("PUBLIC_S3_ACCESS_KEY", PUBLIC_S3_ACCESS_KEY),
            _ev("PUBLIC_S3_SECRET_KEY", PUBLIC_S3_SECRET_KEY),
            _ev("PUBLIC_S3_BUCKET", PUBLIC_S3_BUCKET),
            region=_ev("PUBLIC_S3_REGION", PUBLIC_S3_REGION),
            secure=secure,
        )
    return _public_s3


# ============================= 对外接口 =============================

def upload_to_public_s3(file_path: str | Path) -> tuple[str, str]:
    """
    上传到公网 MinIO，返回 (公网URL, file_key)

    私聊场景使用，生成可公网访问的链接
    """
    file_path = Path(file_path)
    current_day = strftime("%Y-%m-%d", localtime())
    file_key = f"{current_day}/{file_path.name}"
    try:
        url = _get_public_s3().upload_public(str(file_path), file_key)
        return url, file_key
    except Exception:
        logger.exception("上传到公网 MinIO 失败")
        return "", ""


def upload_to_local_s3(
    title: str, image_post: bool, file_path: str | Path, plan: str | None = None
) -> str:
    """
    上传到局域网 S3，返回预签名 URL

    plan="A" → PLANA 桶
    plan="B" → PLANB 桶
    None    → PLANC 桶（默认）
    """
    file_path = Path(file_path)
    current_day = strftime("%Y-%m-%d", localtime())
    if image_post:
        file_key = f"{current_day}/{title}/{file_path.name}"
    else:
        file_key = f"{current_day}/{file_path.name}"

    if plan == "A":
        client = _get_local_s3_plana()
    elif plan == "B":
        client = _get_local_s3_planb()
    else:
        client = _get_local_s3_planc()

    try:
        logger.info(f"文件 {file_key} 已成功上传至本地S3")
        return client.upload_and_get_url(str(file_path), file_key)
    except Exception:
        logger.exception("上传到本地 S3 失败")
        return ""


def delete_from_public_s3(file_key: str) -> bool:
    """从公网 MinIO 删除指定 key 的对象"""
    return _get_public_s3().delete(file_key)


def upload_with_plan(
    file_path: str | Path,
    *,
    plan: str | None = None,
    is_private: bool = False,
    title: str = "",
    image_post: bool = False,
) -> tuple[str, str | None]:
    """
    统一上传入口：按 plan 路由到对应的本地桶，并按需上传公网

    plan="A"  → PLANA（仅本地）
    plan="B"  → PLANB + 公网
    私聊       → PLANB + 公网
    默认       → PLANC（仅本地）

    Returns:
        (local_url, public_url_or_none)
    """
    # 本地上传
    if plan == "A":
        local_url = upload_to_local_s3(title, image_post, file_path, plan="A")
    elif plan == "B":
        local_url = upload_to_local_s3(title, image_post, file_path, plan="B")
    elif is_private:
        local_url = upload_to_local_s3(title, image_post, file_path, plan="B")
    else:
        local_url = upload_to_local_s3(title, image_post, file_path)  # PLANC

    # 公网上传（仅 PLANB 或私聊）
    public_url = None
    if plan == "B" or is_private:
        public_url, _ = upload_to_public_s3(file_path)

    return local_url, public_url