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 urllib.parse import urlparse
from typing import Optional

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

        if settings.oss_access_key_id and settings.oss_access_key_secret:
            if oss2 is not None:
                auth = oss2.Auth(
                    settings.oss_access_key_id,
                    settings.oss_access_key_secret,
                )
                self.bucket = oss2.Bucket(
                    auth,
                    settings.oss_endpoint,
                    settings.oss_bucket_name,
                )
        self.access_key_id = settings.oss_access_key_id
        self.access_key_secret = settings.oss_access_key_secret
        self.endpoint = settings.oss_endpoint

    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
            return self.get_url(self._normalize_storage_key(storage_key_or_url))

        storage_key = self._normalize_storage_key(storage_key_or_url)
        try:
            return self.bucket.sign_url("GET", storage_key, expires_seconds)
        except Exception:
            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()
    return _storage_service


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