feat(watcher,api,frontend): archive who writes to watched bots over mtproto

backend/src/watcher/registry.py

@@ -0,0 +1,156 @@
+import asyncio
+import contextlib
+from dataclasses import dataclass
+
+import asyncpg
+from pyrogram import Client
+from pyrogram.errors import RPCError
+from pyrogram.handlers import (
+    CallbackQueryHandler,
+    ChatMemberUpdatedHandler,
+    EditedMessageHandler,
+    MessageHandler,
+)
+
+from utils.logging import logger
+from utils.read import bots
+from utils.telegram import API_HASH, API_ID
+from watcher import tokens
+from watcher.backfill import backfill
+from watcher.capture import SOURCE_LIVE, callback_row, member_row, message_row, store
+
+DEVICE_MODEL = "Beavergram Watcher"
+TOKEN_MISSING = "Токен не найден на диске"  # noqa: S105
+INVALID_TOKEN_ERRORS = ("ACCESS_TOKEN_INVALID", "ACCESS_TOKEN_EXPIRED")
+
+
+@dataclass
+class RunningBot:
+    client: Client
+    backfill_task: asyncio.Task | None
+
+
+async def _cancel(task: asyncio.Task) -> None:
+    task.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await task
+
+
+def _handlers(pool: asyncpg.Pool, bot_id: int) -> tuple:
+    async def on_message(_client: Client, message) -> None:  # noqa: ANN001
+        if message.empty or message.chat is None or message.date is None:
+            return
+        await store(pool, message_row(bot_id, message, source=SOURCE_LIVE))
+
+    async def on_callback(_client: Client, query) -> None:  # noqa: ANN001
+        await store(pool, callback_row(bot_id, query))
+
+    async def on_member(_client: Client, update) -> None:  # noqa: ANN001
+        await store(pool, member_row(bot_id, update))
+
+    return (
+        MessageHandler(on_message),
+        EditedMessageHandler(on_message),
+        CallbackQueryHandler(on_callback),
+        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) -> None:
+        self._pool = pool
+        self._running: dict[int, RunningBot] = {}
+        self._lock = asyncio.Lock()
+
+    @property
+    def running(self) -> list[int]:
+        return list(self._running)
+
+    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 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 _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}")
+
+    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 = Client(
+            str(bot_id),
+            api_id=API_ID,
+            api_hash=API_HASH,
+            bot_token=token,
+            workdir=str(tokens.sessions_dir()),
+            device_model=DEVICE_MODEL,
+        )
+        for handler in _handlers(self._pool, bot_id):
+            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
+
+        await bots.set_status(
+            self._pool, bot_id, bots.STATUS_OK, username=me.username, title=me.full_name
+        )
+        task = (
+            asyncio.create_task(backfill(client, self._pool, bot_id))
+            if await bots.needs_backfill(self._pool, bot_id)
+            else None
+        )
+        self._running[bot_id] = RunningBot(client, task)
+        logger.info(f"[green]Bot watched:[/] @{me.username} ({bot_id})")
+
+    async def _stop(self, bot_id: int) -> None:
+        running = self._running.pop(bot_id, None)
+        if running is None:
+            return
+        if running.backfill_task is not None:
+            await _cancel(running.backfill_task)
+        with contextlib.suppress(Exception):
+            await running.client.stop()
+        logger.info(f"[yellow]Bot unwatched:[/] {bot_id}")