MusicBot/musicbot/proxy.py

"""文件代理：bot 自建 HTTP 服务，把播放直链换成自己的公开 URL。

QQ 服务器侧上传富媒体（语音 file_type=3 / 群文件 4）要现场抓取 URL 入库，
两类直链它抓不到：tx 等源的 CDN 风控直链（QQ 侧报 40093007，bot 自己的
出口反而常能抓），以及 qs 登录态直链（MusicServer 本地解密流代理，只有
bot 所在网络可达）。开启后 send_song 发出的 URL（上传与直链文本同源）一律
换成 <PUBLIC_URL>/file/<key>/<文件名>：

- key 在发送时登记（歌曲 + 搜索词），随机不可猜、短时有效、总量封顶；
- 请求到来时经 resolve 现场重解析直链——直链本身有时效，登记时探测过的
  那条不能复用；解析自带候选源回退（tx 风控会自动落到 kw 等源）——再从
  上游流式转发字节；
- bot 主机需能被 QQ 服务器 / 群友访问（公网 IP 或端口映射），否则上传
  失败、回退直链文本；服务起不来（端口被占）时 claim 不换链，退回发
  原直链，主流程不受影响。"""

import secrets
import time
from collections import OrderedDict
from collections.abc import Awaitable, Callable
from typing import Any
from urllib.parse import quote

from aiohttp import ClientError, ClientSession, ClientTimeout, web
from loguru import logger

# 现场解析一条当前可抓取的直链（music._fetchable_play，含候选源回退），
# 失败返回 None。歌曲对象对本模块透明（避免与 music 循环导入），透传即可。
Resolver = Callable[[Any, str], Awaitable[tuple[Any, str] | None]]

_TTL = 2 * 60 * 60.0  # 秒：直链文本可能被群友过一阵才点，登记保留久一点
_MAX_ENTRIES = 256  # 条；超出按插入序淘汰最旧，封住内存
_CHUNK = 64 * 1024  # 流式转发的块大小
# 转发超时：连接要快、单次读不能卡死，但不设总时长（整首歌可以下挺久）
_STREAM_TIMEOUT = ClientTimeout(total=None, connect=10, sock_read=60)


class FileProxy:
    """发送时登记换链、请求时解析直链并转发字节的小型文件代理。"""

    def __init__(
        self,
        public_base: str,
        http: ClientSession,
        resolve: Resolver,
        port: int,
    ) -> None:
        self._base = public_base.rstrip("/")
        self._http = http
        self._resolve = resolve
        self._port = port
        self.port = port  # 实际监听端口（port=0 时在 start 后更新为分配值）
        # key -> (过期时刻, 歌曲, 搜索词, 文件名)；OrderedDict 按插入序淘汰
        self._entries: OrderedDict[str, tuple[float, Any, str, str]] = OrderedDict()
        self._runner: web.AppRunner | None = None

    @property
    def ready(self) -> bool:
        """HTTP 服务是否已在跑；没起来时调用方应放弃换链、退回原直链。"""
        return self._runner is not None

    def claim(self, song: Any, keyword: str, filename: str) -> str:
        """登记一首歌并返回它的公开代理 URL（尾段带文件名，供展示与扩展名）。"""
        now = time.monotonic()
        while self._entries:
            key, (expire_at, *_) = next(iter(self._entries.items()))
            if expire_at > now:
                break
            del self._entries[key]
        key = secrets.token_urlsafe(8)
        self._entries[key] = (now + _TTL, song, keyword, filename)
        while len(self._entries) > _MAX_ENTRIES:
            self._entries.popitem(last=False)
        return f"{self._base}/file/{key}/{quote(filename)}"

    async def start(self) -> None:
        """起 HTTP 服务；幂等（重连 READY 会再调），端口被占等失败仅记日志。"""
        if self._runner is not None:
            return
        app = web.Application()
        app.router.add_get("/file/{key}/{name}", self._serve)
        runner = web.AppRunner(app, access_log=None)  # 交付另有日志，不刷访问日志
        try:
            await runner.setup()
            site = web.TCPSite(runner, "0.0.0.0", self._port)
            await site.start()
        except OSError as e:
            logger.warning(f"文件代理监听 0.0.0.0:{self._port} 失败: {e}")
            await runner.cleanup()
            return
        if self._port == 0:
            # sockets 未列入 aiohttp 类型存根，只能从私有属性取实际端口
            server = site._server
            assert server is not None
            self.port = server.sockets[0].getsockname()[1]  # type: ignore[attr-defined]
        self._runner = runner
        logger.info(f"文件代理已监听 0.0.0.0:{self.port}，对外地址 {self._base}")

    async def stop(self) -> None:
        if self._runner is not None:
            await self._runner.cleanup()
            self._runner = None

    async def _serve(self, request: web.Request) -> web.StreamResponse:
        key = request.match_info["key"]
        entry = self._entries.get(key)
        if entry is None or entry[0] < time.monotonic():
            return web.Response(status=404, text="链接已过期或不存在")
        _, song, keyword, filename = entry
        try:
            found = await self._resolve(song, keyword)
        except Exception as e:  # 解析在 handler 边界，别让 500 栈漏给请求方
            logger.warning(f"文件代理解析直链异常: {e!r}")
            found = None
        if found is None:
            return web.Response(status=502, text="取流失败，请重新点歌")
        url = found[1]
        resp: web.StreamResponse | None = None
        try:
            async with self._http.get(url, timeout=_STREAM_TIMEOUT) as upstream:
                if upstream.status != 200:
                    raise ClientError(f"上游 HTTP {upstream.status}")
                resp = web.StreamResponse(headers={
                    # 上游的 Content-Type 透传（音源都是 audio/*），缺省按 mp3
                    "Content-Type": upstream.headers.get("Content-Type", "audio/mpeg"),
                    "Content-Disposition": (
                        f"attachment; filename*=UTF-8''{quote(filename)}"
                    ),
                })
                if upstream.content_length is not None:
                    resp.content_length = upstream.content_length
                await resp.prepare(request)
                async for chunk in upstream.content.iter_chunked(_CHUNK):
                    await resp.write(chunk)
                await resp.write_eof()
                logger.info(f"文件代理已交付：{filename}")
                return resp
        except (ClientError, TimeoutError, OSError) as e:
            logger.warning(f"文件代理转发失败（{url[:80]}）: {e!r}")
            # 字节已发出一半时状态码改不了，就地断流让客户端感知；没开始才给 502
            if resp is not None and resp.prepared:
                return resp
            return web.Response(status=502, text="上游取流失败")