xiaoxia-saas/packages/shared/storage.py

"""统一存储服务 — API 和 Worker 共用的唯一存储入口。

实现 StoragePort 端口接口，整合原来分散在各处的存储能力：
- API端 SharedStorageService 的全部能力（上传/下载/签名URL/直传POST）
- Worker端 oss_helpers 的高级能力（分片上传/超时保护/HTTP下载/Asset路径解析）

所有服务都通过这个统一入口与存储交互，消除重复实现。

P1 (2026-09-28) OSS 双 endpoint 分离：
- 内部 bucket（self.bucket）：使用 internal endpoint（VPC 千兆带宽），
  用于所有 SDK 上传/下载/删除/object_exists 操作；
- 公网 bucket（self.public_bucket）：使用公网 endpoint，仅用于 sign_url
  生成给前端/P4000/MediaKit 等外网访问方用的预签名 URL；
- public_url 永远拼公网域名，不随 internal endpoint 变化。
"""

from __future__ import annotations

import base64
import datetime as dt
import hashlib
import hmac
import json
import logging
import os
import threading
from pathlib import Path
from typing import Optional
from urllib.parse import unquote, urlparse

import requests

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

from packages.config import get_shared_settings
from packages.ports.storage_port import StoragePort

logger = logging.getLogger(__name__)

# ── OSS 高级配置（从 oss_helpers 合并）─────────────────────────────────
OSS_CONNECT_TIMEOUT = 10  # 连接超时（秒）
OSS_UPLOAD_TOTAL_TIMEOUT = 300  # 单文件上传总超时（秒）
OSS_MULTIPART_THRESHOLD = 100 * 1024 * 1024  # 分片上传阈值：100MB
OSS_PART_SIZE = 8 * 1024 * 1024  # 分片大小：8MB
OSS_MULTIPART_NUM_THREADS = 3  # 分片上传并发数
OSS_HTTP_DOWNLOAD_TIMEOUT = 300  # HTTP下载超时（秒）


def _make_bucket(
    auth,
    endpoint: str,
    bucket_name: str,
    *,
    connect_timeout: int = OSS_CONNECT_TIMEOUT,
    app_name: str = "",
):
    """构造 oss2.Bucket，自动补 https:// 前缀。"""
    if not endpoint.startswith(("http://", "https://")):
        endpoint = f"https://{endpoint}"
    kwargs: dict = {"connect_timeout": connect_timeout}
    if app_name:
        kwargs["app_name"] = app_name
    return oss2.Bucket(auth, endpoint, bucket_name, **kwargs)


