Teknium
422 lines
16 KiB
Python
422 lines
16 KiB
Python
"""
|
||
yuanbao_media.py — 元宝平台媒体处理模块
|
||
|
||
提供 COS 上传、文件下载、TIM 媒体消息构建等功能。移植自 TypeScript 版 media.ts
|
||
（yuanbao-openclaw-plugin），使用 httpx 替代 cos-nodejs-sdk-v5。
|
||
|
||
COS 上传流程：
|
||
1. genUploadInfo 获取临时凭证（tmpSecretId/tmpSecretKey/sessionToken）
|
||
2. 临时凭证 HMAC-SHA1 签名构建 Authorization 头
|
||
3. HTTP PUT 上传到 COS
|
||
|
||
TIM 消息体：build_image_msg_body() → TIMImageElem，build_file_msg_body() → TIMFileElem
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import hmac
|
||
import logging
|
||
import os
|
||
import secrets
|
||
import struct
|
||
import time
|
||
import urllib.parse
|
||
from typing import Optional, Any
|
||
|
||
import httpx
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
UPLOAD_INFO_PATH = "/api/resource/genUploadInfo"
|
||
DEFAULT_API_DOMAIN = "yuanbao.tencent.com"
|
||
DEFAULT_MAX_SIZE_MB = 50
|
||
|
||
# COS 加速域名（优先使用全球加速）
|
||
COS_USE_ACCELERATE = True
|
||
|
||
# MIME → image_format 数字（TIM 协议字段）
|
||
_MIME_TO_IMAGE_FORMAT: dict[str, int] = {
|
||
"image/jpeg": 1, "image/jpg": 1, "image/gif": 2, "image/png": 3, "image/bmp": 4,
|
||
"image/webp": 255, "image/heic": 255, "image/tiff": 255,
|
||
}
|
||
|
||
_EXT_TO_MIME: dict[str, str] = {
|
||
".jpg": "image/jpeg",
|
||
".jpeg": "image/jpeg",
|
||
".png": "image/png",
|
||
".gif": "image/gif",
|
||
".webp": "image/webp",
|
||
".bmp": "image/bmp",
|
||
".heic": "image/heic",
|
||
".tiff": "image/tiff",
|
||
".ico": "image/x-icon",
|
||
".pdf": "application/pdf",
|
||
".doc": "application/msword",
|
||
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||
".xls": "application/vnd.ms-excel",
|
||
".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||
".ppt": "application/vnd.ms-powerpoint",
|
||
".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
||
".txt": "text/plain",
|
||
".zip": "application/zip",
|
||
".tar": "application/x-tar",
|
||
".gz": "application/gzip",
|
||
".mp3": "audio/mpeg",
|
||
".mp4": "video/mp4",
|
||
".wav": "audio/wav",
|
||
".ogg": "audio/ogg",
|
||
".webm": "video/webm",
|
||
}
|
||
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".heic", ".tiff", ".ico"}
|
||
|
||
|
||
def guess_mime_type(filename: str) -> str:
|
||
"""根据文件扩展名猜测 MIME 类型。"""
|
||
return _EXT_TO_MIME.get(os.path.splitext(filename)[-1].lower(), "application/octet-stream")
|
||
|
||
|
||
def is_image(filename: str, mime_type: str = "") -> bool:
|
||
"""判断是否为图片类型（MIME 前缀或扩展名）。"""
|
||
return mime_type.startswith("image/") or os.path.splitext(filename)[-1].lower() in _IMAGE_EXTS
|
||
|
||
|
||
def get_image_format(mime_type: str) -> int:
|
||
"""获取 TIM 图片格式编号。"""
|
||
return _MIME_TO_IMAGE_FORMAT.get(mime_type.lower(), 255)
|
||
|
||
|
||
def md5_hex(data: bytes) -> str:
|
||
return hashlib.md5(data).hexdigest()
|
||
|
||
|
||
def generate_file_id() -> str:
|
||
"""随机文件 ID（32 位 hex）。"""
|
||
return secrets.token_hex(16)
|
||
|
||
|
||
# ============ 图片尺寸解析（纯 Python，无需 Pillow） ============
|
||
|
||
def parse_image_size(data: bytes) -> Optional[dict[str, int]]:
|
||
"""解析图片宽高（JPEG/PNG/GIF/WebP），返回 {"width", "height"} 或 None。"""
|
||
return _parse_png_size(data) or _parse_jpeg_size(data) or _parse_gif_size(data) or _parse_webp_size(data)
|
||
|
||
|
||
def _wh(w: int, h: int) -> dict[str, int]:
|
||
return {"width": w, "height": h}
|
||
|
||
|
||
def _parse_png_size(buf: bytes) -> Optional[dict[str, int]]:
|
||
if len(buf) < 24 or buf[:4] != b"\x89PNG":
|
||
return None
|
||
return _wh(*struct.unpack(">II", buf[16:24]))
|
||
|
||
|
||
def _parse_jpeg_size(buf: bytes) -> Optional[dict[str, int]]:
|
||
if len(buf) < 4 or buf[0] != 0xFF or buf[1] != 0xD8:
|
||
return None
|
||
i = 2
|
||
while i < len(buf) - 9:
|
||
if buf[i] != 0xFF:
|
||
i += 1
|
||
continue
|
||
if buf[i + 1] in {0xC0, 0xC2}: # SOF0 / SOF2: height then width
|
||
h, w = struct.unpack(">HH", buf[i + 5: i + 9])
|
||
return _wh(w, h)
|
||
if i + 3 < len(buf):
|
||
i += 2 + struct.unpack(">H", buf[i + 2: i + 4])[0]
|
||
else:
|
||
break
|
||
return None
|
||
|
||
|
||
def _parse_gif_size(buf: bytes) -> Optional[dict[str, int]]:
|
||
if len(buf) < 10 or buf[:6].decode("ascii", errors="replace") not in {"GIF87a", "GIF89a"}:
|
||
return None
|
||
return _wh(*struct.unpack("<HH", buf[6:10]))
|
||
|
||
|
||
def _parse_webp_size(buf: bytes) -> Optional[dict[str, int]]:
|
||
if len(buf) < 16 or buf[:4] != b"RIFF" or buf[8:12] != b"WEBP":
|
||
return None
|
||
chunk = buf[12:16].decode("ascii", errors="replace")
|
||
if chunk == "VP8 ":
|
||
if len(buf) >= 30 and buf[23:26] == b"\x9d\x01\x2a":
|
||
w, h = struct.unpack("<HH", buf[26:30])
|
||
return _wh(w & 0x3FFF, h & 0x3FFF)
|
||
elif chunk == "VP8L":
|
||
if len(buf) >= 25 and buf[20] == 0x2F:
|
||
bits = struct.unpack("<I", buf[21:25])[0]
|
||
return _wh((bits & 0x3FFF) + 1, ((bits >> 14) & 0x3FFF) + 1)
|
||
elif chunk == "VP8X" and len(buf) >= 30:
|
||
return _wh(int.from_bytes(buf[24:27], "little") + 1, int.from_bytes(buf[27:30], "little") + 1)
|
||
return None
|
||
|
||
|
||
# ============ URL 下载 ============
|
||
|
||
async def download_url(url: str, max_size_mb: int = DEFAULT_MAX_SIZE_MB) -> tuple[bytes, str]:
|
||
"""下载 URL 内容，返回 (bytes, content_type)。
|
||
|
||
Raises:
|
||
ValueError: 内容超过 max_size_mb，或 URL / 重定向目标被 SSRF 防护拦截
|
||
httpx.HTTPError: 网络/HTTP 错误
|
||
"""
|
||
# SSRF protection: yuanbao downloads model-supplied and inbound URLs server-side.
|
||
# Reject private/internal targets up front, and re-validate every redirect hop so a
|
||
# public URL can't 302 to http://169.254.169.254/.
|
||
from tools.url_safety import create_ssrf_safe_async_client, is_safe_url
|
||
|
||
if not is_safe_url(url):
|
||
raise ValueError(f"Blocked unsafe URL (SSRF protection): {url}")
|
||
|
||
async def _redirect_guard(response: httpx.Response) -> None:
|
||
if response.is_redirect and response.next_request:
|
||
redirect_url = str(response.next_request.url)
|
||
if not is_safe_url(redirect_url):
|
||
raise ValueError(f"Blocked redirect to private/internal address: {redirect_url}")
|
||
|
||
max_bytes = max_size_mb * 1024 * 1024
|
||
async with create_ssrf_safe_async_client(
|
||
timeout=30.0, follow_redirects=True, event_hooks={"response": [_redirect_guard]},
|
||
) as client:
|
||
# HEAD 预检大小；部分服务器不支持 HEAD，忽略
|
||
try:
|
||
head = await client.head(url)
|
||
content_length = int(head.headers.get("content-length", 0) or 0)
|
||
if content_length > 0 and content_length > max_bytes:
|
||
raise ValueError(f"文件过大: {content_length / 1024 / 1024:.1f} MB > {max_size_mb} MB")
|
||
except httpx.HTTPStatusError:
|
||
pass
|
||
|
||
# 流式 GET，防止超限
|
||
async with client.stream("GET", url) as resp:
|
||
resp.raise_for_status()
|
||
content_type = resp.headers.get("content-type", "").split(";")[0].strip()
|
||
chunks: list[bytes] = []
|
||
downloaded = 0
|
||
async for chunk in resp.aiter_bytes(65536):
|
||
downloaded += len(chunk)
|
||
if downloaded > max_bytes:
|
||
raise ValueError(f"文件过大: 已超过 {max_size_mb} MB 限制")
|
||
chunks.append(chunk)
|
||
return b"".join(chunks), content_type
|
||
|
||
|
||
# ============ COS 鉴权（HMAC-SHA1） ============
|
||
|
||
def _hmac_sha1_hex(key: str, msg: str) -> str:
|
||
return hmac.new(key.encode("utf-8"), msg.encode("utf-8"), hashlib.sha1).hexdigest()
|
||
|
||
|
||
def _sorted_kv(d: dict[str, str]) -> list[tuple[str, str]]:
|
||
"""签名用 (小写 key, URL-encoded value) 按字典序排列"""
|
||
return sorted((k.lower(), urllib.parse.quote(str(v), safe="")) for k, v in d.items())
|
||
|
||
|
||
def _cos_sign(
|
||
method: str,
|
||
path: str,
|
||
params: dict[str, str],
|
||
headers: dict[str, str],
|
||
secret_id: str,
|
||
secret_key: str,
|
||
start_time: Optional[int] = None,
|
||
expire_seconds: int = 3600,
|
||
) -> str:
|
||
"""构建 COS 请求签名（q-sign-algorithm=sha1；https://cloud.tencent.com/document/product/436/7778）。
|
||
|
||
Args:
|
||
method / path: HTTP 方法（小写）与 URL-encoded 路径
|
||
params / headers: 参与签名的查询参数与请求头（key 会被小写）
|
||
secret_id / secret_key: 临时 tmpSecretId / tmpSecretKey
|
||
start_time: 签名起始 Unix 时间戳（默认 now）；expire_seconds: 有效期（秒）
|
||
|
||
Returns:
|
||
Authorization header 值
|
||
"""
|
||
now = int(time.time())
|
||
q_sign_time = f"{start_time or now};{(start_time or now) + expire_seconds}"
|
||
# SignKey = HMAC-SHA1(SecretKey, q-sign-time)
|
||
sign_key = _hmac_sha1_hex(secret_key, q_sign_time)
|
||
sorted_params = _sorted_kv(params)
|
||
sorted_headers = _sorted_kv(headers)
|
||
http_string = "\n".join([
|
||
method.lower(),
|
||
path,
|
||
"&".join(f"{k}={v}" for k, v in sorted_params),
|
||
"&".join(f"{k}={v}" for k, v in sorted_headers),
|
||
"",
|
||
])
|
||
string_to_sign = "\n".join(["sha1", q_sign_time, hashlib.sha1(http_string.encode("utf-8")).hexdigest(), ""])
|
||
signature = _hmac_sha1_hex(sign_key, string_to_sign)
|
||
return (
|
||
f"q-sign-algorithm=sha1"
|
||
f"&q-ak={secret_id}"
|
||
f"&q-sign-time={q_sign_time}"
|
||
f"&q-key-time={q_sign_time}"
|
||
f"&q-header-list={';'.join(k for k, _ in sorted_headers)}"
|
||
f"&q-url-param-list={';'.join(k for k, _ in sorted_params)}"
|
||
f"&q-signature={signature}"
|
||
)
|
||
|
||
|
||
# ============ 主要公开 API ============
|
||
|
||
async def get_cos_credentials(
|
||
app_key: str,
|
||
api_domain: str,
|
||
token: str,
|
||
filename: str = "file",
|
||
file_id: Optional[str] = None,
|
||
bot_id: str = "",
|
||
route_env: str = "",
|
||
) -> dict:
|
||
"""调用 genUploadInfo 获取 COS 临时密钥及上传配置。
|
||
|
||
Args:
|
||
app_key / bot_id: X-ID 头（bot_id 优先）; api_domain: 如 https://bot.yuanbao.tencent.com
|
||
token: 当前有效签票 token（X-Token 头）; file_id: 不传则自动生成
|
||
|
||
Returns:
|
||
接口 data dict：bucketName, region, location (COS key), encryptTmpSecretId,
|
||
encryptTmpSecretKey, encryptToken, startTime, expiredTime, resourceUrl, resourceID(可选)
|
||
|
||
Raises:
|
||
RuntimeError: 接口返回非 0 code 或 bucketName/location 缺失
|
||
"""
|
||
if file_id is None:
|
||
file_id = generate_file_id()
|
||
upload_url = f"{api_domain.rstrip('/')}{UPLOAD_INFO_PATH}"
|
||
headers = {"Content-Type": "application/json", "X-Token": token, "X-ID": bot_id or app_key, "X-Source": "web"}
|
||
if route_env:
|
||
headers["X-Route-Env"] = route_env
|
||
body = {"fileName": filename, "fileId": file_id, "docFrom": "localDoc", "docOpenId": ""}
|
||
|
||
async with httpx.AsyncClient(timeout=15.0) as client:
|
||
resp = await client.post(upload_url, json=body, headers=headers)
|
||
resp.raise_for_status()
|
||
result: dict[str, Any] = resp.json()
|
||
|
||
code = result.get("code")
|
||
if code != 0 and code is not None:
|
||
raise RuntimeError(f"genUploadInfo 失败: code={code}, msg={result.get('msg', '')}")
|
||
data = result.get("data") or result
|
||
missing = [f for f in ["bucketName", "location"] if not data.get(f)]
|
||
if missing:
|
||
raise RuntimeError(f"genUploadInfo 返回字段不完整: 缺少字段 {missing}")
|
||
return data
|
||
|
||
|
||
async def upload_to_cos(
|
||
file_bytes: bytes,
|
||
filename: str,
|
||
content_type: str,
|
||
credentials: dict,
|
||
bucket: str,
|
||
region: str,
|
||
) -> dict:
|
||
"""用临时凭证 HMAC-SHA1 签名，httpx PUT 上传到 COS。
|
||
|
||
Args:
|
||
credentials: get_cos_credentials() 返回的 dict（encryptTmpSecretId/Key, encryptToken,
|
||
location, resourceUrl, startTime, expiredTime）
|
||
bucket / region: 如 chatbot-1234567890 / ap-guangzhou
|
||
|
||
Returns:
|
||
{url, uuid (内容 MD5), size, width?, height? (仅图片)}
|
||
|
||
Raises:
|
||
httpx.HTTPStatusError: COS 返回非 2xx
|
||
RuntimeError: credentials 字段缺失
|
||
"""
|
||
secret_id: str = credentials.get("encryptTmpSecretId", "")
|
||
secret_key: str = credentials.get("encryptTmpSecretKey", "")
|
||
session_token: str = credentials.get("encryptToken", "")
|
||
cos_key: str = credentials.get("location", "")
|
||
resource_url: str = credentials.get("resourceUrl", "")
|
||
start_time: Optional[int] = credentials.get("startTime")
|
||
expired_time: Optional[int] = credentials.get("expiredTime")
|
||
|
||
if not secret_id or not secret_key or not cos_key:
|
||
raise RuntimeError(
|
||
f"COS credentials 不完整: secretId={bool(secret_id)}, "
|
||
f"secretKey={bool(secret_key)}, location={bool(cos_key)}"
|
||
)
|
||
|
||
cos_host = f"{bucket}.cos.accelerate.myqcloud.com" if COS_USE_ACCELERATE else f"{bucket}.cos.{region}.myqcloud.com"
|
||
encoded_key = urllib.parse.quote(cos_key, safe="/").lstrip("/")
|
||
cos_url = f"https://{cos_host}/{encoded_key}"
|
||
|
||
if not content_type or content_type == "application/octet-stream":
|
||
content_type = guess_mime_type(filename) if is_image(filename) else "application/octet-stream"
|
||
|
||
file_uuid = md5_hex(file_bytes)
|
||
file_size = len(file_bytes)
|
||
|
||
now = int(time.time())
|
||
authorization = _cos_sign(
|
||
method="put",
|
||
path=f"/{encoded_key}",
|
||
params={},
|
||
headers={"host": cos_host, "content-type": content_type, "x-cos-security-token": session_token},
|
||
secret_id=secret_id,
|
||
secret_key=secret_key,
|
||
start_time=start_time if start_time else now,
|
||
expire_seconds=(expired_time - now) if expired_time and expired_time > now else 3600,
|
||
)
|
||
put_headers = {"Authorization": authorization, "Content-Type": content_type, "x-cos-security-token": session_token}
|
||
|
||
logger.info("COS PUT: bucket=%s region=%s key=%s size=%d mime=%s", bucket, region, cos_key, file_size, content_type)
|
||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||
resp = await client.put(cos_url, content=file_bytes, headers=put_headers)
|
||
resp.raise_for_status()
|
||
|
||
result: dict[str, Any] = {"url": resource_url or cos_url, "uuid": file_uuid, "size": file_size}
|
||
if content_type.startswith("image/"):
|
||
size_info = parse_image_size(file_bytes)
|
||
if size_info:
|
||
result["width"] = size_info["width"]
|
||
result["height"] = size_info["height"]
|
||
logger.info("COS 上传成功: url=%s size=%d", result["url"], file_size)
|
||
return result
|
||
|
||
|
||
# ============ TIM 媒体消息构建（https://cloud.tencent.com/document/product/269/2720） ============
|
||
|
||
def build_image_msg_body(
|
||
url: str,
|
||
uuid: Optional[str] = None,
|
||
filename: Optional[str] = None,
|
||
size: int = 0,
|
||
width: int = 0,
|
||
height: int = 0,
|
||
mime_type: str = "",
|
||
) -> list[dict]:
|
||
"""构建 TIMImageElem 消息体（可直接放入 msg_body）。uuid 缺省依次退到 filename / URL basename / "image"。"""
|
||
return [{
|
||
"msg_type": "TIMImageElem",
|
||
"msg_content": {
|
||
"uuid": uuid or filename or _basename_from_url(url) or "image",
|
||
"image_format": get_image_format(mime_type) if mime_type else 255,
|
||
"image_info_array": [
|
||
{"type": 1, "size": size, "width": width, "height": height, "url": url}, # type 1 = 原图
|
||
],
|
||
},
|
||
}]
|
||
|
||
|
||
def build_file_msg_body(url: str, filename: str, uuid: Optional[str] = None, size: int = 0) -> list[dict]:
|
||
"""构建 TIMFileElem 消息体（可直接放入 msg_body）。uuid 缺省使用 filename。"""
|
||
return [{
|
||
"msg_type": "TIMFileElem",
|
||
"msg_content": {"uuid": uuid or filename, "file_name": filename, "file_size": size, "url": url},
|
||
}]
|
||
|
||
|
||
def _basename_from_url(url: str) -> str:
|
||
try:
|
||
return os.path.basename(urllib.parse.urlparse(url).path)
|
||
except Exception:
|
||
return ""
|