import asyncio import json import time import logging from collections import deque from typing import Set, Dict from fastapi import FastAPI, WebSocket, WebSocketDisconnect # ==================== 配置日志 ==================== logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', handlers=[ logging.StreamHandler(), logging.FileHandler('danmaku.log', encoding='utf-8') ] ) logger = logging.getLogger("danmaku") MSG_NO_LOGIN = "未登录，无法加入弹幕" MSG_BAD_TOKEN = "登录已失效，请刷新页面重试" NICK_GUEST = "观众" # ==================== 弹幕登录校验(live token -> live-api) ==================== import os as _os VERIFY_URL = 'https://api.rrrosmontis.icu/v1/live/api/auth/verify' try: _cfg = json.load(open(_os.path.join(_os.path.dirname(_os.path.abspath(__file__)), 'danmaku_auth.json'), encoding='utf-8')) VERIFY_URL = _cfg.get('verify_url', VERIFY_URL) except Exception: pass def _verify_token_sync(token): try: import urllib.request as _ur req = _ur.Request(VERIFY_URL, data=json.dumps({'token': token}).encode('utf-8'), headers={'Content-Type': 'application/json'}) with _ur.urlopen(req, timeout=6) as resp: d = json.loads(resp.read().decode('utf-8')) return d.get('user') if d.get('ok') else None except Exception: return None async def _verify_token(token): return await asyncio.to_thread(_verify_token_sync, token) # ==================== 直播间类 ==================== class LiveRoom: def __init__(self, room_id: int): self.room_id = room_id self.start_time = time.time() self.timeline: list = [] self.timeline_lock = asyncio.Lock() self.recent_danmaku = deque(maxlen=20) self.connections: Set[WebSocket] = set() self.connections_lock = asyncio.Lock() logger.info(f"创建新房间 {room_id}，开始时间: {self.start_time}") # ==================== 全局房间管理 ==================== rooms: Dict[int, LiveRoom] = {} rooms_lock = asyncio.Lock() async def get_or_create_room(room_id: int) -> LiveRoom: async with rooms_lock: if room_id not in rooms: rooms[room_id] = LiveRoom(room_id) return rooms[room_id] async def remove_connection_from_room(websocket: WebSocket): async with rooms_lock: for room in rooms.values(): async with room.connections_lock: if websocket in room.connections: room.connections.remove(websocket) logger.info(f"从房间 {room.room_id} 移除连接，当前连接数: {len(room.connections)}") break # ==================== FastAPI 应用 ==================== app = FastAPI() @app.websocket("/danmaku") async def websocket_endpoint(websocket: WebSocket): await websocket.accept() client_host = websocket.client.host if websocket.client else "unknown" token = (websocket.query_params.get('token') or '').strip() logger.info(f"新WebSocket连接: {client_host}") current_room: LiveRoom = None current_user: str = None try: async for message in websocket.iter_text(): logger.info(f"收到消息 from {client_host}: {message[:200]}") data = json.loads(message) msg_type = data.get("type") if msg_type == "join": # 房间号处理 room_id = data.get("room") try: room_id = int(room_id) except (TypeError, ValueError): await websocket.send_json({"type": "error", "message": "房间号必须是整数"}) continue # 如果已经加入过房间，先从原房间移除 if current_room: async with current_room.connections_lock: current_room.connections.discard(websocket) # 加入新房间 # 登录校验: 弹幕连接必须携带有效 live token (?token=) if not token: await websocket.send_json({"type": "error", "message": MSG_NO_LOGIN}) await websocket.close(code=4401) return user_info = await _verify_token(token) if not user_info: await websocket.send_json({"type": "error", "message": MSG_BAD_TOKEN}) await websocket.close(code=4401) return current_room = await get_or_create_room(room_id) current_user = user_info.get("nickname") or NICK_GUEST async with current_room.connections_lock: current_room.connections.add(websocket) logger.info(f"客户端 {client_host} 加入房间 {room_id}，用户名 {current_user}") # 发送 init 消息（历史弹幕等） elapsed = time.time() - current_room.start_time init_msg = { "type": "init", "start_time": current_room.start_time, "elapsed": elapsed, "recent_danmaku": list(current_room.recent_danmaku) } await websocket.send_json(init_msg) continue elif msg_type == "danmaku": if current_room is None: await websocket.send_json({"type": "error", "message": "请先发送 join 消息加入房间"}) continue text = data.get("text", "") user = current_user or NICK_GUEST danmaku = { "text": text, "user": user, "time": time.time() } # 记录到房间的最近弹幕和 timeline current_room.recent_danmaku.append(danmaku) async with current_room.timeline_lock: current_room.timeline.append({ "time": time.time() - current_room.start_time, "type": "danmaku", "user": danmaku["user"], "text": danmaku["text"] }) logger.info(f"弹幕 (房间 {current_room.room_id}): {danmaku['user']}: {danmaku['text']}") # 广播给同一房间的所有客户端 await broadcast_danmaku(current_room, danmaku) continue else: logger.warning(f"未知消息类型: {msg_type}") await websocket.send_json({"type": "error", "message": f"未知消息类型: {msg_type}"}) except WebSocketDisconnect: logger.info(f"WebSocket断开: {client_host}") except Exception as e: logger.error(f"WebSocket处理异常: {e}", exc_info=True) finally: if current_room: async with current_room.connections_lock: current_room.connections.discard(websocket) logger.info(f"从房间 {current_room.room_id} 移除连接，当前连接数: {len(current_room.connections)}") else: await remove_connection_from_room(websocket) logger.info(f"当前所有房间连接状态: {[(rid, len(room.connections)) for rid, room in rooms.items()]}") async def broadcast_danmaku(room: LiveRoom, danmaku: dict): async with room.connections_lock: for conn in room.connections: try: await conn.send_json({"type": "danmaku", "data": danmaku}) except: pass logger.info(f"广播弹幕到房间 {room.room_id}: {danmaku['user']}: {danmaku['text']}") @app.get("/") async def root(): return { "status": "running", "rooms": [ { "room_id": rid, "connections": len(room.connections), "timeline_count": len(room.timeline), "start_time": room.start_time } for rid, room in rooms.items() ] } # ==================== 启动服务 ==================== if __name__ == "__main__": import uvicorn logger.info("启动 uvicorn 服务器，端口 55556") uvicorn.run(app, host="0.0.0.0", port=55556)