xiaoxia-saas/packages/shared/storage.py

"""Shared OSS storage service for API and Worker."""

import base64
import datetime as dt
import hashlib
import hmac
import json
import logging
import os
from typing import Optional
from urllib.parse import urlparse

try:
    import oss2
except ImportError:  # pragma: no cover
    oss2 = None

from packages.shared.config import get_shared_settings

logger = logging.getLogger(__name__)


class SharedStorageService:
    """Shared OSS storage service."""

    def __init__(self):
        settings = get_shared_settings()
        self.bucket_name = settings.oss_bucket_name
        self.public_url = f"https://{settings.oss_bucket_name}.{settings.oss_endpoint}"
        self.local_url_prefix = os.getenv("GENERATED_FILES_URL_PREFIX", "/generated-files")
        self.bucket = None

        has_key_id = bool(settings.oss_access_key_id)
        has_key_secret = bool(settings.oss_access_key_secret)

        if has_key_id and has_key_secret:
            if oss2 is not None:
                try:
                    # P0-2 修复：oss2.Bucket 的 endpoint 必须带 https:// 前缀，
                    # 否则 sign_url 默认生成 HTTP URL。
                    bucket_endpoint = settings.oss_endpoint
                    if not bucket_endpoint.startswith(("http://", "https://")):
                        bucket_endpoint = f"https://{bucket_endpoint}"
                    auth = oss2.Auth(
                        settings.oss_access_key_id,
                        settings.oss_access_key_secret,
                    )
                    self.bucket = oss2.Bucket(
                        auth,
                        bucket_endpoint,
                        settings.oss_bucket_name,
                    )
                    logger.info(
                        "OSS initialized: endpoint=%s bucket=%s",
                        settings.oss_endpoint,
                        settings.oss_bucket_name,
                    )
                except Exception as error:
                    logger.error("Failed to initialize OSS bucket client: %s", error)
            else:
                logger.error("oss2 SDK is not installed — OSS operations will fail")
        else:
            missing = []
            if not has_key_id:
                missing.append("OSS_ACCESS_KEY_ID")
            if not has_key_secret:
                missing.append("OSS_ACCESS_KEY_SECRET")
            logger.error("OSS credentials not configured — missing: %s", ", ".join(missing))

        self.access_key_id = settings.oss_access_key_id
        self.access_key_secret = settings.oss_access_key_secret
        self.endpoint = settings.oss_endpoint

    def diagnose(self) -> None:
        """启动诊断：输出 OSS 配置状态，帮助排查预签名 URL 问题。"""
        key_id_display = (
            f"{self.access_key_id[:4]}...{self.access_key_id[-4:]}" if len(self.access_key_id) > 8 else "(empty)"
        )
        logger.info(
            "[OSS诊断] endpoint=%s bucket_name=%s access_key_id=%s",
            self.endpoint,
            self.bucket_name,
            key_id_display,
        )
        if self.bucket is None:
            logger.error(
                "[OSS诊断] ❌ bucket=None — 预签名URL不可用！"
                "原因: OSS_ACCESS_KEY_ID/OSS_ACCESS_KEY_SECRET 未配置或 oss2 未安装。"
                "请检查服务器 .env 文件（如 /var/lib/xiaoxia-saas-staging/.env）"
            )
        else:
            logger.info("[OSS诊断] ✅ bucket 已配置，预签名URL可用")

    def _is_local_generated_url(self, storage_key_or_url: str) -> bool:
        parsed = urlparse(storage_key_or_url)
        path = parsed.path if parsed.scheme else storage_key_or_url
        return path.startswith(f"{self.local_url_prefix}/")

    def upload_file(
        self,
        file_or_path,
        storage_key: str,
        content_type: str = "application/octet-stream",
    ) -> str:
        """Upload file to OSS."""
        if self.bucket is None:
            raise RuntimeError("OSS storage is not configured")

        try:
            if isinstance(file_or_path, str):
                self.bucket.put_object_from_file(storage_key, file_or_path, headers={"Content-Type": content_type})
            else:
                file_or_path.seek(0)
                self.bucket.put_object(storage_key, file_or_path, headers={"Content-Type": content_type})
            return f"{self.public_url}/{storage_key}"
        except Exception as e:
            raise Exception(f"Failed to upload file to OSS: {e}")

    def get_url(self, storage_key: str) -> str:
        """Get public URL for a file."""
        return f"{self.public_url}/{storage_key}"

    def download_file(self, storage_key: str, local_path: str):
        """Download file from OSS to local path."""
        if self.bucket is None:
            raise RuntimeError("OSS storage is not configured")

        try:
            os.makedirs(os.path.dirname(local_path), exist_ok=True)
            self.bucket.get_object_to_file(storage_key, local_path)
        except Exception as e:
            raise Exception(f"Failed to download file from OSS: {e}")

    def get_download_url(self, storage_key_or_url: str, expires_seconds: int = 3600) -> str:
        """Get signed download URL."""
        if self.bucket is None:
            if self._is_local_generated_url(storage_key_or_url):
                return storage_key_or_url
            logger.warning(
                "get_download_url: OSS bucket not configured, returning raw URL. storage_key_or_url=%s",
                storage_key_or_url[:200],
            )
            return self.get_url(self._normalize_storage_key(storage_key_or_url))

        storage_key = self._normalize_storage_key(storage_key_or_url)
        try:
            signed = self.bucket.sign_url("GET", storage_key, expires_seconds)
            logger.info(
                "get_download_url: signed URL generated. storage_key=%s url_prefix=%s",
                storage_key[:80],
                signed[:60],
            )
            return signed
        except Exception:
            logger.exception(
                "get_download_url: sign_url failed, falling back to raw URL. storage_key=%s",
                storage_key[:200],
            )
            return self.get_url(storage_key)

    def _normalize_storage_key(self, storage_key_or_url: str) -> str:
        """Extract storage key from URL."""
        if storage_key_or_url.startswith("http://") or storage_key_or_url.startswith("https://"):
            parsed = urlparse(storage_key_or_url)
            return parsed.path.lstrip("/")
        return storage_key_or_url.lstrip("/")

    def create_direct_upload_post(
        self,
        storage_key: str,
        content_type: str,
        max_size_bytes: int,
        expires_seconds: int,
    ) -> dict[str, object]:
        """Create browser direct upload POST form."""
        if not self.access_key_id or not self.access_key_secret:
            raise RuntimeError("OSS storage is not configured")
        normalized_key = self._normalize_storage_key(storage_key)
        if not normalized_key.startswith("uploads/"):
            raise ValueError("direct upload key must be under uploads/")

        expiration = (dt.datetime.now(dt.timezone.utc) + dt.timedelta(seconds=expires_seconds)).strftime(
            "%Y-%m-%dT%H:%M:%S.000Z"
        )
        policy = {
            "expiration": expiration,
            "conditions": [
                {"bucket": self.bucket_name},
                {"key": normalized_key},
                ["content-length-range", 1, max_size_bytes],
                ["starts-with", "$Content-Type", content_type.split("/", 1)[0] + "/" if "/" in content_type else ""],
            ],
        }
        encoded_policy = base64.b64encode(json.dumps(policy, separators=(",", ":")).encode("utf-8")).decode("ascii")
        signature = base64.b64encode(
            hmac.new(self.access_key_secret.encode("utf-8"), encoded_policy.encode("utf-8"), hashlib.sha1).digest()
        ).decode("ascii")

        return {
            "url": self.public_url,
            "method": "POST",
            "storage_key": normalized_key,
            "expires_at": expiration,
            "fields": {
                "key": normalized_key,
                "OSSAccessKeyId": self.access_key_id,
                "policy": encoded_policy,
                "Signature": signature,
                "success_action_status": "201",
                "Content-Type": content_type,
            },
        }

    def delete_file(self, storage_key: str):
        """Delete file from OSS."""
        if self.bucket is None:
            return
        try:
            self.bucket.delete_object(storage_key)
        except Exception as error:
            logger.warning("Failed to delete file from OSS", extra={"storage_key": storage_key, "error": str(error)})

    def file_exists(self, storage_key: str) -> bool:
        """Check if file exists."""
        if self.bucket is None:
            return False
        return self.bucket.object_exists(storage_key)


_storage_service: Optional[SharedStorageService] = None


def get_shared_storage_service() -> SharedStorageService:
    """Get shared storage service instance (global singleton)."""
    global _storage_service
    if _storage_service is None:
        _storage_service = SharedStorageService()
        _storage_service.diagnose()
    return _storage_service


# Backward compatibility alias
def get_storage_service() -> SharedStorageService:
    """Backward compatibility: returns shared storage service."""
    return get_shared_storage_service()