video-gen/video-gen-api/app/services/resource_signed_url_service.py

from __future__ import annotations

import hashlib
import hmac
import time
from typing import Any, Optional
from urllib.parse import urlsplit

from app.config import settings


class ResourceSignedUrlError(RuntimeError):
    """
    资源签名 URL 生成异常。
    """
    pass


def _to_int(value: Any, default: int) -> int:
    """
    安全转换整数。
    """
    try:
        return int(value)
    except (TypeError, ValueError):
        return default


def _get_sign_secret(secret: Optional[str] = None) -> str:
    """
    获取资源签名密钥。

    优先级：
    1. 函数传入 secret
    2. app.config.settings.RESOURCE_SIGN_SECRET
    """
    sign_secret = secret or getattr(settings, "RESOURCE_SIGN_SECRET", "")

    if not sign_secret or not str(sign_secret).strip():
        raise ResourceSignedUrlError("RESOURCE_SIGN_SECRET 未配置")

    return str(sign_secret)


def _get_sign_expire_seconds(expire_seconds: Optional[int] = None) -> int:
    """
    获取资源签名有效期秒数。
    """
    if expire_seconds is not None:
        seconds = _to_int(expire_seconds, 3600)
    else:
        seconds = _to_int(getattr(settings, "RESOURCE_SIGN_EXPIRE_SECONDS", 3600), 3600)

    if seconds <= 0:
        seconds = 3600

    return seconds


def _get_expire_arg_name() -> str:
    """
    获取过期时间参数名。
    默认 exp。
    """
    name = getattr(settings, "RESOURCE_SIGN_ARG_EXPIRE", "exp")
    name = str(name or "exp").strip()
    return name or "exp"


def _get_signature_arg_name() -> str:
    """
    获取签名参数名。
    默认 sign。
    """
    name = getattr(settings, "RESOURCE_SIGN_ARG_SIGNATURE", "sign")
    name = str(name or "sign").strip()
    return name or "sign"


def _extract_sign_uri(resource_url: str) -> str:
    """
    提取用于签名的 URI path。

    例如：
    http://www.test6.com/generation/video/a.mp4?x=1

    用于签名的是：
    /generation/video/a.mp4

    注意：
    OpenResty/Lua 侧建议使用 ngx.var.uri 或 r.uri 参与签名，
    不要使用完整 URL，也不要包含 query string。
    """
    if not resource_url or not str(resource_url).strip():
        raise ResourceSignedUrlError("resource_url 不能为空")

    url = str(resource_url).strip()
    parsed = urlsplit(url)

    sign_uri = parsed.path or url

    if not sign_uri.startswith("/"):
        sign_uri = "/" + sign_uri

    return sign_uri


def generate_resource_signature(
    resource_url: str,
    expires_at: int,
    secret: Optional[str] = None,
) -> str:
    """
    生成资源 URL 签名。

    签名规则：

        message = "{uri}:{exp}"
        sign = hmac_sha256(secret, message).hexdigest()

    例如：

        uri = "/generation/video/a.mp4"
        exp = 1780000000
        message = "/generation/video/a.mp4:1780000000"

    OpenResty/Lua 侧必须使用完全一致的 message 规则。
    """
    sign_secret = _get_sign_secret(secret)
    sign_uri = _extract_sign_uri(resource_url)

    expires_at = _to_int(expires_at, 0)
    if expires_at <= 0:
        raise ResourceSignedUrlError("expires_at 必须是有效的 Unix 时间戳")

    message = f"generation_resource_controller:{sign_uri}:{expires_at}"

    return hmac.new(
        sign_secret.encode("utf-8"),
        message.encode("utf-8"),
        hashlib.sha256,
    ).hexdigest()


def _append_query_params(resource_url: str, params: dict[str, Any]) -> str:
    """
    按用户要求追加 URL 参数：

    - 原 URL 有 ?，用 &param=value 追加
    - 原 URL 没有 ?，用 ?param=value 追加

    同时兼容 #fragment，参数会追加到 # 前面。
    """
    url = str(resource_url).strip()

    if not url:
        raise ResourceSignedUrlError("resource_url 不能为空")

    base_url = url
    fragment = ""

    if "#" in url:
        base_url, fragment_part = url.split("#", 1)
        fragment = "#" + fragment_part

    separator = "&" if "?" in base_url else "?"

    query_string = "&".join(
        f"{key}={value}"
        for key, value in params.items()
        if key and value is not None
    )

    if not query_string:
        return url

    return f"{base_url}{separator}{query_string}{fragment}"


def build_resource_signed_url(
    resource_url: str | None,
    expire_seconds: Optional[int] = None,
    secret: Optional[str] = None,
    now_ts: Optional[int] = None,
) -> str:
    """
    生成带时效签名的资源 URL。

    参数：
        resource_url:
            外部传入的资源 URL，可以是完整 URL，也可以是 path。

            例如：
                http://www.test6.com/generation/video/a.mp4
                http://www.test6.com/generation/video/a.mp4?from=history
                /generation/video/a.mp4

        expire_seconds:
            有效期秒数，不传则使用 settings.RESOURCE_SIGN_EXPIRE_SECONDS。

        secret:
            可选，自定义签名密钥。不传则使用 settings.RESOURCE_SIGN_SECRET。

        now_ts:
            可选，当前时间戳。主要用于单元测试，正常业务不需要传。

    返回：
        带 exp 和 sign 参数的 URL。

    示例：
        http://www.test6.com/generation/video/a.mp4?exp=1780000000&sign=xxxx

        http://www.test6.com/generation/video/a.mp4?from=history&exp=1780000000&sign=xxxx
    """
    if not resource_url or not resource_url.startswith("/generate/"):
        return resource_url
    
    seconds = _get_sign_expire_seconds(expire_seconds)
    current_ts = _to_int(now_ts, int(time.time())) if now_ts is not None else int(time.time())

    expires_at = current_ts + seconds

    expire_arg_name = _get_expire_arg_name()
    signature_arg_name = _get_signature_arg_name()

    signature = generate_resource_signature(
        resource_url=resource_url,
        expires_at=expires_at,
        secret=secret,
    )

    return _append_query_params(
        resource_url=resource_url,
        params={
            expire_arg_name: expires_at,
            signature_arg_name: signature,
        },
    )


def build_resource_signed_urls(
    resource_urls: list[str],
    expire_seconds: Optional[int] = None,
    secret: Optional[str] = None,
) -> list[str]:
    """
    批量生成资源签名 URL。

    用于历史记录列表、资源列表等场景。
    """
    if not resource_urls:
        return []

    now_ts = int(time.time())

    return [
        build_resource_signed_url(
            resource_url=url,
            expire_seconds=expire_seconds,
            secret=secret,
            now_ts=now_ts,
        )
        for url in resource_urls
        if url
    ]