import asyncio import sqlite3 import time from contextlib import closing from dataclasses import dataclass, field from pathlib import Path from pyrogram import raw, utils from pyrogram.errors import FloodPremiumWait, FloodWait, RPCError from userbot.modules.capture.chat_meta import meta_from_peer from userbot.modules.capture.context import CaptureContext from userbot.modules.catchup import repository from userbot.modules.catchup.repository import UpdateCursor from userbot.modules.client import PyroClient from utils.logging import logger from utils.policy.models import ChatKind BOX_SLICES = 50 CHANNEL_SLICES = 10 CHANNEL_LIMIT = 100 MAX_CHANNELS = 50 QUEUE_HIGH_WATER = 200 QUEUE_WAIT_SECONDS = 0.5 NO_QTS = -1 MAX_AGE_SECONDS = 3 * 24 * 3600 @dataclass(frozen=True) class SavedState: box: UpdateCursor | None = None channels: dict[int, int] = field(default_factory=dict) @dataclass class BoxWalk: cursor: UpdateCursor done: bool = False updates: int = 0 channels: dict[int, raw.base.Chat] = field(default_factory=dict) def read_saved_state(path: Path) -> SavedState: try: uri = f"{path.resolve().as_uri()}?mode=ro" with closing(sqlite3.connect(uri, uri=True)) as conn: rows = conn.execute( "SELECT id, pts, qts, date FROM update_state" ).fetchall() except Exception as exc: logger.warning(f"[yellow]Session {path.stem}: no saved update state ({exc})[/]") return SavedState() box = None channels: dict[int, int] = {} for state_id, pts, qts, date in rows: if pts is None: continue if state_id == 0: box = UpdateCursor( pts, NO_QTS if qts is None else qts, date or int(time.time()) ) else: channels[state_id] = pts return SavedState(box, channels) async def _dispatch( client: PyroClient, new_messages: list, other_updates: list, users: list, chats: list, ) -> int: await client.fetch_peers(users) await client.fetch_peers(chats) user_map = {user.id: user for user in users} chat_map = {chat.id: chat for chat in chats} queue = client.dispatcher.updates_queue for message in new_messages: update = raw.types.UpdateNewMessage(message=message, pts=0, pts_count=0) queue.put_nowait((update, user_map, chat_map)) for update in other_updates: queue.put_nowait((update, user_map, chat_map)) while queue.qsize() > QUEUE_HIGH_WATER: # noqa: ASYNC110 await asyncio.sleep(QUEUE_WAIT_SECONDS) return len(new_messages) + len(other_updates) def _channel_hints(other_updates: list, chats: list) -> dict[int, raw.base.Chat]: by_id = {chat.id: chat for chat in chats} return { utils.get_channel_id(update.channel_id): by_id[update.channel_id] for update in other_updates if isinstance(update, raw.types.UpdateChannelTooLong) and update.channel_id in by_id } async def _walk_box( client: PyroClient, ctx: CaptureContext, cursor: UpdateCursor ) -> BoxWalk: walk = BoxWalk(cursor) for _ in range(BOX_SLICES): diff = await client.invoke( raw.functions.updates.GetDifference( pts=walk.cursor.pts, date=walk.cursor.date, qts=walk.cursor.qts ) ) if isinstance(diff, raw.types.updates.DifferenceEmpty): walk.done = True return walk if isinstance(diff, raw.types.updates.DifferenceTooLong): logger.warning( f"[yellow]Account {ctx.account_id}: too much missed to catch up " f"(pts {walk.cursor.pts} to {diff.pts}), that gap stays[/]" ) walk.done = True return walk walk.updates += await _dispatch( client, diff.new_messages, diff.other_updates, diff.users, diff.chats ) walk.channels.update(_channel_hints(diff.other_updates, diff.chats)) final = isinstance(diff, raw.types.updates.Difference) state = diff.state if final else diff.intermediate_state moved = UpdateCursor(state.pts, state.qts, state.date or walk.cursor.date) if moved == walk.cursor and not final: logger.warning( f"[yellow]Account {ctx.account_id}: catch-up stopped advancing " f"at pts={moved.pts}[/]" ) return walk walk.cursor = moved await repository.save_cursor(ctx.pool, ctx.account_id, moved) if final: walk.done = True return walk return walk async def _walk_channel(client: PyroClient, peer_id: int, pts: int) -> int: peer = await client.resolve_peer(peer_id) if not isinstance(peer, raw.types.InputPeerChannel): return 0 channel = raw.types.InputChannel( channel_id=peer.channel_id, access_hash=peer.access_hash ) updates = 0 for _ in range(CHANNEL_SLICES): diff = await client.invoke( raw.functions.updates.GetChannelDifference( channel=channel, filter=raw.types.ChannelMessagesFilterEmpty(), pts=pts, limit=CHANNEL_LIMIT, ) ) if isinstance(diff, raw.types.updates.ChannelDifferenceEmpty): return updates if isinstance(diff, raw.types.updates.ChannelDifferenceTooLong): logger.info(f"Channel {peer_id}: too much missed to catch up") return updates updates += await _dispatch( client, diff.new_messages, diff.other_updates, diff.users, diff.chats ) if diff.final: return updates pts = diff.pts return updates def _captured_channels( ctx: CaptureContext, hints: dict[int, raw.base.Chat], saved: dict[int, int] ) -> list[int]: wanted: list[tuple[bool, int]] = [] for peer_id, chat in hints.items(): if peer_id not in saved: continue meta = meta_from_peer( raw.types.PeerChannel(channel_id=chat.id), {chat.id: chat}, {}, ctx.contacts.ids, ) toggles = ctx.resolve(meta) if toggles.messages or toggles.track_edits_deletes: wanted.append((meta.kind is ChatKind.CHANNEL, peer_id)) return [peer_id for _, peer_id in sorted(wanted)[:MAX_CHANNELS]] async def _walk_channels( client: PyroClient, ctx: CaptureContext, walk: BoxWalk, saved: dict[int, int] ) -> int: peer_ids = _captured_channels(ctx, walk.channels, saved) updates = 0 for peer_id in peer_ids: try: updates += await _walk_channel(client, peer_id, saved[peer_id]) except (FloodWait, FloodPremiumWait) as exc: logger.warning( f"[yellow]Account {ctx.account_id}: channel catch-up stopped " f"by a flood wait:[/] {exc}" ) break except RPCError as exc: logger.warning( f"[yellow]Account {ctx.account_id}: channel {peer_id} " f"catch-up failed:[/] {exc}" ) skipped = len(walk.channels) - len(peer_ids) if skipped: logger.info(f"Account {ctx.account_id}: {skipped} channel gap(s) left as is") return updates async def catch_up(client: PyroClient, saved: SavedState) -> None: ctx = client.capture if ctx is None: return try: pending = await repository.load_cursor(ctx.pool, ctx.account_id) cursor = pending or saved.box if cursor is None: logger.info(f"Account {ctx.account_id}: no saved update state to catch up") return if time.time() - cursor.date > MAX_AGE_SECONDS: logger.warning( f"[yellow]Account {ctx.account_id}: last update seen too long ago " f"(pts={cursor.pts}), not catching up[/]" ) await repository.clear_cursor(ctx.pool, ctx.account_id) return await repository.save_cursor(ctx.pool, ctx.account_id, cursor) logger.info( f"Account {ctx.account_id}: catching up from pts={cursor.pts}" f"{' (resumed)' if pending else ''}" ) walk = await _walk_box(client, ctx, cursor) if not walk.done: logger.warning( f"[yellow]Account {ctx.account_id}: catch-up unfinished at " f"pts={walk.cursor.pts},[/] {walk.updates} updates replayed, " f"resumes on the next start" ) return from_channels = await _walk_channels(client, ctx, walk, saved.channels) await repository.clear_cursor(ctx.pool, ctx.account_id) logger.info( f"[green]Account {ctx.account_id}: caught up,[/] {walk.updates} updates " f"from the main box, {from_channels} from channels" ) except Exception: logger.exception( f"[red]Account {ctx.account_id}: catch-up failed,[/] " f"resumes on the next start" )