"""Shared context, permission guard, message helpers and post send/edit operations.""" from __future__ import annotations import asyncio import functools import logging import shutil import tempfile import time from dataclasses import dataclass, field from types import SimpleNamespace from typing import Any, Callable, Optional from telethon import TelegramClient, errors, events from .buttons_sync import ensure_markup, link_urls from .config import Config from .db import Database from .replace_engine import fetch_messages, is_skippable from .richmsg import rich_of, send_rich, with_rich from .tgutil import build_markup, de_entities, edit_raw, esc, explain_rpc, fetch_media, media_info, peer_of, with_media from .userbot import UserbotError, UserbotRefused, edit_problem_text log = logging.getLogger(__name__) class UserError(Exception): """An error whose message is safe and useful to show to the user.""" @dataclass class Ctx: cfg: Config db: Database client: TelegramClient admins: set = field(default_factory=set) state: dict = field(default_factory=dict) # user id -> {"mode": ..., ...} pending: dict = field(default_factory=dict) # tokens for confirm dialogs (replace, repost) seen_channels: dict = field(default_factory=dict) # channel id -> Channel entity seen in updates text_handlers: dict = field(default_factory=dict) # mode -> async fn(event, state) callbacks: dict = field(default_factory=dict) # prefix -> async fn(event, parts) pickers: dict = field(default_factory=dict) # channel chooser kind -> async fn(event, arg, page) that draws it again lock: asyncio.Lock = field(default_factory=asyncio.Lock) # one long job (replace / repost) at a time userbot: Any = None # app.userbot.Userbot or None reporter: Any = None # app.errorlog.ErrorReporter or None expirer: Any = None # app.expiry.Expirer (timed deletions of broadcasts) or None syncer: Any = None # app.syncer.Syncer (keeps My posts in step with the channels) or None sync_ignore: set = field(default_factory=set) # (channel id, message id) of posts the owner made the bot forget def is_cb(event) -> bool: return isinstance(event, events.CallbackQuery.Event) or getattr(event, "_is_callback", False) def guard(ctx: Ctx, owner: bool = False) -> Callable: """Only allowed users may use a handler; errors are turned into friendly messages.""" def deco(fn): @functools.wraps(fn) async def wrapper(event): uid = event.sender_id ok = ctx.cfg.is_owner(uid) if owner else (ctx.cfg.is_owner(uid) or uid in ctx.admins or uid in ctx.cfg.admins) if not ok: try: if is_cb(event): await event.answer("Not allowed.", alert=True) elif getattr(event, "is_private", False): await event.respond("⛔ This bot is private.") except Exception: pass return try: await fn(event) except events.StopPropagation: raise except errors.MessageNotModifiedError: if is_cb(event): try: await event.answer() except Exception: pass except UserError as e: await notify_error(event, str(e)) except errors.RPCError as e: log.error("Telegram refused a request in handler %s (user %s)", fn.__name__, uid, exc_info=True) await notify_error(event, "❌ " + explain_rpc(e)) except Exception as e: log.exception("Unhandled exception in handler %s (user %s)", fn.__name__, uid) await notify_error(event, f"⚠️ Something went wrong: {type(e).__name__}: {str(e)[:200]}") return wrapper return deco async def notify_error(event, text: str) -> None: try: if is_cb(event): try: await event.answer("That didn't work - see the message below.", alert=False) except Exception: pass await event.respond(esc(text), link_preview=False) except Exception: # pragma: no cover log.exception("could not report an error") async def say(event, text: str, buttons=None): return await event.respond(text, buttons=buttons, link_preview=False) async def edit_callback_message(event, text: str, buttons=None): """Edit the message a button was pressed on. Telethon's CallbackQuery.edit() quietly starts a second `answer` task; the dispatcher has already answered the query, so we call edit_message directly instead. """ query = getattr(event, "query", None) if query is not None and hasattr(event, "client"): return await event.client.edit_message( await event.get_input_chat(), query.msg_id, text, buttons=buttons, link_preview=False ) return await event.edit(text, buttons=buttons, link_preview=False) # test doubles async def show(event, text: str, buttons=None): """Edit the message a button was pressed on; for normal messages send a new one.""" if is_cb(event): try: return await edit_callback_message(event, text, buttons) except errors.MessageNotModifiedError: return None return await say(event, text, buttons) def require_owner(ctx: Ctx, event) -> None: if not ctx.cfg.is_owner(event.sender_id): raise UserError("⛔ Only the bot owner can do that.") # ------------------------------------------------------------------ sending and editing async def send_post(client, peer, post): ents = de_entities(post.entities) markup = build_markup(post.buttons) rich = rich_of(post) if rich is not None and not post.media_file_id: # a rich message: Telegram builds its blocks from what was kept return await send_rich(client, peer, rich, markup=markup) if post.media_file_id: return await with_media( post.media_file_id, lambda media: client.send_file( peer, media, caption=post.text or "", formatting_entities=ents, buttons=markup ), ) return await client.send_message( peer, post.text, formatting_entities=ents, buttons=markup, link_preview=bool(post.link_preview) ) async def push_edit(client, ch, p, *, replace_media: bool = False) -> None: """Make the live channel message look like `p` (any object with the Post fields).""" peer = peer_of(ch) ents = de_entities(p.entities) markup = build_markup(p.buttons) rich = rich_of(p) if rich is not None and not p.media_kind: await with_rich(rich, lambda inp: edit_raw(client, peer, p.message_id, rich_message=inp, markup=markup)) return if replace_media and p.media_file_id: await with_media( p.media_file_id, lambda media: client.edit_message( peer, p.message_id, text=p.text or "", formatting_entities=ents, buttons=markup, file=media ), ) else: await edit_raw( client, peer, p.message_id, text=p.text or "", entities=ents, markup=markup, preview=None if p.media_kind else bool(p.link_preview), ) # What the userbot can change in a post. It is a user account: it can't set buttons, and it fails the way the bot did # when Telegram dislikes the text itself. USERBOT_EDITS = {"text", "entities", "link_preview"} NOT_THE_USERBOTS_JOB = { "EntityBoundsInvalidError", "EntitiesTooLongError", "MessageTooLongError", "MediaCaptionTooLongError", "MessageEmptyError", "ButtonUrlInvalidError", "ReplyMarkupInvalidError", } @dataclass class EditReport: """How the edit of a published post went (for the message the owner gets).""" by_userbot: bool = False # Telegram refused the bot, so the userbot made the edit bot_error: Optional[str] = None # the name of the error Telegram gave the bot buttons: Optional[str] = None # after the userbot's edit: "kept" | "restored" | "lost" (None: no link buttons saved) notes: list = field(default_factory=list) # what the userbot had to do in the channel (joined, got a right ...) media: Optional[tuple] = None # (kind, reference) of the photo / file the post shows now, as the bot sees it def edit_report_text(report: Optional[EditReport]) -> str: """The lines to add to the owner's message when the userbot had to make the edit (empty for an ordinary edit).""" if report is None or not report.by_userbot: return "" lines = [f"🤖 Telegram refused the bot ({esc(report.bot_error or 'error')}), so the userbot made the edit."] if report.buttons == "kept": lines.append("The buttons are still on the post.") elif report.buttons == "restored": lines.append("Telegram dropped the buttons with that edit, so the bot put them back.") elif report.buttons == "lost": lines.append( "Telegram dropped the buttons with that edit and the bot may not put them back. They stay saved in " "My posts; 🔧 Check buttons tries again." ) if report.notes: text = "; ".join(report.notes) lines.append("ℹ️ " + esc(text[:1].upper() + text[1:]) + ".") return "\n".join(lines) async def put_buttons_back(client, ch, p) -> Optional[str]: """After the userbot's edit: make the post show its saved link buttons again where Telegram dropped them. "kept" (they were there), "restored", "lost" (Telegram won't let the bot set them), None (none are saved).""" markup = build_markup(p.buttons) if not link_urls(markup): return None try: how = await ensure_markup(client, peer_of(ch), p.message_id, markup) except Exception as e: # noqa: BLE001 - the edit itself worked; only the buttons are missing log.info("could not put the buttons back on %s: %s", p.message_id, type(e).__name__) return "lost" return "kept" if how == "fine" else "restored" async def live_media(client, ch, message_id: int) -> Optional[tuple]: """(kind, reference) of the photo / file the channel post shows now, as the bot sees it; None if it can't be read.""" try: msgs = await fetch_messages(client, peer_of(ch), [message_id]) except Exception: # noqa: BLE001 return None m = msgs[0] if msgs else None if is_skippable(m): return None kind, ref = media_info(m) return (kind, ref) if kind else None async def edit_by_userbot(ctx: Ctx, ch, p, *, replace_media: bool, bot_error: str, notify=None) -> Optional[EditReport]: """The userbot makes the live post look like `p`. None when there is no userbot to ask; UserbotError when it fails. A new photo / file goes through the userbot's own chat with the bot: the bot sends it there and the userbot edits the post with that copy, so nothing is downloaded or uploaded - a big file is as quick as a small one and keeps its thumbnail. Only when that does not work does the bot download the file (with its thumbnail) and the userbot upload it again, which is slow for a big file. `notify(text)` tells the owner what is going on.""" ub = ctx.userbot editor = await ub.editor_for(ch) if ub is not None else None if editor is None: return None report = EditReport(by_userbot=True, bot_error=bot_error) folder = None async def tell(text: str) -> None: if notify is not None: await notify(text) fields = dict( text=p.text or "", entities=de_entities(p.entities), preview=None if p.media_kind or rich_of(p) else bool(p.link_preview), rich=None if p.media_kind else rich_of(p), ) try: await tell("⏳ Telegram refused the bot, so the userbot is making the edit.") if not (replace_media and p.media_file_id): await editor.edit(p.message_id, **fields) else: done = False media = await editor.media_by_reference(p.media_file_id) if media is not None: try: await editor.edit(p.message_id, media=media, **fields) done = True editor.notes.append("the new file went through the userbot's chat with the bot, so nothing was uploaded") except UserbotRefused as e: if not e.about_media: raise editor.dm_problem = f"Telegram would not take the userbot's copy of the file ({e.name})" if not done: why = editor.dm_problem or "that did not work" editor.notes.append(f"the file was downloaded and uploaded again, because {why}") await tell( "⏳ The file can't go through the userbot's chat with the bot, so it is downloaded and uploaded " f"again - a big file takes a while. (Reason: {esc(why)}.)" ) folder = tempfile.mkdtemp(prefix="cpb-media-") try: media = await fetch_media(ctx.client, p.media_file_id, p.media_kind, folder) except (errors.RPCError, OSError, ValueError) as e: raise UserbotError(f"I could not download the new media for the userbot ({type(e).__name__}).") await editor.edit(p.message_id, media=media, **fields) finally: await editor.close() report.notes = list(editor.notes) if folder: shutil.rmtree(folder, ignore_errors=True) if replace_media: report.media = await live_media(ctx.client, ch, p.message_id) report.buttons = await put_buttons_back(ctx.client, ch, p) return report async def apply_edit_reported(ctx: Ctx, post, ch, *, replace_media: bool = False, notify=None, **changes): """Edit the channel message first (if the post is published), and only then save the change. The bot edits. If Telegram refuses it (a post that belongs to another bot) and the change is to the text, caption or media, the userbot makes the edit instead. Returns (the saved post, an EditReport).""" report = EditReport() if post.status == "sent" and post.message_id: new = SimpleNamespace(**{**post.as_dict(), **changes}) try: await push_edit(ctx.client, ch, new, replace_media=replace_media) except errors.MessageNotModifiedError: pass except errors.RPCError as e: refused = "❌ " + explain_rpc(e) if type(e).__name__ in NOT_THE_USERBOTS_JOB or not (replace_media or USERBOT_EDITS & set(changes)): raise UserError(refused + "\nNothing was changed.") try: done = await edit_by_userbot( ctx, ch, new, replace_media=replace_media, bot_error=type(e).__name__, notify=notify ) except UserbotError as ue: log.info("the userbot could not make the edit either: %s", ue) raise UserError(f"{refused} {edit_problem_text(ctx.userbot, fallback_error=str(ue))}\nNothing was changed.") if done is None: raise UserError(f"{refused} {edit_problem_text(ctx.userbot)}\nNothing was changed.") report = done if report.media: # the post now shows the userbot's upload: keep My posts in step with it changes = {**changes, "media_kind": report.media[0], "media_file_id": report.media[1]} return await ctx.db.update_post(post.id, **changes), report async def apply_edit(ctx: Ctx, post, ch, *, replace_media: bool = False, **changes): """apply_edit_reported for the callers that only want the saved post.""" saved, _ = await apply_edit_reported(ctx, post, ch, replace_media=replace_media, **changes) return saved def purge_pending(ctx: Ctx, ttl: float = 3600.0) -> None: """Forget confirm dialogs (replace plans, repost choices...) nobody answered for a while.""" now = time.monotonic() stale = [ k for k, v in ctx.pending.items() if not getattr(v, "running", False) and now - getattr(v, "created", now) > ttl ] for k in stale: ctx.pending.pop(k, None) def on_cb(ctx: Ctx, prefix: str): def deco(fn): ctx.callbacks[prefix] = fn return fn return deco def on_text(ctx: Ctx, mode: str): def deco(fn): ctx.text_handlers[mode] = fn return fn return deco def set_state(ctx: Ctx, uid: int, mode: str, **data: Any) -> dict: ctx.state[uid] = {"mode": mode, **data} return ctx.state[uid] def clear_state(ctx: Ctx, uid: int) -> Optional[dict]: return ctx.state.pop(uid, None) def cmd(name: str, args: bool = False): """Event builder for a private-chat command such as /new or /replace .""" pat = rf"(?is)^/{name}(?:@\w+)?" + (r"(?:\s+(.*))?$" if args else r"\s*$") return events.NewMessage(incoming=True, pattern=pat, func=lambda e: e.is_private) async def discard_temp(ctx: Ctx, uid: int) -> None: """Drop a half-made draft when the user abandons the flow that was creating it.""" st = ctx.state.get(uid) if st and st.get("temp") and st.get("pid"): p = await ctx.db.get_post(st["pid"]) if p is not None and p.status == "draft": await ctx.db.delete_post(p.id)