LoveACE-EndF/loveace/service/remote/s3/__init__.py

import os
from contextlib import asynccontextmanager
from typing import Any, AsyncGenerator, BinaryIO, Dict, Optional

import aioboto3
from botocore.client import Config as BotoCoreConfig
from types_aiobotocore_s3 import S3Client

from loveace.config.logger import logger
from loveace.config.manager import config_manager
from loveace.service.model.service import Service
from loveace.service.remote.s3.model.s3 import (
    S3CopyResult,
    S3ListResult,
    S3Object,
    S3UploadResult,
)

s3_config = config_manager.get_settings().s3

# Boto3 很诡异的问题，不把这两个参数设为 when_required 他会把 check 直接塞到 rawfile 里
# 阅读了一下应该是国内的一些 S3 兼容服务不能识读 checksum 导致的

os.environ["AWS_REQUEST_CHECKSUM_CALCULATION"] = "when_required"
os.environ["AWS_RESPONSE_CHECKSUM_VALIDATION"] = "when_required"

# 验证 S3 配置
if not all(
    [
        s3_config.endpoint_url,
        s3_config.access_key_id,
        s3_config.secret_access_key,
        s3_config.bucket_name,
    ]
):
    logger.warning("S3 配置不完整，S3 功能将不可用")
    raise ValueError("S3 配置不完整，S3 功能将不可用")


