ayaka.driver.ayakabot_driver.driver

  1from importlib import import_module
  2import json
  3import asyncio
  4import uvicorn
  5from loguru import logger
  6from fastapi import FastAPI, WebSocket
  7from pathlib import Path
  8import re
  9
 10from .bot import Bot
 11from .event import json_to_event, MessageEvent
 12from .websocket import FastAPIWebSocket
 13from ...config import ayaka_root_config
 14
 15app = FastAPI()
 16
 17
 18class Config:
 19    port = ayaka_root_config.ayaka_port
 20    ayaka_prefix = "#"
 21    ayaka_separate = " "
 22    ayaka_exclude_old = True
 23    fastapi_reload = True
 24
 25
 26class Driver:
 27    connect_calls = []
 28    disconnect_calls = []
 29    deal_event = None
 30    config = Config()
 31
 32    def deal(self, bot, event):
 33        if self.deal_event:
 34            if isinstance(event, MessageEvent):
 35                logger.success(str(event))
 36                asyncio.create_task(self.deal_event(bot, event))
 37
 38    async def bot_connect(self, bot: Bot):
 39        ts = []
 40        for call in self.connect_calls:
 41            ts.append(asyncio.create_task(call(bot)))
 42        await asyncio.gather(*ts)
 43
 44    async def bot_disconnect(self, bot: Bot):
 45        ts = []
 46        for call in self.disconnect_calls:
 47            ts.append(asyncio.create_task(call(bot)))
 48        await asyncio.gather(*ts)
 49
 50    def on_bot_connect(self, func):
 51        self.connect_calls.append(func)
 52
 53    def on_bot_disconnect(self, func):
 54        self.disconnect_calls.append(func)
 55
 56    def on_startup(self, func):
 57        app.on_event("startup")(func)
 58
 59    def on_shutdown(self, func):
 60        app.on_event("shutdown")(func)
 61
 62
 63driver = Driver()
 64
 65
 66# 启动服务
 67def run(host="127.0.0.1", port=driver.config.port, reload=True):
 68    tip = f'''
 69您目前正在使用AYAKA BOT！
 70注意设置cqhttp连接配置为
 71- 反向ws
 72- ws://127.0.0.1:{port}/onebot/v11/ws
 73'''
 74    print(tip)
 75    uvicorn.run(
 76        app=f"{__name__}:app",
 77        host=host,
 78        port=port,
 79        reload=reload,
 80    )
 81
 82
 83def load_plugins(path):
 84    path = Path(path)
 85    for p in path.iterdir():
 86        if p.name.startswith("_"):
 87            continue
 88
 89        load_plugin(p)
 90
 91
 92def load_plugin(path):
 93    if isinstance(path, Path):
 94        p = path
 95    else:
 96        p = Path(path)
 97
 98    name = re.sub(r"\\|/", ".", str(p))
 99    name = re.sub(r".py$", "", name)
100    try:
101        import_module(name)
102        logger.opt(colors=True).success(f"导入成功 \"<y>{p.stem}</y>\"")
103    except:
104        logger.opt(colors=True).exception(f"导入失败 \"<y>{p.stem}</y>\"")
105
106
107@app.websocket("/onebot/v11/ws")
108async def endpoint(websocket: WebSocket):
109    self_id = websocket.headers.get("x-self-id")
110    ws = FastAPIWebSocket(websocket)
111    bot = Bot(ws, self_id)
112
113    # 建立ws连接
114    await ws.accept()
115    await driver.bot_connect(bot)
116
117    try:
118        # 监听循环
119        while True:
120            data = await ws.receive()
121            json_data = json.loads(data)
122            # 将json解析为对应的event
123            event = json_to_event(json_data)
124
125            if not event:
126                continue
127
128            driver.deal(bot, event)
129
130    except:
131        logger.exception("连接中断")
132    finally:
133        # 结束ws连接
134        await driver.bot_disconnect(bot)
135        await ws.close()
136
137
138def get_driver():
139    return driver
140
141
142def on_message(priority, block, handlers):
143    driver.deal_event = handlers[0]