282 lines
9.8 KiB
Python
282 lines
9.8 KiB
Python
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}")
|