WechatOnCloud/bridge/woc_bridge/routes/media.py

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)