Spaces:
Paused
Paused
| """/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 | |
| 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 | |
| ) | |
| class ChangeInfo: | |
| message_id: int | |
| buttons: int = 0 | |
| links: int = 0 | |
| typed: int = 0 | |
| mine: bool = False | |
| callbacks: bool = False | |
| 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 | |
| def total(self) -> int: | |
| return len(self.infos) | |
| def foreign(self) -> int: | |
| return sum(1 for i in self.infos.values() if not i.mine) | |
| 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 | |
| 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 | |