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 有 ?,用 ¶m=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: 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 ]