import contextlib from asyncio import sleep from typing import Optional from aiohttp import ClientConnectorError from mipa.exception import WebSocketNotConnected from mipa.ext import commands from mipa.router import Router from mipac import ( Note, NotificationFollow, NotificationFollowRequest, ChatMessage, NotificationAchievement, ) from mipac.client import Client as MisskeyClient from defs.chat import send_chat_message from defs.misskey import send_update from defs.notice import ( send_user_followed, send_follow_request, send_follow_request_accept, send_achievement_earned, ) from glover import misskey_url, misskey_host from models.models.user import User, TokenStatusEnum from models.services.user import UserAction from init import bot, logs, sqlite class MisskeyBot(commands.Bot): def __init__(self, user: User): super().__init__() self.user_id: int = user.user_id self.tg_user: User = user async def on_ready(self, ws): await Router(ws).connect_channel(["main", "home"]) logs.info(f"成功启动 Misskey Bot WS 任务 {self.user_id}") async def on_reconnect(self, ws): await Router(ws).connect_channel(["main", "home"]) async def on_note(self, note: Note): if self.tg_user: await send_update(self.tg_user.chat_id, note, self.tg_user.timeline_topic) async def on_user_followed(self, notice: NotificationFollow): if self.tg_user: await send_user_followed( self.tg_user.chat_id, notice, self.tg_user.notice_topic ) async def on_follow_request(self, notice: NotificationFollowRequest): if self.tg_user: await send_follow_request( self.tg_user.chat_id, notice, self.tg_user.notice_topic ) async def on_follow_request_accept(self, notice: NotificationFollowRequest): if self.tg_user: await send_follow_request_accept( self.tg_user.chat_id, notice, self.tg_user.notice_topic ) async def on_chat(self, message: ChatMessage): if self.tg_user: await send_chat_message( self.tg_user.chat_id, message, self.tg_user.notice_topic ) async def on_chat_unread_message(self, message: ChatMessage): await message.api.read() async def on_achievement_earned(self, notice: NotificationAchievement): if self.tg_user: await send_achievement_earned( self.tg_user.chat_id, notice, self.tg_user.notice_topic ) misskey_bot_map: dict[int, MisskeyBot] = {} def get_misskey_bot(user_id: int) -> Optional[MisskeyBot]: return None if user_id not in misskey_bot_map else misskey_bot_map[user_id] async def create_or_get_misskey_bot(user: User) -> MisskeyBot: if user.user_id not in misskey_bot_map: misskey_bot_map[user.user_id] = MisskeyBot(user) return misskey_bot_map[user.user_id] async def run(user: User): misskey = await create_or_get_misskey_bot(user) try: logs.info(f"尝试启动 Misskey Bot WS 任务 {user.user_id}") await misskey.start(misskey_url, user.token) except ClientConnectorError: await sleep(3) await run(user) async def test_token(token: str) -> bool: try: logs.info(f"验证 Token {token}") client = MisskeyClient(misskey_host, token) await client.http.login() await client.http.close_session() return True except Exception: return False async def rerun_misskey_bot(user_id: int) -> bool: if misskey := get_misskey_bot(user_id): with contextlib.suppress(WebSocketNotConnected): await misskey.disconnect() misskey_bot_map.pop(user_id) user = await UserAction.get_user_if_ok(user_id) if not user: return False if not await test_token(user.token): await UserAction.set_user_status(user_id, TokenStatusEnum.INVALID_TOKEN) return False bot.loop.create_task(run(user)) return True async def init_misskey_bot(): await sqlite.create_db_and_tables() count = 0 for user in await UserAction.get_all_token_ok_users(): if not await test_token(user.token): user.status = TokenStatusEnum.INVALID_TOKEN await UserAction.update_user(user) continue count += 1 bot.loop.create_task(run(user)) logs.info(f"初始化 Misskey Bot 完成,共启动 {count} 个 WS 任务")