xiaoxia-saas/apps/worker/video_processing/oss_helpers.py

"""OSS 工具函数 — 从 generation.py 提取的共享 OSS 操作.

提供 OSS 配置读取、Bucket 创建、素材上传/下载、asset_id → 本地路径解析
等能力，供 render_edit_plan 和 generate_video 共同复用。
"""

from __future__ import annotations

import hashlib
import logging
import os
import threading
from pathlib import Path
from urllib.parse import urlparse

import oss2
import requests

from packages.shared.config import get_shared_settings

logger = logging.getLogger(__name__)

# OSS 上传配置
OSS_CONNECT_TIMEOUT = 10  # 连接超时（秒），防止 TCP 握手挂死
OSS_UPLOAD_TOTAL_TIMEOUT = 900  # 单文件上传总超时（秒），防止网络慢时无限卡住
OSS_MULTIPART_THRESHOLD = 100 * 1024 * 1024  # 分片上传阈值：100MB 以上走分片
OSS_PART_SIZE = 8 * 1024 * 1024  # 分片大小：8MB
OSS_MULTIPART_NUM_THREADS = 3  # 分片上传并发数


# ── OSS 配置 ──────────────────────────────────────────────────────────────────


def oss_settings() -> tuple[str, str, str, str] | None:
    """获取 OSS 配置。

    统一使用 SharedSettings 读取配置，与 SharedStorageService 保持一致，
    支持从 .env 文件加载，避免两套配置路径不一致。

    Returns:
        (access_key_id, access_key_secret, endpoint, bucket_name) 元组，
        配置缺失时返回 None。
    """
    settings = get_shared_settings()
    access_key_id = settings.oss_access_key_id
    access_key_secret = settings.oss_access_key_secret
    endpoint = settings.oss_endpoint
    bucket_name = settings.oss_bucket_name
    if not all([access_key_id, access_key_secret, endpoint, bucket_name]):
        return None
    return access_key_id, access_key_secret, endpoint, bucket_name


def oss_bucket() -> oss2.Bucket | None:
    """获取 OSS Bucket 实例。

    P0-2 修复：endpoint 不带 scheme 时自动补 https:// 前缀，
    确保 sign_url 等依赖 scheme 的方法返回 HTTPS URL。

    P0-staging 修复：增加 connect_timeout=10s，防止网络抖动时
    TCP 握手阶段无限挂死，导致 worker 进程卡死。

    Returns:
        oss2.Bucket 实例，配置缺失时返回 None。
    """
    settings = oss_settings()
    if settings is None:
        return None
    access_key_id, access_key_secret, endpoint, bucket_name = settings
    # endpoint 无 scheme 时补 https://，与 API 端 storage.py 保持一致
    if not endpoint.startswith(("http://", "https://")):
        endpoint = f"https://{endpoint}"
    return oss2.Bucket(
        oss2.Auth(access_key_id, access_key_secret),
        endpoint,
        bucket_name,
        connect_timeout=OSS_CONNECT_TIMEOUT,
    )


def normalize_storage_key(storage_key_or_url: str) -> str:
    """标准化存储键 — 如果是完整 URL 则提取 path 部分。

    Examples:
        "https://bucket.oss-cn-hangzhou.aliyuncs.com/path/to/file.mp4"
        → "path/to/file.mp4"
        "path/to/file.mp4" → "path/to/file.mp4"
    """
    if storage_key_or_url.startswith(("http://", "https://")):
        return urlparse(storage_key_or_url).path.lstrip("/")
    return storage_key_or_url.lstrip("/")


# ── 上传 / 下载 ───────────────────────────────────────────────────────────────


def download_asset(asset_storage_key: str, local_path: Path) -> bool:
    """从 OSS 下载素材文件到本地路径。

    自动识别输入类型：
    - 完整 URL（http:// 或 https:// 开头）→ 走 HTTP 下载（支持预签名URL）
    - OSS 存储键 → 走 oss2 SDK 下载

    Args:
        asset_storage_key: 素材的存储键或完整 URL
        local_path: 本地保存路径

    Returns:
        True 表示下载成功，False 表示失败。
    """
    # 完整URL走HTTP下载（兼容预签名URL）
    if asset_storage_key.startswith(("http://", "https://")):
        return _download_via_http(asset_storage_key, local_path)

    # OSS存储键走SDK
    bucket = oss_bucket()
    if bucket is None:
        return False
    try:
        bucket.get_object_to_file(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(url: str, local_path: Path) -> bool:
    """通过 HTTP 下载文件（支持预签名 URL）。

    使用流式下载避免大文件内存溢出，超时 900s。
    """
    try:
        resp = requests.get(url, stream=True, timeout=900)
        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)
        return False


