from __future__ import annotations import asyncio import logging import mimetypes import urllib.error import urllib.parse import urllib.request from typing import Optional from fastapi import APIRouter from fastapi.responses import Response from woc_bridge.config import _require_db_reader from woc_bridge.db.coordinator import _check_db_readable, with_db_retry from woc_bridge.models import BridgeError logger = logging.getLogger("woc-bridge") router = APIRouter() # 头像下载安全限制 _AVATAR_MAX_BYTES = 10 * 1024 * 1024 # 10MB 上限,头像通常 < 1MB _AVATAR_ALLOWED_HOSTS = ( "wx.qlogo.cn", "thirdwx.qlogo.cn", "wxhead.clouddn.com", "thirdqq.qlogo.cn", "q.qlogo.cn", ) # --------------------------------------------------------------------------- # 路由:GET /api/media/avatar/{wxid} # --------------------------------------------------------------------------- # 注意:此路由必须在 /api/media/{msg_id} 之前声明,否则会被 {msg_id} # 匹配到(FastAPI 按声明顺序匹配,先匹配先得)。 @router.get("/api/media/avatar/{wxid}") @with_db_retry async def get_avatar(wxid: str) -> Response: """获取联系人头像。 最小实现:优先从 contact.db 读取 big_head_url / small_head_url / avatar_url, 若存在 HTTP(S) 链接则拉取并透传二进制;否则返回 MEDIA_NOT_FOUND。 本地 .dat 头像缓存解密仍留给后续 W-04 完善。 Args: wxid: 路径参数,联系人 wxid Returns: Response:头像二进制内容 + Content-Type Raises: BridgeError(MEDIA_NOT_FOUND): 未找到联系人或无有效头像 URL(HTTP 404) BridgeError(DB_ENCRYPTED): DB 已加密,无法读取联系人表(HTTP 503) Notes: - 路由声明顺序敏感:必须在 /api/media/{msg_id} 之前,否则 FastAPI 会把 "avatar" 当作 msg_id 匹配 """ db_reader = _require_db_reader() # DB 不可读时统一返回 503,行为与 /api/media/{msg_id} 一致 await _check_db_readable() contact = await asyncio.to_thread(db_reader.get_contact_detail, wxid) url: Optional[str] = None if contact: url = ( contact.get("big_head_url") or contact.get("small_head_url") or contact.get("avatar_url") ) if not url or not url.startswith(("http://", "https://")): raise BridgeError( code="MEDIA_NOT_FOUND", message=f"未找到联系人头像 URL: {wxid}", ) # 下载前校验 URL 安全性(SSRF 防护:拒绝内网/非白名单域名) try: await asyncio.to_thread(_validate_avatar_url, url) except ValueError as e: logger.warning("avatar: wxid=%s url=%s 校验失败: %s", wxid, url, e) raise BridgeError( code="MEDIA_NOT_FOUND", message=f"头像 URL 校验失败: {wxid}", ) try: data = await asyncio.to_thread(_fetch_avatar_url, url) except Exception as e: logger.warning("avatar: wxid=%s url=%s 拉取失败: %s", wxid, url, e) raise BridgeError( code="MEDIA_NOT_FOUND", message=f"头像下载失败: {wxid}", ) mime = mimetypes.guess_type(url)[0] or "image/jpeg" logger.info("avatar: wxid=%s url=%s → %d bytes (%s)", wxid, url, len(data), mime) return Response(content=data, media_type=mime) def _validate_avatar_url(url: str) -> None: """校验头像 URL 安全性:仅允许已知头像 CDN 域名,拒绝内网地址。""" parsed = urllib.parse.urlparse(url) host = parsed.hostname or "" if not host: raise ValueError("头像 URL 缺少 hostname") # 白名单域名检查 if host not in _AVATAR_ALLOWED_HOSTS: raise ValueError(f"头像 URL host 不在白名单内: {host}") def _fetch_avatar_url(url: str) -> bytes: """同步拉取头像 URL 二进制内容(限制 10MB 防止 OOM)。""" with urllib.request.urlopen(url, timeout=15) as resp: # 读取 Content-Length 预检 content_length = resp.headers.get("Content-Length") if content_length and int(content_length) > _AVATAR_MAX_BYTES: raise ValueError( f"头像文件过大 ({content_length} bytes),超过 {_AVATAR_MAX_BYTES} 上限" ) data = resp.read(_AVATAR_MAX_BYTES + 1) if len(data) > _AVATAR_MAX_BYTES: raise ValueError(f"头像文件过大,超过 {_AVATAR_MAX_BYTES} 上限") return data # --------------------------------------------------------------------------- # 路由:GET /api/media/{msg_id} # --------------------------------------------------------------------------- @router.get("/api/media/{msg_id}") @with_db_retry async def get_media(msg_id: str) -> Response: """下载消息媒体文件(图片/语音/视频/文件)。 Args: msg_id: 路径参数,消息 ID(来自 /api/messages/since 返回的 msg_id) Returns: Response:二进制内容 + Content-Type(image/jpeg / audio/amr / video/mp4 / application/octet-stream 等) Raises: BridgeError(MEDIA_NOT_FOUND): 媒体文件未找到或解析失败(HTTP 404) BridgeError(DB_ENCRYPTED): DB 已加密,无法查 media 表(HTTP 503) BridgeError(DB_NOT_FOUND): 未找到微信消息 DB(HTTP 500) Notes: - 前置调 _check_db_readable(),DB 加密时直接返回 503 DB_ENCRYPTED, 不再走 sqlite3.DatabaseError 被兜底为 500 的旧路径 - DbReader.get_media_bytes 内部:DB 查 msg_id → 找对应 .dat 文件 → XOR 解密 → 返回 (data, mime) - 大文件可能阻塞较久,建议调用方设置合理的 timeout """ db_reader = _require_db_reader() # 前置 DB 可达性检查:加密时直接返回 503,避免 sqlite3 异常被兜底 500 await _check_db_readable() result = await asyncio.to_thread(db_reader.get_media_bytes, msg_id) if result is None: logger.info("media: msg_id=%s → 未找到", msg_id) raise BridgeError( code="MEDIA_NOT_FOUND", message=f"未找到媒体文件: {msg_id}", ) data, mime = result logger.info("media: msg_id=%s → %d bytes (%s)", msg_id, len(data), mime) return Response(content=data, media_type=mime)