Spaces:
Paused
Paused
| """Drive the real handlers with fake events and a fake Telegram client (no network).""" | |
| import itertools | |
| import re | |
| from types import SimpleNamespace | |
| from telethon import events, functions, types | |
| from app.common import Ctx | |
| from app.config import Config | |
| from app.db import Database | |
| from bot import register_all | |
| from .dbutil import db_url_for | |
| from .fakes import FakeChannelClient, rpc | |
| class FakeSent: | |
| """A message the bot sent to the user (so `.edit()` works on it).""" | |
| def __init__(self, out, text, buttons): | |
| self.out, self.text, self.buttons = out, text, buttons | |
| async def edit(self, text=None, buttons=None, **kw): | |
| if text is not None: | |
| self.text = text | |
| self.buttons = buttons | |
| self.out.log.append(("edit", self.text, buttons)) | |
| class FakeEvent: | |
| def __init__(self, uid, out, text="", data=None, message=None, reply=None): | |
| self.sender_id, self.out = uid, out | |
| self.is_private = True | |
| self.is_channel = False | |
| self.is_group = False | |
| self.raw_text = text | |
| self.text = text | |
| self.message = message or SimpleNamespace( | |
| raw_text=text, entities=[], media=None, grouped_id=None, forward=None | |
| ) | |
| self.pattern_match = None | |
| self._reply = reply | |
| self._is_callback = data is not None | |
| if data is not None: | |
| self.data = data.encode() | |
| async def respond(self, text=None, buttons=None, file=None, **kw): | |
| self.out.log.append(("say", text, buttons)) | |
| return FakeSent(self.out, text, buttons) | |
| async def edit(self, text=None, buttons=None, **kw): | |
| self.out.log.append(("edit", text, buttons)) | |
| async def answer(self, *a, **kw): | |
| return None | |
| async def get_input_chat(self): | |
| return types.InputPeerUser(self.sender_id, 0) | |
| async def get_reply_message(self): | |
| return self._reply | |
| class Outbox: | |
| def __init__(self): | |
| self.log = [] | |
| def last_text(self): | |
| return self.log[-1][1] if self.log else None | |
| def last_buttons(self): | |
| b = self.log[-1][2] | |
| return [[getattr(x, "text", None) for x in row] for row in (b or [])] | |
| def _data_of(self, btn): | |
| t = btn.type if hasattr(btn, "type") else btn | |
| return t.data.decode() | |
| def callback_data(self, label_part): | |
| """Data of the first inline button (in the last message) whose label contains `label_part`. | |
| A channel chooser writes the names in the message and puts only numbers on its buttons: when no button has | |
| the label, `label_part` is looked for in the numbered list and the button with that number is taken - what a | |
| person does when they find the channel in the list and tap its number.""" | |
| for row in self.log[-1][2] or []: | |
| for btn in row: | |
| if label_part in btn.text: | |
| return self._data_of(btn) | |
| listed = re.search(rf"(?m)^(\d+)\. [^\n]*{re.escape(label_part)}", norm(self.last_text)) | |
| if listed: | |
| for row in self.log[-1][2] or []: | |
| for btn in row: | |
| if btn.text == listed.group(1): | |
| return self._data_of(btn) | |
| raise AssertionError(f"no button {label_part!r} in {self.last_buttons()}") | |
| def callback_exact(self, label): | |
| """Data of the button whose label is exactly `label` (the number buttons of a chooser, ◀ ▶ ...).""" | |
| for row in self.log[-1][2] or []: | |
| for btn in row: | |
| if btn.text == label: | |
| return self._data_of(btn) | |
| raise AssertionError(f"no button {label!r} in {self.last_buttons()}") | |
| class FakeTG(FakeChannelClient): | |
| """The bot's Telegram connection: a simulated channel plus a log of what the bot sent to the user.""" | |
| def __init__(self): | |
| super().__init__() | |
| self.handlers = [] | |
| self.sent = [] | |
| self.edits = [] | |
| self.deleted = [] | |
| self.ids = itertools.count(100) | |
| self.dm_rich_ids = [] # ids of the rich messages the bot sent to a person (previews and checks) | |
| self.parse_mode = "html" | |
| self.rights = dict(admin=True, post=True, edit=True, delete=True, invite=False, add_admins=False) | |
| def on(self, builder): | |
| def deco(fn): | |
| self.handlers.append((builder, fn)) | |
| return fn | |
| return deco | |
| async def send_message(self, peer, message="", **kw): | |
| m = SimpleNamespace(id=next(self.ids)) | |
| self.sent.append(("text", peer, message, kw)) | |
| return m | |
| async def send_file(self, peer, file, **kw): | |
| if self._is_dm(peer): # to the userbot account: the chat simulation | |
| return await super().send_file(peer, file, **kw) | |
| m = SimpleNamespace(id=next(self.ids)) | |
| self.sent.append(("file", peer, file, kw)) | |
| return m | |
| async def edit_message(self, peer, message, **kw): | |
| self.edits.append(("edit_message", message, kw)) | |
| err = self.fail.get("edit_message") # Telegram refuses the bot's high-level edit (text / caption / new media) | |
| if err is not None: | |
| raise err | |
| chan = self.registry.get(getattr(peer, "channel_id", None), self) | |
| m = chan.msgs.get(message) | |
| if m is not None and getattr(m, "via_other_bot", False): # its keyboard belongs to another bot: only that bot may edit | |
| raise rpc("MessageIdInvalidError") | |
| async def delete_messages(self, peer, ids, **kw): | |
| if self._is_dm(peer): | |
| return await super().delete_messages(peer, ids, **kw) | |
| self.deleted.append(list(ids)) | |
| return [SimpleNamespace(pts_count=len(ids))] | |
| async def get_entity(self, what): | |
| """Public channels by @username, or any channel by PeerChannel - the ones put next to this one.""" | |
| for chan in self.registry.values(): | |
| if isinstance(what, str) and chan.username and chan.username.lower() == what.lstrip("@").lower(): | |
| return chan._channel_obj() | |
| if isinstance(what, types.PeerChannel) and what.channel_id == chan.cid: | |
| return chan._channel_obj() | |
| raise ValueError(f"Cannot find any entity corresponding to {what!r}") | |
| async def __call__(self, req): | |
| if isinstance(req, functions.messages.CheckChatInviteRequest): | |
| for chan in self.registry.values(): | |
| if chan.invite_hash == req.hash: | |
| return types.ChatInviteAlready(chat=chan._channel_obj()) | |
| raise ValueError("INVITE_HASH_INVALID") | |
| target = self._target(req) | |
| if target is not self: | |
| return await target(req) | |
| if isinstance(req, functions.messages.EditMessageRequest): | |
| self.edits.append(("raw", req.id, req)) | |
| err = self.fail.get("EditMessageRequest") | |
| if err is not None: | |
| raise err | |
| if req.id not in self.msgs: # a post the simulation doesn't hold: just note the edit | |
| return None | |
| if isinstance(req, functions.channels.DeleteMessagesRequest): | |
| self.deleted.append(list(req.id)) | |
| if isinstance(req, functions.messages.SendMessageRequest) and req.rich_message is not None: | |
| self.sent.append(("rich", req.peer, req.rich_message, {"reply_markup": req.reply_markup, "silent": req.silent})) | |
| if isinstance(req.peer, types.InputPeerUser): # a chat with a person (the preview): nothing lands in a channel | |
| if self.fail.get("SendMessageRequest") is not None: | |
| raise self.fail["SendMessageRequest"] | |
| bytes(req) # it must serialise | |
| self._rich(req.rich_message) # Telegram has to be able to make blocks of it | |
| mid = next(self.ids) | |
| self.dm_rich_ids.append(mid) | |
| return types.UpdateShortSentMessage(id=mid, pts=0, pts_count=0, date=None, out=True) | |
| return await super().__call__(req) | |
| class App: | |
| def __init__(self, tmp_path, owner=1, **cfg_overrides): | |
| base = dict(api_id=1, api_hash="x", bot_token="1:x", owners=frozenset({owner}), admins=frozenset(), | |
| database_url=db_url_for(tmp_path, "h.db"), db_pool="null", session_path="x", session_string="", | |
| edit_delay=0.3, replace_typed_links=True, port=0, log_level="INFO") | |
| base.update(cfg_overrides) | |
| cfg = Config(**base) | |
| self.tg = FakeTG() | |
| self.db = Database(cfg.database_url) | |
| self.ctx = Ctx(cfg=cfg, db=self.db, client=self.tg) | |
| self.uid = owner | |
| self.out = Outbox() | |
| async def start(self): | |
| from app.db import Base | |
| async with self.db.engine.begin() as conn: | |
| await conn.run_sync(Base.metadata.drop_all) | |
| await self.db.init() | |
| register_all(self.ctx) | |
| async def text(self, s, message=None, reply=None): | |
| ev = FakeEvent(self.uid, self.out, s, message=message, reply=reply) | |
| for builder, cb in self.tg.handlers: | |
| if not isinstance(builder, events.NewMessage): | |
| continue | |
| if builder.pattern: | |
| m = builder.pattern(s) | |
| if not m: | |
| continue | |
| ev.pattern_match = m | |
| if builder.func and not builder.func(ev): | |
| continue | |
| await cb(ev) | |
| return ev | |
| async def raw(self, update): | |
| """Hand a raw Telegram update (a new / edited / deleted channel message ...) to the handlers that listen for it.""" | |
| for builder, cb in self.tg.handlers: | |
| if isinstance(builder, events.Raw) and builder.filter(update) is not None: | |
| await cb(update) | |
| async def press(self, data): | |
| ev = FakeEvent(self.uid, self.out, data=data) | |
| for builder, cb in self.tg.handlers: | |
| if isinstance(builder, events.CallbackQuery): | |
| await cb(ev) | |
| return ev | |
| def norm(s): | |
| import html | |
| return html.unescape(re.sub(r"<[^>]+>", "", s or "")) | |