class S3Service(Service):
    """类型提示完善的 aioboto3 S3 管理器"""

    def __init__(self):
        self._session: aioboto3.Session = aioboto3.Session()
        self._bucket_name = s3_config.bucket_name
        self._endpoint_url = s3_config.endpoint_url
        self._client_config = {
            "aws_access_key_id": s3_config.access_key_id,
            "aws_secret_access_key": s3_config.secret_access_key,
            "endpoint_url": s3_config.endpoint_url,
            "region_name": s3_config.region_name,
            "use_ssl": s3_config.use_ssl,
            "config": BotoCoreConfig(
                s3={
                    "addressing_style": s3_config.addressing_style,
                    "signature_version": s3_config.signature_version,
                }
            ),
        }

    @asynccontextmanager
    async def get_client(self) -> AsyncGenerator[S3Client, None]:
        """获取 S3 客户端上下文管理器"""
        async with self._session.client("s3", **self._client_config) as client:  # type: ignore
            yield client

    def _get_object_url(self, s3_key: str, bucket: Optional[str] = None) -> str:
        """
        生成对象的直链 URL（非预签名）

        Args:
            s3_key: S3 对象键
            bucket: 存储桶名称

        Returns:
            str: 直链 URL
        """
        bucket_name = bucket or self._bucket_name
        # 根据寻址风格构建 URL
        if s3_config.addressing_style == "virtual":
            # 虚拟主机风格：https://bucket-name.endpoint/key
            return f"https://{bucket_name}.{self._endpoint_url.replace('https://', '').replace('http://', '')}/{s3_key}"
        else:
            # 路径风格：https://endpoint/bucket-name/key
            return f"{self._endpoint_url}/{bucket_name}/{s3_key}"

    async def upload_obj(
        self,
        file_obj: BinaryIO,
        s3_key: str,
        bucket: Optional[str] = None,
        extra_args: Optional[Dict[str, Any]] = None,
    ) -> S3UploadResult:
        """
        上传文件对象到 S3

        Args:
            file_obj: 文件对象
            s3_key: S3 对象键
            bucket: 存储桶名称
            extra_args: 额外参数

        Returns:
            S3UploadResult: 上传结果，包含成功状态和直链 URL
        """
        bucket_name = bucket or self._bucket_name

        try:
            async with self.get_client() as s3:
                logger.info(f"开始上传文件对象到 S3: {s3_key}")
                await s3.upload_fileobj(
                    file_obj, bucket_name, s3_key, ExtraArgs=extra_args
                )
                logger.info(f"文件对象上传成功: {s3_key}")
                obj_url = self._get_object_url(s3_key, bucket_name)
                return S3UploadResult(
                    success=True,
                    url=obj_url,
                    key=s3_key,
                )
        except Exception as e:
            logger.error(f"文件对象上传失败 -> {s3_key}: {e}")
            return S3UploadResult(
                success=False,
                key=s3_key,
                error=str(e),
            )

    async def delete_object(self, s3_key: str, bucket: Optional[str] = None) -> bool:
        """
        删除单个 S3 对象

        Args:
            s3_key: S3 对象键
            bucket: 存储桶名称

        Returns:
            bool: 删除成功返回 True
        """
        bucket_name = bucket or self._bucket_name

        try:
            async with self.get_client() as s3:
                await s3.delete_object(Bucket=bucket_name, Key=s3_key)
                logger.info(f"对象删除成功: {s3_key}")
                return True
        except Exception as e:
            logger.error(f"对象删除失败 {s3_key}: {e}")
            return False

    async def list_objects(
        self,
        prefix: str = "",
        bucket: Optional[str] = None,
        max_keys: int = 1000,
        continuation_token: Optional[str] = None,
    ) -> S3ListResult:
        """
        列出 S3 对象

        Args:
            prefix: 对象键前缀
            bucket: 存储桶名称
            max_keys: 最大返回数量
            continuation_token: 继续令牌，用于分页

        Returns:
            S3ListResult: 对象列表结果
        """
        bucket_name = bucket or self._bucket_name

        try:
            async with self.get_client() as s3:
                params: Dict[str, Any] = {
                    "Bucket": bucket_name,
                    "Prefix": prefix,
                    "MaxKeys": max_keys,
                }

                if continuation_token:
                    params["ContinuationToken"] = continuation_token

                response = await s3.list_objects_v2(**params)

                objects = []
                if contents := response.get("Contents"):
                    for item in contents:
                        if key := item.get("Key"):
                            size = item.get("Size", 0)
                            last_mod = item.get("LastModified")
                            last_modified_str = last_mod.isoformat() if last_mod else ""
                            objects.append(
                                S3Object(
                                    key=key,
                                    size=size or 0,
                                    last_modified=last_modified_str,
                                )
                            )

                return S3ListResult(
                    success=True,
                    objects=objects,
                    prefix=prefix,
                    is_truncated=response.get("IsTruncated", False),
                    continuation_token=response.get("NextContinuationToken"),
                )

        except Exception as e:
            logger.error(f"列出对象失败，前缀: {prefix}: {e}")
            return S3ListResult(
                success=False,
                prefix=prefix,
                error=str(e),
            )

    async def generate_presigned_url(
        self,
        s3_key: str,
        bucket: Optional[str] = None,
        expiration: int = 3600,
        method: str = "get_object",
    ) -> Optional[str]:
        """
        生成预签名 URL

        Args:
            s3_key: S3 对象键
            bucket: 存储桶名称
            expiration: URL 有效期（秒）
            method: HTTP 方法（get_object, put_object 等）

        Returns:
            Optional[str]: 预签名 URL，生成失败返回 None
        """
        bucket_name = bucket or self._bucket_name

        try:
            async with self.get_client() as s3:
                url = await s3.generate_presigned_url(
                    ClientMethod=method,
                    Params={"Bucket": bucket_name, "Key": s3_key},
                    ExpiresIn=expiration,
                )
                logger.info(f"预签名 URL 生成成功: {s3_key}")
                return url
        except Exception as e:
            logger.error(f"生成预签名 URL 失败 {s3_key}: {e}")
            return None

    async def generate_presigned_url_from_direct_url(
        self,
        direct_url: str,
        expiration: int = 3600,
    ) -> Optional[str]:
        """
        从直链 URL 生成预签名 URL

        Args:
            direct_url: 直链 URL
            expiration: URL 有效期（秒）

        Returns:
            Optional[str]: 预签名 URL，生成失败返回 None
        """
        try:
            # 解析出 bucket 和 key
            if s3_config.addressing_style == "virtual":
                # 虚拟主机风格：https://bucket-name.endpoint/key
                url_without_protocol = direct_url.replace("https://", "").replace(
                    "http://", ""
                )
                first_slash = url_without_protocol.find("/")
                bucket_name = self._bucket_name
                s3_key = url_without_protocol[first_slash + 1 :]
            else:
                # 路径风格：https://endpoint/bucket-name/key
                url_without_protocol = direct_url.replace("https://", "").replace(
                    "http://", ""
                )
                path_parts = url_without_protocol.split("/")
                bucket_name = self._bucket_name
                s3_key = "/".join(path_parts[2:])

            return await self.generate_presigned_url(
                s3_key=s3_key,
                bucket=bucket_name,
                expiration=expiration,
                method="get_object",
            )
        except Exception as e:
            logger.error(f"从直链 URL 生成预签名 URL 失败 {direct_url}: {e}")
            return None

    async def object_exists(self, s3_key: str, bucket: Optional[str] = None) -> bool:
        """
        检查 S3 对象是否存在

        Args:
            s3_key: S3 对象键
            bucket: 存储桶名称

        Returns:
            bool: 存在返回 True
        """
        bucket_name = bucket or self._bucket_name

        try:
            async with self.get_client() as s3:
                await s3.head_object(Bucket=bucket_name, Key=s3_key)
                return True
        except Exception:
            return False

    async def copy_object(
        self,
        source_key: str,
        dest_key: str,
        source_bucket: Optional[str] = None,
        dest_bucket: Optional[str] = None,
    ) -> S3CopyResult:
        """
        复制 S3 对象

        Args:
            source_key: 源对象键
            dest_key: 目标对象键
            source_bucket: 源存储桶名称
            dest_bucket: 目标存储桶名称

        Returns:
            S3CopyResult: 复制结果，包含成功状态和目标直链 URL
        """
        src_bucket_name = source_bucket or self._bucket_name
        dst_bucket_name = dest_bucket or self._bucket_name

        copy_source = {"Bucket": src_bucket_name, "Key": source_key}

        try:
            async with self.get_client() as s3:
                await s3.copy_object(
                    CopySource=copy_source,  # type: ignore
                    Bucket=dst_bucket_name,
                    Key=dest_key,  # type: ignore
                )
                logger.info(f"对象复制成功: {source_key} -> {dest_key}")
                return S3CopyResult(
                    success=True,
                    source_key=source_key,
                    dest_key=dest_key,
                    dest_url=self._get_object_url(dest_key, dst_bucket_name),
                )
        except Exception as e:
            logger.error(f"对象复制失败 {source_key} -> {dest_key}: {e}")
            return S3CopyResult(
                success=False,
                source_key=source_key,
                dest_key=dest_key,
                error=str(e),
            )

    async def initialize(self):
        """初始化 S3 服务"""
        logger.info("S3 服务初始化完成")

    async def shutdown(self):
        """关闭 S3 服务"""
        logger.info("S3 服务已关闭")