class SharedStorageService(StoragePort):
    """统一存储服务 — 实现 StoragePort，API 和 Worker 共用。"""

    # 类级默认值，方便单测 mock __init__ 后实例仍有这些属性
    bucket: Optional[object] = None
    public_bucket: Optional[object] = None
    public_endpoint: str = ""
    internal_endpoint: str = ""
    public_url: str = ""
    local_url_prefix: str = "/generated-files"
    bucket_name: str = ""

    def __init__(self):
        settings = get_shared_settings()
        self.bucket_name = settings.oss_bucket_name
        self.public_endpoint = settings.oss_endpoint  # 公网 endpoint，用于签名 URL
        self.internal_endpoint = settings.effective_oss_internal_endpoint  # 内网 endpoint，SDK 用
        self.public_url = f"https://{settings.oss_bucket_name}.{self._public_host()}"
        self.local_url_prefix = os.getenv("GENERATED_FILES_URL_PREFIX", "/generated-files")
        self.bucket: Optional[object] = None  # internal: SDK 上传/下载/删除
        self.public_bucket: Optional[object] = None  # public: sign_url 给外网

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

        has_key_id = bool(self.access_key_id)
        has_key_secret = bool(self.access_key_secret)

        if has_key_id and has_key_secret:
            if oss2 is not None:
                try:
                    auth = oss2.Auth(self.access_key_id, self.access_key_secret)
                    self.bucket = _make_bucket(
                        auth,
                        self.internal_endpoint,
                        self.bucket_name,
                        app_name="xiaoxia-internal",
                    )
                    self.public_bucket = _make_bucket(
                        auth,
                        self.public_endpoint,
                        self.bucket_name,
                        app_name="xiaoxia-public",
                    )
                    same_ep = self.internal_endpoint == self.public_endpoint
                    logger.info(
                        "OSS initialized: public_ep=%s internal_ep=%s bucket=%s dual=%s",
                        self.public_endpoint,
                        self.internal_endpoint,
                        self.bucket_name,
                        "no" if same_ep else "yes",
                    )
                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))

    def _public_host(self) -> str:
        ep = self.public_endpoint
        if ep.startswith("https://"):
            return ep[len("https://") :]
        if ep.startswith("http://"):
            return ep[len("http://") :]
        return ep

    # ── 诊断 ───────────────────────────────────────────────────────────

    def diagnose(self) -> None:
        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诊断] public_ep=%s internal_ep=%s bucket=%s ak=%s",
            self.public_endpoint,
            self.internal_endpoint,
            self.bucket_name,
            key_id_display,
        )
        if self.bucket is None:
            logger.error("[OSS诊断] ❌ bucket(internal)=None")
        else:
            logger.info("[OSS诊断] ✅ bucket(internal) 就绪")
        if self.public_bucket is None:
            logger.error("[OSS诊断] ❌ public_bucket=None")
        else:
            logger.info("[OSS诊断] ✅ public_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 _normalize_storage_key(self, storage_key_or_url: str) -> str:
        if storage_key_or_url.startswith("http://") or storage_key_or_url.startswith("https://"):
            parsed = urlparse(storage_key_or_url)
            return unquote(parsed.path.lstrip("/"))
        return storage_key_or_url.lstrip("/")

    def normalize_storage_key(self, storage_key_or_url: str) -> str:
        return self._normalize_storage_key(storage_key_or_url)

    # ── 上传（SDK 走 internal endpoint）───────────────────────────────

    def upload_file(
        self,
        file_or_path: str | Path | object,
        storage_key: str,
        content_type: str = "application/octet-stream",
    ) -> str:
        if self.bucket is None:
            raise RuntimeError("OSS storage is not configured")
        try:
            if isinstance(file_or_path, (str, Path)):
                self.bucket.put_object_from_file(
                    storage_key,
                    str(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}") from e

    def upload_file_smart(
        self,
        local_path: str | Path,
        storage_key: str,
    ) -> Optional[str]:
        local_path = Path(local_path)
        if not local_path.exists():
            logger.error("上传文件不存在: %s", local_path)
            return None
        if self.bucket is None:
            logger.error("OSS未配置，无法上传: %s", storage_key[:80])
            return None

        result: dict = {"url": None, "error": None, "file_size": 0}
        done = threading.Event()

        def _do_upload():
            try:
                try:
                    file_size = local_path.stat().st_size
                    result["file_size"] = file_size
                    use_multipart = file_size >= OSS_MULTIPART_THRESHOLD
                except OSError:
                    use_multipart = False
                    file_size = 0

                if use_multipart:
                    logger.info(
                        "大文件分片上传(internal): key=%s size=%.1fMB part=%dMB threads=%d",
                        storage_key[:80],
                        file_size / 1048576,
                        OSS_PART_SIZE // 1048576,
                        OSS_MULTIPART_NUM_THREADS,
                    )
                    oss2.resumable_upload(
                        self.bucket,
                        storage_key,
                        str(local_path),
                        multipart_threshold=OSS_MULTIPART_THRESHOLD,
                        part_size=OSS_PART_SIZE,
                        num_threads=OSS_MULTIPART_NUM_THREADS,
                    )
                else:
                    self.bucket.put_object_from_file(storage_key, str(local_path))
                result["url"] = f"{self.public_url}/{storage_key}"
            except Exception as e:
                result["error"] = e
                logger.exception("上传 OSS 失败: %s", storage_key)
            finally:
                done.set()

        t = threading.Thread(target=_do_upload, daemon=True)
        t.start()
        finished = done.wait(timeout=OSS_UPLOAD_TOTAL_TIMEOUT)

        if not finished:
            logger.error(
                "OSS 上传超时(%ds): key=%s size=%.1fMB",
                OSS_UPLOAD_TOTAL_TIMEOUT,
                storage_key[:80],
                result["file_size"] / 1048576 if result["file_size"] else 0,
            )
            return None
        return None if result["error"] else result["url"]

    # ── 下载（SDK 走 internal endpoint）───────────────────────────────

    def download_file(self, storage_key: str, local_path: str | Path) -> None:
        if self.bucket is None:
            raise RuntimeError("OSS storage is not configured")
        local_path = Path(local_path)
        os.makedirs(local_path.parent, exist_ok=True)
        try:
            self.bucket.get_object_to_file(
                self._normalize_storage_key(storage_key),
                str(local_path),
            )
        except Exception as e:
            raise Exception(f"Failed to download file from OSS: {e}") from e

    def download_asset(self, asset_storage_key: str, local_path: str | Path) -> bool:
        local_path = Path(local_path)
        os.makedirs(local_path.parent, exist_ok=True)
        if asset_storage_key.startswith(("http://", "https://")):
            return self._download_via_http(asset_storage_key, local_path)
        if self.bucket is None:
            logger.error("OSS not configured: %s", asset_storage_key[:80])
            return False
        try:
            self.bucket.get_object_to_file(
                self._normalize_storage_key(asset_storage_key),
                str(local_path),
            )
            return local_path.exists() and local_path.stat().st_size > 0
        except Exception:
            logger.exception("下载素材失败: %s", asset_storage_key)
            return False

    def _download_via_http(self, url: str, local_path: Path) -> bool:
        try:
            resp = requests.get(url, stream=True, timeout=OSS_HTTP_DOWNLOAD_TIMEOUT)
            resp.raise_for_status()
            with open(local_path, "wb") as f:
                for chunk in resp.iter_content(chunk_size=8 * 1024 * 1024):
                    if chunk:
                        f.write(chunk)
            return local_path.exists() and local_path.stat().st_size > 0
        except Exception:
            logger.exception("HTTP下载失败: %s", url[:100])
            return False

    # ── URL 生成（sign_url 用 public_bucket 签公网域名）───────────────

    def get_url(self, storage_key: str) -> str:
        return f"{self.public_url}/{storage_key}"

    def _sign_bucket(self):
        """签名优先用 public_bucket，回退到 bucket。"""
        return self.public_bucket or self.bucket

    def get_download_url(self, storage_key_or_url: str, expires_seconds: int = 3600) -> str:
        sign_bucket = self._sign_bucket()
        if sign_bucket is None:
            if self._is_local_generated_url(storage_key_or_url):
                return storage_key_or_url
            logger.warning("OSS bucket not configured, returning raw 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 = sign_bucket.sign_url("GET", storage_key, expires_seconds)
            logger.info(
                "signed URL generated for key=%s prefix=%s",
                storage_key[:80],
                signed[:60],
            )
            return signed
        except Exception:
            logger.exception("get_download_url: sign_url 失败，返回 raw URL: %s", storage_key[:200])
            return self.get_url(storage_key)

    def get_upload_url(
        self,
        storage_key_or_url: str,
        expires_seconds: int = 3600,
        content_type: str = "video/mp4",
    ) -> str:
        sign_bucket = self._sign_bucket()
        if sign_bucket is None:
            if self._is_local_generated_url(storage_key_or_url):
                return storage_key_or_url
            logger.warning("get_upload_url: OSS 未配置，返回 raw 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:
            headers = {"Content-Type": content_type} if content_type else None
            signed = sign_bucket.sign_url("PUT", storage_key, expires_seconds, headers=headers)
            logger.info(
                "get_upload_url: 公网签名PUT URL已生成 key=%s prefix=%s",
                storage_key[:80],
                signed[:60],
            )
            return signed
        except Exception:
            logger.exception("get_upload_url: sign_url 失败，返回 raw URL: %s", storage_key[:200])
            return self.get_url(storage_key)

    def create_direct_upload_post(
        self,
        storage_key: str,
        content_type: str,
        max_size_bytes: int,
        expires_seconds: int,
    ) -> dict[str, object]:
        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.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,
            },
        }

    # ── 文件操作（internal endpoint）──────────────────────────────────

    def delete_file(self, storage_key: str) -> None:
        if self.bucket is None:
            return
        try:
            self.bucket.delete_object(storage_key)
        except Exception as error:
            logger.warning("OSS delete 失败", extra={"storage_key": storage_key, "error": str(error)})

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

    # ── Asset 路径解析（Worker 用）────────────────────────────────────

    def resolve_asset_path(self, asset_id: str, work_dir: str | Path) -> Optional[Path]:
        from video_processing.path_security import (  # type: ignore[import-not-found]
            PathSecurityError,
            get_allowed_local_dirs,
            is_in_allowed_dirs,
            sanitize_filename,
        )

        if not asset_id or not isinstance(asset_id, str):
            return None
        work_dir = Path(work_dir)
        os.makedirs(work_dir, exist_ok=True)
        if "\x00" in asset_id:
            logger.warning("asset_id 含空字节，拒绝: %s", asset_id[:50])
            return None
        if asset_id.startswith("/") and os.path.exists(asset_id):
            try:
                resolved = Path(asset_id).resolve()
                if is_in_allowed_dirs(resolved, get_allowed_local_dirs()):
                    return resolved
                logger.warning(
                    "本地素材路径不在允许目录: %s allowed=%s",
                    asset_id[:80],
                    get_allowed_local_dirs(),
                )
                return None
            except (OSError, PathSecurityError):
                return None
        cache_hash = hashlib.sha256(asset_id.encode()).hexdigest()[:16]
        safe_name = sanitize_filename(cache_hash)
        cached_path = work_dir / f"{safe_name}.mp4"
        if cached_path.exists() and cached_path.stat().st_size > 0:
            return cached_path
        safe_key = self.normalize_storage_key(asset_id)
        if ".." in safe_key or safe_key.startswith("/"):
            logger.warning("asset_id 含路径遍历: %s", asset_id[:80])
            return None
        if self.download_asset(safe_key, cached_path):
            return cached_path
        return None

    def resolve_asset_ids_to_paths(
        self,
        asset_ids: list[str],
        work_dir: str | Path,
    ) -> dict[str, Path]:
        result: dict[str, Path] = {}
        for aid in asset_ids:
            p = self.resolve_asset_path(aid, work_dir)
            if p:
                result[aid] = p
        return result


# ── 单例 ────────────────────────────────────────────────────────────────

_storage_service: Optional[SharedStorageService] = None


def get_shared_storage_service() -> SharedStorageService:
    global _storage_service
    if _storage_service is None:
        _storage_service = SharedStorageService()
        _storage_service.diagnose()
    return _storage_service


def get_storage_service() -> SharedStorageService:
    """向后兼容别名。"""
    return get_shared_storage_service()