"""/swap engine: scan channels message by message, rewrite links, log every change, undo.""" from __future__ import annotations import asyncio import logging from dataclasses import dataclass, field from typing import Any, Awaitable, Callable, Optional from telethon import errors, functions, types from .handleswap import compute_handle_changes from .linkswap import compute_message_changes from .richmsg import entry as rich_entry, is_rich_message, plain_of_message, rich_from_received, rich_of from .tgutil import ( build_markup, de_entities, edit_raw, flood_retry, media_info, peer_of, ser_entities, ser_markup, state_of, ) log = logging.getLogger(__name__) CHUNK = 100 # Telegram returns at most 100 messages per request FALLBACK_EMPTY_BATCHES = 15 # used only when the newest message id can't be determined ABORT_AFTER = 5 # stop a channel after this many identical Telegram refusals in a row @dataclass class ScanOptions: old: str new: str include_typed: bool = True include_posts: bool = False last: Optional[int] = None mode: str = "links" # "links": the username inside t.me links (/swap); "handles": '@name' written in a post (/handleswap) def compute_changes(m, opts: ScanOptions) -> tuple: """What this run would change in message `m`: (MsgChange or None, skipped).""" if opts.mode == "handles": return compute_handle_changes(m, opts.old, opts.new) return compute_message_changes( m, opts.old, opts.new, include_typed=opts.include_typed, include_posts=opts.include_posts ) @dataclass class ChangeInfo: message_id: int buttons: int = 0 links: int = 0 typed: int = 0 mine: bool = False callbacks: bool = False @dataclass class ChannelScan: channel: Any can_edit_others: bool = True top: int = 0 first: int = 1 scanned: int = 0 infos: dict = field(default_factory=dict) skipped_post_links: int = 0 not_editable: int = 0 error: Optional[str] = None rights: Any = None @property def total(self) -> int: return len(self.infos) @property def foreign(self) -> int: return sum(1 for i in self.infos.values() if not i.mine) @dataclass class ApplyResult: edited: int = 0 unchanged: int = 0 missing: int = 0 no_perm: int = 0 failed: list = field(default_factory=list) aborted: Optional[str] = None # error name that made us stop early @dataclass class UndoResult: restored: int = 0 failed: list = field(default_factory=list) async def fetch_messages(client, peer, ids: list) -> list: res = await flood_retry(lambda: client.get_messages(peer, ids=ids)) return list(res) async def get_top_id(client, peer) -> Optional[int]: """Upper bound for the newest message id: the channel's event counter (pts) is never lower.""" try: full = await flood_retry(lambda: client(functions.channels.GetFullChannelRequest(peer))) pts = getattr(full.full_chat, "pts", None) if pts: return int(pts) except Exception as e: log.info("could not read channel counter: %s", e) return None def is_skippable(m) -> bool: return m is None or isinstance(m, (types.MessageEmpty, types.MessageService)) async def scan_channel( client, ch, opts: ScanOptions, *, can_edit_others: bool = True, mine_ids: Optional[set] = None, progress: Optional[Callable[[ChannelScan], Awaitable[None]]] = None, ) -> ChannelScan: scan = ChannelScan(channel=ch, can_edit_others=can_edit_others) peer = peer_of(ch) mine_ids = mine_ids or set() top = await get_top_id(client, peer) first = 1 if opts.last and top: first = max(1, top - opts.last + 1) scan.first, scan.top = first, top or 0 a, empty_run = first, 0 while True: if top is not None and a > top: break hi = a + CHUNK - 1 if top is None else min(a + CHUNK - 1, top) msgs = await fetch_messages(client, peer, list(range(a, hi + 1))) found = False for m in msgs: if m is None or isinstance(m, types.MessageEmpty): continue found = True if isinstance(m, types.MessageService): continue change, skipped = compute_changes(m, opts) scan.skipped_post_links += skipped if not change: continue mine = change.mine or m.id in mine_ids if not mine and not can_edit_others: scan.not_editable += 1 continue scan.infos[m.id] = ChangeInfo( m.id, change.n_buttons, change.n_links, change.n_typed, mine, change.callbacks ) scan.scanned = hi if top is None: empty_run = 0 if found else empty_run + 1 if empty_run >= FALLBACK_EMPTY_BATCHES: break if progress: await progress(scan) a = hi + 1 if top is None: scan.top = scan.scanned return scan def preview_flag(m): if isinstance(m.media, types.MessageMediaWebPage): return True return False if m.media is None else None async def _do_edit(client, peer, m, change, markup, preview) -> None: if change.text_changed or change.entities_changed: await edit_raw( client, peer, m.id, text=change.new_text, entities=change.new_entities, markup=markup, preview=preview ) else: # only button links changed: leave the text untouched await edit_raw(client, peer, m.id, markup=markup) async def apply_channel( client, db, scan: ChannelScan, opts: ScanOptions, batch_id: str, user_id: int, *, edit_delay: float = 1.2, progress: Optional[Callable[[ApplyResult], Awaitable[None]]] = None, ) -> ApplyResult: ch = scan.channel peer = peer_of(ch) mine_ids = await db.mine_message_ids(ch.id) res = ApplyResult() ids = sorted(scan.infos) buf: list = [] streak, last_name = 0, None async def flush() -> None: nonlocal buf if buf: batch, buf = buf, [] await db.record_changes(batch_id, ch.id, batch, user_id) try: for i in range(0, len(ids), CHUNK): chunk = ids[i:i + CHUNK] msgs = await fetch_messages(client, peer, chunk) for mid, m in zip(chunk, msgs): if is_skippable(m): res.missing += 1 continue # re-check against the live message: it may have been edited since the scan change, _ = compute_changes(m, opts) if not change: res.unchanged += 1 continue if not (change.mine or m.id in mine_ids) and not scan.can_edit_others: res.no_perm += 1 continue preview = preview_flag(m) markup = change.new_markup if change.new_markup is not None else m.reply_markup old_state = state_of(m.message, m.entities, m.reply_markup, preview) rich = is_rich_message(m) if rich: # its text is blocks this run never touches: undo puts the buttons back and nothing else old_state["rich"] = True try: await flood_retry(lambda: _do_edit(client, peer, m, change, markup, preview)) except errors.MessageNotModifiedError: res.unchanged += 1 continue except errors.RPCError as e: name = type(e).__name__ log.warning("edit of %s/%s failed: %s", ch.id, m.id, name) res.failed.append((m.id, name)) streak = streak + 1 if name == last_name else 1 last_name = name if streak >= ABORT_AFTER: res.aborted = name break continue streak, last_name = 0, None res.edited += 1 kind, fid = media_info(m) adopt = { "text": change.new_text, "entities": ser_entities(change.new_entities), "buttons": ser_markup(markup), "link_preview": bool(preview), } if kind: adopt["media_kind"], adopt["media_file_id"] = kind, fid new_state = state_of(change.new_text, change.new_entities, markup, preview) if rich: new_state["rich"] = True saved = await db.find_by_message(ch.id, m.id) if saved is not None and rich_of(saved) is not None: del adopt["text"], adopt["entities"] # the saved copy keeps what it was sent from else: # a rich message nobody saved yet: it is taken over as its blocks adopt["text"] = plain_of_message(m) adopt["entities"] = [rich_entry(rich_from_received(m.rich_message))] buf.append( { "message_id": m.id, "old_state": old_state, "new_state": new_state, "adopt": adopt, } ) if len(buf) >= 10: await flush() if progress: await progress(res) await asyncio.sleep(edit_delay) if res.aborted: break finally: await flush() return res async def undo_batch( client, db, batch_id: str, *, edit_delay: float = 1.2, progress: Optional[Callable[[UndoResult], Awaitable[None]]] = None, ) -> UndoResult: """Put every post touched by a /swap run back exactly as it was.""" res = UndoResult() channels: dict = {} done_ids: list = [] restored: list = [] for c in await db.batch_changes(batch_id): if c.channel_id not in channels: channels[c.channel_id] = await db.get_channel(c.channel_id) ch = channels[c.channel_id] if ch is None: res.failed.append((c.message_id, "ChannelRemoved")) continue old = c.old_state or {} try: if old.get("rich"): # a rich message: only its buttons were changed, so only they are put back markup = build_markup(old.get("buttons")) if markup is not None: await flood_retry(lambda: edit_raw(client, peer_of(ch), c.message_id, markup=markup)) else: await flood_retry( lambda: edit_raw( client, peer_of(ch), c.message_id, text=old.get("text", ""), entities=de_entities(old.get("entities")), markup=build_markup(old.get("buttons")), preview=old.get("preview"), ) ) except errors.MessageNotModifiedError: pass except errors.RPCError as e: log.warning("undo of %s/%s failed: %s", c.channel_id, c.message_id, e) res.failed.append((c.message_id, type(e).__name__)) continue res.restored += 1 done_ids.append(c.id) restored.append((c.channel_id, c.message_id, old)) if progress: await progress(res) await asyncio.sleep(edit_delay) await db.finish_undo(batch_id, done_ids, restored, all_done=not res.failed) return res