mirror of
https://github.com/nonebot/nonebot2.git
synced 2025-06-06 20:35:45 +00:00
114 lines
3.5 KiB
Python
114 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
import json
|
|
import logging
|
|
from typing import Optional
|
|
from ipaddress import IPv4Address
|
|
|
|
import uvicorn
|
|
from starlette.websockets import WebSocketDisconnect
|
|
from fastapi import Body, FastAPI, WebSocket
|
|
|
|
from nonebot.log import logger
|
|
from nonebot.config import Config
|
|
from nonebot.drivers import BaseDriver
|
|
from nonebot.adapters.coolq import Bot as CoolQBot
|
|
|
|
|
|
class Driver(BaseDriver):
|
|
|
|
def __init__(self, config: Config):
|
|
self._server_app = FastAPI(
|
|
debug=config.debug,
|
|
openapi_url=None,
|
|
docs_url=None,
|
|
redoc_url=None,
|
|
)
|
|
|
|
self.config = config
|
|
|
|
self._server_app.post("/{adapter}/")(self._handle_http)
|
|
self._server_app.post("/{adapter}/http")(self._handle_http)
|
|
self._server_app.websocket("/{adapter}/ws")(self._handle_ws_reverse)
|
|
self._server_app.websocket("/{adapter}/ws/")(self._handle_ws_reverse)
|
|
|
|
@property
|
|
def server_app(self):
|
|
return self._server_app
|
|
|
|
@property
|
|
def asgi(self):
|
|
return self._server_app
|
|
|
|
@property
|
|
def logger(self):
|
|
return logging.getLogger("fastapi")
|
|
|
|
def run(self,
|
|
host: Optional[IPv4Address] = None,
|
|
port: Optional[int] = None,
|
|
*,
|
|
app: Optional[str] = None,
|
|
**kwargs):
|
|
LOGGING_CONFIG = {
|
|
"version": 1,
|
|
"disable_existing_loggers": False,
|
|
"formatters": {
|
|
"default": {
|
|
"()": "logging.Formatter",
|
|
"fmt": "[%(asctime)s %(name)s] %(levelname)s: %(message)s",
|
|
},
|
|
},
|
|
"handlers": {
|
|
"default": {
|
|
"formatter": "default",
|
|
"class": "logging.StreamHandler",
|
|
"stream": "ext://sys.stdout",
|
|
},
|
|
},
|
|
"loggers": {
|
|
"uvicorn.error": {
|
|
"handlers": ["default"],
|
|
"level": "INFO"
|
|
},
|
|
"uvicorn.access": {
|
|
"handlers": ["default"],
|
|
"level": "INFO",
|
|
},
|
|
},
|
|
}
|
|
uvicorn.run(app or self.server_app,
|
|
host=host or str(self.config.host),
|
|
port=port or self.config.port,
|
|
reload=app and self.config.debug,
|
|
debug=self.config.debug,
|
|
log_config=LOGGING_CONFIG,
|
|
**kwargs)
|
|
|
|
async def _handle_http(self, adapter: str, data: dict = Body(...)):
|
|
# TODO: Check authorization
|
|
logger.debug(f"Received message: {data}")
|
|
if adapter == "coolq":
|
|
bot = CoolQBot("http", self.config)
|
|
await bot.handle_message(data)
|
|
return {"status": 200, "message": "success"}
|
|
|
|
async def _handle_ws_reverse(self, adapter: str, websocket: WebSocket):
|
|
# TODO: Check authorization
|
|
await websocket.accept()
|
|
while True:
|
|
try:
|
|
data = await websocket.receive_json()
|
|
except json.decoder.JSONDecodeError as e:
|
|
logger.exception(e)
|
|
continue
|
|
except WebSocketDisconnect:
|
|
logger.error("WebSocket Disconnect")
|
|
return
|
|
|
|
logger.debug(f"Received message: {data}")
|
|
if adapter == "coolq":
|
|
bot = CoolQBot("websocket", self.config, websocket=websocket)
|
|
await bot.handle_message(data)
|