132 lines
4.9 KiB
Python
132 lines
4.9 KiB
Python
"""OneBot 11 反向 WebSocket 服务端（aiohttp 实现）。"""
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
|
|
from aiohttp import WSMsgType, web
|
|
|
|
log = logging.getLogger("onebot")
|
|
|
|
|
|
class OneBotServer:
|
|
"""承载 OneBot 11 反向 WS 连接：鉴权、事件回调、API 调用。"""
|
|
|
|
def __init__(self, host: str, port: int, access_token: str = ""):
|
|
self.host = host
|
|
self.port = port
|
|
self.access_token = access_token
|
|
self.ws = None # 当前连接（单实现端场景足够）
|
|
self._echo_id = 0
|
|
self._api_waiters: dict[str, asyncio.Future] = {}
|
|
self._event_handler = None
|
|
self._runner: web.AppRunner | None = None
|
|
|
|
# ---------- 生命周期 ----------
|
|
|
|
async def start(self, event_handler):
|
|
"""启动 WS 服务端。event_handler 为 async 函数，接收事件 dict。"""
|
|
self._event_handler = event_handler
|
|
app = web.Application()
|
|
app.router.add_get("/ws", self._handle_ws)
|
|
self._runner = web.AppRunner(app)
|
|
await self._runner.setup()
|
|
site = web.TCPSite(self._runner, self.host, self.port)
|
|
await site.start()
|
|
log.info("反向WS服务端已启动 ws://%s:%d/ws", self.host, self.port)
|
|
|
|
async def stop(self):
|
|
if self._runner:
|
|
await self._runner.cleanup()
|
|
|
|
@property
|
|
def connected(self) -> bool:
|
|
return self.ws is not None and not self.ws.closed
|
|
|
|
# ---------- WebSocket 处理 ----------
|
|
|
|
async def _handle_ws(self, request: web.Request):
|
|
if self.access_token:
|
|
auth = request.headers.get("Authorization", "")
|
|
qs = request.query.get("access_token", "")
|
|
if auth != f"Bearer {self.access_token}" and qs != self.access_token:
|
|
log.warning("连接被拒绝：access_token 校验失败 (来自 %s)", request.remote)
|
|
return web.Response(status=401)
|
|
|
|
log.info("OneBot 实现端已连接: %s", request.remote)
|
|
ws = web.WebSocketResponse(heartbeat=30)
|
|
await ws.prepare(request)
|
|
self.ws = ws
|
|
try:
|
|
async for msg in ws:
|
|
if msg.type == WSMsgType.TEXT:
|
|
await self._dispatch(msg.data)
|
|
elif msg.type == WSMsgType.ERROR:
|
|
log.error("WS 错误: %s", ws.exception())
|
|
finally:
|
|
if self.ws is ws:
|
|
self.ws = None
|
|
log.warning("OneBot 实现端断开: %s", request.remote)
|
|
return ws
|
|
|
|
async def _dispatch(self, raw: str):
|
|
try:
|
|
data = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
log.warning("收到非 JSON 消息: %s", raw[:200])
|
|
return
|
|
|
|
echo = data.get("echo")
|
|
if echo is not None and echo in self._api_waiters:
|
|
fut = self._api_waiters.pop(echo)
|
|
if not fut.done():
|
|
fut.set_result(data.get("data"))
|
|
return
|
|
|
|
if "post_type" in data and self._event_handler:
|
|
# 用任务处理事件，避免阻塞消息循环导致 API 回执无法送达
|
|
asyncio.create_task(self._run_event_handler(data))
|
|
|
|
# ---------- API 调用 ----------
|
|
|
|
async def _run_event_handler(self, data: dict):
|
|
try:
|
|
await self._event_handler(data)
|
|
except Exception:
|
|
log.exception("处理事件时出错: %s", data)
|
|
|
|
async def call_api(self, action: str, **params) -> dict:
|
|
"""调用 OneBot API，返回 data 字段。"""
|
|
if not self.connected:
|
|
raise ConnectionError("OneBot 实现端未连接")
|
|
self._echo_id += 1
|
|
echo = f"echo-{self._echo_id}"
|
|
payload = {"action": action, "params": params, "echo": echo}
|
|
fut = asyncio.get_running_loop().create_future()
|
|
self._api_waiters[echo] = fut
|
|
await self.ws.send_str(json.dumps(payload, ensure_ascii=False))
|
|
try:
|
|
return await asyncio.wait_for(fut, timeout=15)
|
|
finally:
|
|
self._api_waiters.pop(echo, None)
|
|
|
|
async def send_group_msg(self, group_id: int, message: str):
|
|
try:
|
|
await self.call_api("send_group_msg", group_id=group_id, message=message)
|
|
log.info("已发送群消息 -> 群 %s", group_id)
|
|
except Exception as e:
|
|
log.error("发送群消息失败 (群 %s): %s", group_id, e)
|
|
|
|
async def send_private_msg(self, user_id: int, message: str):
|
|
try:
|
|
await self.call_api("send_private_msg", user_id=user_id, message=message)
|
|
log.info("已发送私聊消息 -> %s", user_id)
|
|
except Exception as e:
|
|
log.error("发送私聊消息失败 (%s): %s", user_id, e)
|
|
|
|
async def broadcast(self, message: str, groups: list, friends: list):
|
|
"""向所有配置的群和好友发送消息。"""
|
|
for gid in groups:
|
|
await self.send_group_msg(gid, message)
|
|
for uid in friends:
|
|
await self.send_private_msg(uid, message)
|