import asyncio import contextlib from dataclasses import dataclass from datetime import timedelta import asyncpg from pyrogram import Client, enums from pyrogram.errors import RPCError from pyrogram.handlers import ( ChatMemberUpdatedHandler, EditedMessageHandler, MessageHandler, ) from userbot.modules.capture import capture_message from userbot.modules.capture.chat_meta import meta_from_chat from userbot.modules.client import PyroClient from userbot.modules.jobs import JobConsumer from userbot.modules.watches.seen import seen_message from utils.events import notify_bg_event from utils.jobs import enqueue from utils.logging import logger from utils.read import accounts, bots from utils.storage import ContentAddressedStorage from watcher import tokens from watcher.context import build_bot_capture_context from watcher.drain import DrainCursor, drain_backlog, first_cursor DEVICE_MODEL = "Beavergram Bot" TOKEN_MISSING = "Токен не найден на диске" # noqa: S105 INVALID_TOKEN_ERRORS = ("ACCESS_TOKEN_INVALID", "ACCESS_TOKEN_EXPIRED") BACKLOG_ATTEMPTS = 3 BACKLOG_RETRY_PAUSE = 60 ENRICH_COOLDOWN = timedelta(hours=6) @dataclass class RunningBot: client: PyroClient account_id: int tasks: list[asyncio.Task] backlog: asyncio.Task | None = None async def _cancel(task: asyncio.Task) -> None: task.cancel() with contextlib.suppress(asyncio.CancelledError): await task def _handlers(pool: asyncpg.Pool) -> tuple: async def on_message(client: Client, message) -> None: # noqa: ANN001 ctx = getattr(client, "capture", None) if ctx is None or message.empty or message.chat is None or message.date is None: return meta = meta_from_chat(message.chat, ctx.contacts.ids) seen = seen_message(meta.chat_id, message) if message.edit_date is None: await ctx.watches.on_message(seen) else: await ctx.watches.on_edit(seen) toggles = ctx.resolve(meta) if not toggles.messages: return await capture_message(client, message, ctx, toggles) await notify_bg_event( ctx.pool, "message", ctx.account_id, chat_id=meta.chat_id, message_id=message.id, ) if message.chat.type is not enums.ChatType.PRIVATE: await enqueue( pool, ctx.account_id, "enrich_chat", {"chat_id": meta.chat_id}, cooldown=ENRICH_COOLDOWN, ) async def on_member(client: Client, update) -> None: # noqa: ANN001 ctx = getattr(client, "capture", None) if ctx is None or update.chat is None: return await enqueue( pool, ctx.account_id, "enrich_chat", {"chat_id": update.chat.id}, cooldown=ENRICH_COOLDOWN, ) return ( MessageHandler(on_message), EditedMessageHandler(on_message), ChatMemberUpdatedHandler(on_member), ) def _is_invalid_token(exc: RPCError) -> bool: return any(marker in str(exc) for marker in INVALID_TOKEN_ERRORS) class BotRegistry: def __init__(self, pool: asyncpg.Pool, storage: ContentAddressedStorage) -> None: self._pool = pool self._storage = storage self._running: dict[int, RunningBot] = {} self._lock = asyncio.Lock() @property def running(self) -> list[int]: return list(self._running) @property def connected_accounts(self) -> list[int]: return [ bot.account_id for bot in self._running.values() if bot.client.is_connected ] async def sync(self) -> None: async with self._lock: await self._adopt_dropped_in() enabled = await bots.enabled_bot_ids(self._pool) for bot_id in set(self._running) - enabled: await self._stop(bot_id) for bot_id in enabled - set(self._running): await self._start(bot_id) for bot_id in enabled & set(self._running): await self._resume_backlog(bot_id) for bot_id in await bots.removed_bot_ids(self._pool): tokens.drop(bot_id) async def close(self) -> None: for bot_id in list(self._running): await self._stop(bot_id) async def refresh_watches(self) -> None: for running in list(self._running.values()): capture = getattr(running.client, "capture", None) if capture is not None: await capture.watches.refresh() async def _adopt_dropped_in(self) -> None: unknown = tokens.stored_bot_ids() - await bots.known_bot_ids(self._pool) for bot_id in unknown: await bots.register_bot(self._pool, bot_id) logger.info(f"[green]Bot token adopted from disk:[/] {bot_id}") def _client(self, bot_id: int, token: str) -> PyroClient: return PyroClient( str(bot_id), workdir=str(tokens.sessions_dir()), device_model=DEVICE_MODEL, load_handlers=False, bot_token=token, skip_updates=False, ) async def _start(self, bot_id: int) -> None: token = tokens.load(bot_id) if token is None: await bots.set_status( self._pool, bot_id, bots.STATUS_UNAUTHORIZED, detail=TOKEN_MISSING ) return client = self._client(bot_id, token) for handler in _handlers(self._pool): client.add_handler(handler) try: await client.start() me = client.me if me is None: msg = "bot session is not authorized" raise RuntimeError(msg) except RPCError as exc: status = ( bots.STATUS_UNAUTHORIZED if _is_invalid_token(exc) else bots.STATUS_ERROR ) await bots.set_status(self._pool, bot_id, status, detail=str(exc)) logger.warning(f"[yellow]Bot {bot_id} failed to start:[/] {exc}") with contextlib.suppress(Exception): await client.stop() return except Exception as exc: await bots.set_status( self._pool, bot_id, bots.STATUS_ERROR, detail=str(exc) ) logger.exception(f"[red]Bot {bot_id} failed to start[/]") with contextlib.suppress(Exception): await client.stop() return try: account_id = await accounts.sync_account( self._pool, me, str(bot_id), is_bot=True ) client.capture = await build_bot_capture_context( client, self._pool, self._storage, account_id ) await bots.link_account(self._pool, bot_id, account_id) await bots.set_status( self._pool, bot_id, bots.STATUS_OK, username=me.username, title=me.full_name, ) await accounts.notify_accounts_changed(self._pool) except Exception as exc: logger.exception( f"[red]Bot {bot_id} failed to start,[/] retrying on the next sync" ) with contextlib.suppress(Exception): await bots.set_status( self._pool, bot_id, bots.STATUS_ERROR, detail=str(exc) ) with contextlib.suppress(Exception): await client.stop() return consumer = JobConsumer(client, self._pool, account_id) self._running[bot_id] = RunningBot( client, account_id, [asyncio.create_task(consumer.run())] ) await self._resume_backlog(bot_id) logger.info( f"[green]Bot account {account_id} watched:[/] @{me.username} ({bot_id})" ) async def _resume_backlog(self, bot_id: int) -> None: running = self._running[bot_id] if running.backlog is not None and not running.backlog.done(): return if not await bots.needs_backfill(self._pool, bot_id): return running.backlog = asyncio.create_task( self._backlog(running.client, bot_id, running.account_id) ) async def _backlog(self, client: PyroClient, bot_id: int, account_id: int) -> None: async def save(at: DrainCursor) -> None: await bots.save_backlog_cursor(self._pool, bot_id, at.pts, at.qts, at.date) saved = await bots.backlog_cursor(self._pool, bot_id) cursor = DrainCursor(*saved) if saved else first_cursor() stored = 0 for attempt in range(BACKLOG_ATTEMPTS): if attempt: await asyncio.sleep(BACKLOG_RETRY_PAUSE) walk = await drain_backlog(client, account_id, cursor, save) stored += walk.stored cursor = walk.cursor if walk.done: await bots.mark_backfilled(self._pool, bot_id) logger.info( f"[green]Account {account_id}: backlog complete,[/] " f"{stored} messages stored" ) break else: logger.warning( f"[yellow]Account {account_id}: backlog unfinished at " f"pts={cursor.pts},[/] {stored} messages stored, resumes on restart" ) if stored: await accounts.notify_accounts_changed(self._pool) async def _stop(self, bot_id: int) -> None: running = self._running.pop(bot_id, None) if running is None: return for task in (*running.tasks, running.backlog): if task is not None: await _cancel(task) with contextlib.suppress(Exception): await running.client.stop() logger.info(f"[yellow]Bot unwatched:[/] {bot_id}")