feat(watcher,api,frontend): archive who writes to watched bots over mtproto
This commit is contained in:
hh committed
1 parent
089a31db8d
commit
4bd031f21e
30 files changed
+1863
-10
No files matched your search
Whitespace-only changes.
@@ -0,0 +1,4 @@
|
||||
from .runner import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,71 @@
|
||||
import asyncio
|
||||
|
||||
import asyncpg
|
||||
from pyrogram import Client
|
||||
from pyrogram.errors import FloodWait, RPCError
|
||||
|
||||
from utils.env import env
|
||||
from utils.logging import logger
|
||||
from utils.read import bots
|
||||
from watcher.capture import SOURCE_HISTORY, Row, message_row, store_many
|
||||
|
||||
BATCH = 200
|
||||
MAX_FLOOD_WAIT = 300
|
||||
HISTORY_UNAVAILABLE = "История недоступна: {}"
|
||||
|
||||
|
||||
async def _flush(pool: asyncpg.Pool, rows: list[Row]) -> int:
|
||||
await store_many(pool, rows)
|
||||
stored = len(rows)
|
||||
rows.clear()
|
||||
return stored
|
||||
|
||||
|
||||
async def _chat_history(
|
||||
client: Client, pool: asyncpg.Pool, bot_id: int, chat_id: int
|
||||
) -> int:
|
||||
rows: list[Row] = []
|
||||
stored = 0
|
||||
async for message in client.get_chat_history(
|
||||
chat_id, limit=env.bots.backfill_per_chat
|
||||
):
|
||||
if message.empty or message.chat is None or message.date is None:
|
||||
continue
|
||||
rows.append(message_row(bot_id, message, source=SOURCE_HISTORY))
|
||||
if len(rows) >= BATCH:
|
||||
stored += await _flush(pool, rows)
|
||||
return stored + await _flush(pool, rows)
|
||||
|
||||
|
||||
async def backfill(client: Client, pool: asyncpg.Pool, bot_id: int) -> None:
|
||||
try:
|
||||
dialogs = [
|
||||
dialog.chat.id
|
||||
async for dialog in client.get_dialogs(limit=env.bots.backfill_chats)
|
||||
if dialog.chat is not None and dialog.chat.id is not None
|
||||
]
|
||||
except RPCError as exc:
|
||||
logger.warning(f"[yellow]Bot {bot_id}: dialogs unavailable:[/] {exc}")
|
||||
await bots.set_status(
|
||||
pool, bot_id, bots.STATUS_OK, detail=HISTORY_UNAVAILABLE.format(exc)
|
||||
)
|
||||
return
|
||||
|
||||
stored = 0
|
||||
for chat_id in dialogs:
|
||||
try:
|
||||
stored += await _chat_history(client, pool, bot_id, chat_id)
|
||||
except FloodWait as exc:
|
||||
wait = exc.value if isinstance(exc.value, int) else MAX_FLOOD_WAIT + 1
|
||||
if wait > MAX_FLOOD_WAIT:
|
||||
logger.warning(f"[yellow]Bot {bot_id}: flood wait {wait}s, stop[/]")
|
||||
break
|
||||
await asyncio.sleep(wait)
|
||||
except RPCError as exc:
|
||||
logger.warning(f"[yellow]Bot {bot_id}: chat {chat_id} skipped:[/] {exc}")
|
||||
|
||||
await bots.mark_backfilled(pool, bot_id)
|
||||
logger.info(
|
||||
f"[green]Bot {bot_id}: backfilled[/] {stored} messages "
|
||||
f"from {len(dialogs)} chats"
|
||||
)
|
||||
@@ -0,0 +1,146 @@
|
||||
from datetime import UTC, datetime
|
||||
|
||||
import asyncpg
|
||||
from pyrogram.types import CallbackQuery, ChatMemberUpdated, Message, User
|
||||
|
||||
from userbot.modules.raw_json import raw_json
|
||||
from utils.read import bots
|
||||
|
||||
SOURCE_HISTORY = "history"
|
||||
SOURCE_LIVE = "live"
|
||||
|
||||
KIND_MESSAGE = "message"
|
||||
KIND_CALLBACK = "callback_query"
|
||||
KIND_MEMBER = "member"
|
||||
|
||||
MEDIA_KINDS = (
|
||||
"photo",
|
||||
"video",
|
||||
"animation",
|
||||
"audio",
|
||||
"voice",
|
||||
"video_note",
|
||||
"document",
|
||||
"sticker",
|
||||
"contact",
|
||||
"location",
|
||||
"venue",
|
||||
"poll",
|
||||
"dice",
|
||||
"story",
|
||||
"game",
|
||||
)
|
||||
|
||||
Row = tuple
|
||||
|
||||
|
||||
def _media_kind(message: Message) -> str | None:
|
||||
return next(
|
||||
(kind for kind in MEDIA_KINDS if getattr(message, kind, None) is not None), None
|
||||
)
|
||||
|
||||
|
||||
def _full_name(user: User | None) -> str | None:
|
||||
if user is None:
|
||||
return None
|
||||
parts = [user.first_name, user.last_name]
|
||||
return " ".join(part for part in parts if part) or None
|
||||
|
||||
|
||||
def _chat_title(message: Message) -> str | None:
|
||||
chat = message.chat
|
||||
if chat is None:
|
||||
return None
|
||||
parts = [chat.first_name, chat.last_name]
|
||||
return chat.title or " ".join(part for part in parts if part) or None
|
||||
|
||||
|
||||
def _chat_type(message: Message) -> str | None:
|
||||
kind = getattr(message.chat, "type", None)
|
||||
return kind.value if kind is not None else None
|
||||
|
||||
|
||||
def message_row(bot_id: int, message: Message, *, source: str) -> Row:
|
||||
user = message.from_user
|
||||
chat = message.chat
|
||||
sender = user.id if user is not None else None
|
||||
if sender is None and message.sender_chat is not None:
|
||||
sender = message.sender_chat.id
|
||||
return (
|
||||
bot_id,
|
||||
KIND_MESSAGE,
|
||||
chat.id if chat is not None else 0,
|
||||
message.id,
|
||||
message.date,
|
||||
source,
|
||||
bool(message.outgoing),
|
||||
sender,
|
||||
user.username if user is not None else None,
|
||||
_full_name(user),
|
||||
user.is_bot if user is not None else None,
|
||||
_chat_type(message),
|
||||
_chat_title(message),
|
||||
message.text or message.caption,
|
||||
_media_kind(message),
|
||||
message.edit_date,
|
||||
raw_json(message),
|
||||
)
|
||||
|
||||
|
||||
def callback_row(bot_id: int, query: CallbackQuery) -> Row:
|
||||
user = query.from_user
|
||||
message = query.message
|
||||
chat = message.chat if message is not None else None
|
||||
return (
|
||||
bot_id,
|
||||
KIND_CALLBACK,
|
||||
chat.id if chat is not None else 0,
|
||||
message.id if message is not None else 0,
|
||||
datetime.now(UTC),
|
||||
SOURCE_LIVE,
|
||||
False,
|
||||
user.id if user is not None else None,
|
||||
user.username if user is not None else None,
|
||||
_full_name(user),
|
||||
user.is_bot if user is not None else None,
|
||||
chat.type.value if chat is not None and chat.type is not None else None,
|
||||
chat.title if chat is not None else None,
|
||||
query.data,
|
||||
None,
|
||||
None,
|
||||
raw_json(query),
|
||||
)
|
||||
|
||||
|
||||
def member_row(bot_id: int, update: ChatMemberUpdated) -> Row:
|
||||
member = update.new_chat_member
|
||||
user = member.user if member is not None else None
|
||||
chat = update.chat
|
||||
status = member.status.value if member is not None else None
|
||||
return (
|
||||
bot_id,
|
||||
KIND_MEMBER,
|
||||
chat.id if chat is not None else 0,
|
||||
user.id if user is not None else 0,
|
||||
update.date or datetime.now(UTC),
|
||||
SOURCE_LIVE,
|
||||
False,
|
||||
update.from_user.id if update.from_user is not None else None,
|
||||
update.from_user.username if update.from_user is not None else None,
|
||||
_full_name(update.from_user),
|
||||
update.from_user.is_bot if update.from_user is not None else None,
|
||||
chat.type.value if chat is not None and chat.type is not None else None,
|
||||
chat.title if chat is not None else None,
|
||||
status,
|
||||
None,
|
||||
None,
|
||||
raw_json(update),
|
||||
)
|
||||
|
||||
|
||||
async def store(pool: asyncpg.Pool, row: Row) -> None:
|
||||
await bots.store_update(pool, row)
|
||||
|
||||
|
||||
async def store_many(pool: asyncpg.Pool, rows: list[Row]) -> None:
|
||||
await bots.store_updates(pool, rows)
|
||||
@@ -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}")
|
||||
@@ -0,0 +1,72 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import signal
|
||||
|
||||
import asyncpg
|
||||
import uvloop
|
||||
|
||||
from dependencies.container import container
|
||||
from utils.env import env
|
||||
from utils.logging import logger, setup_logging
|
||||
from utils.read.bots import BOTS_CHANGED_CHANNEL
|
||||
from watcher.registry import BotRegistry
|
||||
|
||||
setup_logging()
|
||||
|
||||
RETRY_INTERVAL = 300
|
||||
|
||||
|
||||
async def _retry_loop(registry: BotRegistry) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(RETRY_INTERVAL)
|
||||
with contextlib.suppress(Exception):
|
||||
await registry.sync()
|
||||
|
||||
|
||||
async def _listen_changes(
|
||||
registry: BotRegistry, tasks: set[asyncio.Task]
|
||||
) -> asyncpg.Connection:
|
||||
def on_change(
|
||||
_conn: asyncpg.Connection, _pid: int, _channel: str, _payload: str
|
||||
) -> None:
|
||||
task = asyncio.create_task(registry.sync())
|
||||
tasks.add(task)
|
||||
task.add_done_callback(tasks.discard)
|
||||
|
||||
conn = await asyncpg.connect(dsn=env.db.connection_url)
|
||||
await conn.add_listener(BOTS_CHANGED_CHANNEL, on_change)
|
||||
return conn
|
||||
|
||||
|
||||
async def runner() -> None:
|
||||
pool = await container.get(asyncpg.Pool)
|
||||
registry = BotRegistry(pool)
|
||||
tasks: set[asyncio.Task] = set()
|
||||
listen_conn: asyncpg.Connection | None = None
|
||||
stop = asyncio.Event()
|
||||
asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, stop.set)
|
||||
retry = asyncio.create_task(_retry_loop(registry))
|
||||
try:
|
||||
await registry.sync()
|
||||
if not registry.running:
|
||||
logger.warning("[yellow]No bots yet. Add a bot token in the web UI.[/]")
|
||||
listen_conn = await _listen_changes(registry, tasks)
|
||||
logger.info("[green]Bot watcher running.[/]")
|
||||
await stop.wait()
|
||||
finally:
|
||||
retry.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await retry
|
||||
if listen_conn is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await listen_conn.close()
|
||||
await registry.close()
|
||||
await container.close()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
uvloop.install()
|
||||
logger.info("Starting bot watcher...")
|
||||
with contextlib.suppress(KeyboardInterrupt):
|
||||
asyncio.run(runner())
|
||||
logger.info("[red]Bot watcher stopped.[/]")
|
||||
@@ -0,0 +1,42 @@
|
||||
from pathlib import Path
|
||||
|
||||
from utils.env import env
|
||||
|
||||
TOKEN_SUFFIX = ".token" # noqa: S105
|
||||
SESSION_SUFFIX = ".session"
|
||||
TOKEN_MODE = 0o600
|
||||
|
||||
|
||||
def sessions_dir() -> Path:
|
||||
path = Path(env.bots.sessions_dir)
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def bot_id_of(token: str) -> int | None:
|
||||
head, _, secret = token.strip().partition(":")
|
||||
return int(head) if head.isdigit() and secret else None
|
||||
|
||||
|
||||
def save(bot_id: int, token: str) -> None:
|
||||
path = sessions_dir() / f"{bot_id}{TOKEN_SUFFIX}"
|
||||
path.write_text(token.strip(), encoding="utf-8")
|
||||
path.chmod(TOKEN_MODE)
|
||||
|
||||
|
||||
def load(bot_id: int) -> str | None:
|
||||
path = sessions_dir() / f"{bot_id}{TOKEN_SUFFIX}"
|
||||
return path.read_text(encoding="utf-8").strip() if path.is_file() else None
|
||||
|
||||
|
||||
def drop(bot_id: int) -> None:
|
||||
for suffix in (TOKEN_SUFFIX, SESSION_SUFFIX):
|
||||
(sessions_dir() / f"{bot_id}{suffix}").unlink(missing_ok=True)
|
||||
|
||||
|
||||
def stored_bot_ids() -> set[int]:
|
||||
return {
|
||||
int(path.stem)
|
||||
for path in sessions_dir().glob(f"*{TOKEN_SUFFIX}")
|
||||
if path.stem.isdigit()
|
||||
}
|
||||
Reference in new issue
Block a user