def upload_to_oss(local_path: Path | str, storage_key: str) -> str | None:
    """上传文件到 OSS，返回公开 URL。

    大文件（>100MB）自动走分片上传，降低内存峰值，减少 OOM 风险。
    上传加总超时保护（默认 900s），防止网络异常时无限挂死。

    Args:
        local_path: 本地文件路径（Path 或 str 均可）
        storage_key: 目标存储键

    Returns:
        公开访问 URL，上传失败或 OSS 未配置时返回 None。
    """
    local_path = Path(local_path)  # 统一转 Path，兼容 str 调用
    bucket = oss_bucket()
    if bucket is None:
        return None

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

    def _do_upload():
        try:
            # 尝试获取文件大小，用于分片判断和日志；stat 失败时 fallback 走普通上传
            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:
                # 分片上传：降低内存峰值，每片 8MB，3 线程并发
                logger.info(
                    "大文件分片上传: storage_key=%s, size=%.1fMB, part_size=%dMB, threads=%d",
                    storage_key[:80],
                    file_size / 1024 / 1024,
                    OSS_PART_SIZE // 1024 // 1024,
                    OSS_MULTIPART_NUM_THREADS,
                )
                oss2.resumable_upload(
                    bucket,
                    storage_key,
                    str(local_path),
                    multipart_threshold=OSS_MULTIPART_THRESHOLD,
                    part_size=OSS_PART_SIZE,
                    num_threads=OSS_MULTIPART_NUM_THREADS,
                )
            else:
                bucket.put_object_from_file(storage_key, str(local_path))

            # 构造返回 URL
            settings = oss_settings()
            if settings:
                _, _, endpoint, bucket_name = settings
                endpoint_clean = endpoint.replace("https://", "").replace("http://", "")
                result["url"] = f"https://{bucket_name}.{endpoint_clean}/{storage_key}"
        except Exception as e:
            result["error"] = e
            logger.exception("上传 OSS 失败: %s", storage_key)
        finally:
            done.set()

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

    if not finished:
        logger.error(
            "OSS 上传超时（%.0fs），强制中止: storage_key=%s, size=%.1fMB",
            OSS_UPLOAD_TOTAL_TIMEOUT,
            storage_key[:80],
            result["file_size"] / 1024 / 1024 if result["file_size"] else 0,
        )
        return None

    if result["error"]:
        return None

    return result["url"]


def get_signed_download_url(storage_key_or_url: str, expires_seconds: int = 3600) -> str | None:
    """生成预签名下载 URL（用于私有 bucket 的 URL 校验或临时下载）。

    Args:
        storage_key_or_url: 存储键或完整 URL（URL 会自动提取 path）
        expires_seconds: 签名有效期（秒）

    Returns:
        预签名 URL，失败或 OSS 未配置时返回 None。
    """
    bucket = oss_bucket()
    if bucket is None:
        return None
    try:
        storage_key = normalize_storage_key(storage_key_or_url)
        signed = bucket.sign_url("GET", storage_key, expires_seconds)
        logger.info("生成预签名URL: key=%s url_prefix=%s", storage_key[:80], signed[:60])
        return signed
    except Exception:
        logger.exception("生成预签名URL失败: %s", storage_key_or_url[:80])
        return None


# ── Asset 解析 ────────────────────────────────────────────────────────────────


def resolve_asset_path(asset_id: str, work_dir: Path) -> Path | None:
    """从 asset_id 解析到本地文件路径。

    策略（按优先级）：
    1. 如果 asset_id 是本地绝对路径（/var/storage/...）→ 安全校验后返回
    2. 如果 work_dir 下已有缓存文件 → 返回缓存路径
    3. 从 OSS 下载到 work_dir/{hash}.mp4 → 返回下载路径
    4. 下载失败 → 返回 None

    缓存策略：以 asset_id 的 SHA256 前 16 位为文件名，避免重复下载。

    安全：
    - 本地绝对路径必须在 ASSET_ALLOWED_DIRS 环境变量指定的目录内
    - 文件名经过 sanitize，防止路径遍历
    - 禁止空字节、控制字符
    """
    from video_processing.path_security import (
        PathSecurityError,
        get_allowed_local_dirs,
        is_in_allowed_dirs,
        sanitize_filename,
    )

    if not asset_id or not isinstance(asset_id, str):
        return None

    # 空字节检测
    if "\x00" in asset_id:
        logger.warning("asset_id 包含空字节，拒绝: %s", asset_id[:50])
        return None

    # 1. 本地绝对路径 — 必须在允许的目录内
    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
            else:
                logger.warning(
                    "本地素材路径不在允许目录内，拒绝: %s (allowed=%s)",
                    asset_id[:80],
                    get_allowed_local_dirs(),
                )
                return None
        except (OSError, PathSecurityError):
            return None

    # 2. 缓存命中（使用 hash 而非原始 ID，防止路径遍历）
    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

    # 3. 从 OSS 下载（先标准化 key，防止路径遍历注入）
    safe_key = normalize_storage_key(asset_id)
    # 额外校验：存储键不能包含 ../ 或绝对路径
    if ".." in safe_key or safe_key.startswith("/"):
        logger.warning("asset_id 包含路径遍历模式，拒绝下载: %s", asset_id[:80])
        return None

    if download_asset(safe_key, cached_path):
        return cached_path

    return None


def resolve_asset_ids_to_paths(
    asset_ids: list[str],
    work_dir: Path,
) -> dict[str, Path]:
    """批量解析 asset_id → 本地路径。

    Args:
        asset_ids: 素材 ID 列表
        work_dir: 工作目录

    Returns:
        {asset_id: local_path} 映射，仅包含成功解析的条目。
    """
    result: dict[str, Path] = {}
    for aid in asset_ids:
        local_path = resolve_asset_path(aid, work_dir)
        if local_path:
            result[aid] = local_path
    return result