#!/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)