Blame for finbot/telegram.py

"""Telegram Bot API client and update parsing.

Webhook rather than long polling, because free Spaces sleep after inactivity
and a sleeping container cannot poll. Telegram's POST wakes it instead.
"""

from __future__ import annotations

import logging
import re
from collections import OrderedDict
from dataclasses import dataclass
from typing import Any, Optional

import httpx

log = logging.getLogger(__name__)

API_ROOT = "https://api.telegram.org"

#: getFile / downloads are capped at 20 MB by the Bot API itself.
MAX_DOWNLOAD_BYTES = 20 * 1024 * 1024

#: Telegram rejects messages over 4096 characters.
MAX_MESSAGE_CHARS = 4096


@dataclass
class ParsedUpdate:
    update_id: int
    user_id: int
    chat_id: int
    #: Current @handle without the @, empty if the user has not set one.
    #: Mutable -- never use as a storage key, only for access checks and display.
    username: str = ""
    first_name: str = ""
    text: str = ""
    command: str = ""
    args: str = ""
    caption: str = ""
    file_id: str = ""
    file_name: str = ""
    mime_type: str = ""
    file_kind: str = ""
    file_size: int = 0

    @property
    def has_file(self) -> bool:
        return bool(self.file_id)


def parse_update(update: dict[str, Any], bot_username: str = "") -> Optional[ParsedUpdate]:
    """Normalise a Telegram update. Returns None for anything we don't handle."""
    message = update.get("message") or update.get("edited_message")
    if not isinstance(message, dict):
        return None

    sender = message.get("from") or {}
    chat = message.get("chat") or {}
    if not sender.get("id") or not chat.get("id"):
        return None

    parsed = ParsedUpdate(
        update_id=int(update.get("update_id", 0)),
        user_id=int(sender["id"]),
        chat_id=int(chat["id"]),
        username=(sender.get("username") or "").lstrip("@"),
        first_name=(sender.get("first_name") or "").strip(),
        caption=(message.get("caption") or "").strip(),
    )

    text = (message.get("text") or "").strip()
    parsed.text = text

    if text.startswith("/"):
        head, _, rest = text.partition(" ")
        command = head[1:]
        # Group chats deliver "/month@MyBot"; strip the mention.
        if "@" in command:
            command, _, mentioned = command.partition("@")
            if bot_username and mentioned.lower() != bot_username.lower():
                return None
        parsed.command = command.lower()
        parsed.args = rest.strip()

    if photos := message.get("photo"):
        # PhotoSize list runs smallest to largest; the largest reads best.
        largest = max(photos, key=lambda p: p.get("file_size") or 0)
        parsed.file_id = largest.get("file_id", "")
        parsed.file_kind = "photo"
        parsed.mime_type = "image/jpeg"
        parsed.file_size = int(largest.get("file_size") or 0)
    elif document := message.get("document"):
        parsed.file_id = document.get("file_id", "")
        parsed.file_kind = "document"
        parsed.file_name = document.get("file_name", "")
        parsed.mime_type = document.get("mime_type", "")
        parsed.file_size = int(document.get("file_size") or 0)

    return parsed


def chunk_message(text: str, limit: int = MAX_MESSAGE_CHARS) -> list[str]:
    """Split a long reply on line boundaries so HTML tags stay balanced."""
    if len(text) <= limit:
        return [text]

    chunks: list[str] = []
    current: list[str] = []
    length = 0
    for line in text.split("\n"):
        # +1 for the newline we will rejoin with.
        if length + len(line) + 1 > limit and current:
            chunks.append("\n".join(current))
            current, length = [], 0
        current.append(line)
        length += len(line) + 1
    if current:
        chunks.append("\n".join(current))
    return chunks


class SeenUpdates:
    """Bounded set of processed update ids.

    We answer Telegram's webhook immediately and process in the background, so
    a slow reply can cause Telegram to redeliver. Without this, one statement
    could be imported twice.
    """

    def __init__(self, capacity: int = 512):
        self.capacity = capacity
        self._seen: OrderedDict[int, None] = OrderedDict()

    def add_if_new(self, update_id: int) -> bool:
        if update_id in self._seen:
            return False
        self._seen[update_id] = None
        while len(self._seen) > self.capacity:
            self._seen.popitem(last=False)
        return True


class TelegramClient:
    def __init__(self, token: str):
        self.token = token
        self._client: Optional[httpx.AsyncClient] = None
        self.username = ""

    @property
    def enabled(self) -> bool:
        return bool(self.token)

    def _http(self) -> httpx.AsyncClient:
        if self._client is None:
            self._client = httpx.AsyncClient(
                base_url=f"{API_ROOT}/bot{self.token}",
                timeout=httpx.Timeout(60.0, connect=10.0),
            )
        return self._client

    async def close(self) -> None:
        if self._client is not None:
            await self._client.aclose()
            self._client = None

    async def _call(self, method: str, **payload: Any) -> dict:
        response = await self._http().post(f"/{method}", json=payload)
        data = response.json()
        if not data.get("ok"):
            raise RuntimeError(f"{method} failed: {data.get('description', data)}")
        return data.get("result") or {}

    async def get_me(self) -> dict:
        result = await self._call("getMe")
        self.username = result.get("username", "")
        return result

    async def send_message(
        self, chat_id: int, text: str, parse_mode: Optional[str] = "HTML"
    ) -> None:
        for chunk in chunk_message(text):
            payload: dict[str, Any] = {
                "chat_id": chat_id,
                "text": chunk,
                "link_preview_options": {"is_disabled": True},
            }
            if parse_mode:
                payload["parse_mode"] = parse_mode
            try:
                await self._call("sendMessage", **payload)
            except RuntimeError as exc:
                # Malformed HTML from a merchant name shouldn't swallow the
                # reply -- retry once as plain text so the user still sees it.
                if parse_mode:
                    log.warning("HTML send failed (%s); retrying as plain text", exc)
                    await self._call(
                        "sendMessage",
                        chat_id=chat_id,
                        text=_strip_tags(chunk),
                    )
                else:
                    raise

    async def send_document(
        self, chat_id: int, filename: str, data: bytes, caption: str = ""
    ) -> None:
        files = {"document": (filename, data, "text/csv")}
        form: dict[str, Any] = {"chat_id": str(chat_id)}
        if caption:
            form["caption"] = caption[:1024]
            form["parse_mode"] = "HTML"
        response = await self._http().post("/sendDocument", data=form, files=files)
        payload = response.json()
        if not payload.get("ok"):
            raise RuntimeError(f"sendDocument failed: {payload.get('description')}")

    async def send_typing(self, chat_id: int) -> None:
        try:
            await self._call("sendChatAction", chat_id=chat_id, action="typing")
        except RuntimeError:
            pass  # cosmetic only

    async def download(self, file_id: str) -> tuple[bytes, str]:
        """Fetch a file by id. Returns (bytes, telegram_file_path)."""
        info = await self._call("getFile", file_id=file_id)
        file_path = info.get("file_path")
        if not file_path:
            raise RuntimeError("Telegram did not return a file path")

        url = f"{API_ROOT}/file/bot{self.token}/{file_path}"
        async with httpx.AsyncClient(timeout=httpx.Timeout(120.0)) as client:
            response = await client.get(url)
            if response.status_code >= 400:
                # raise_for_status() would put the token-bearing URL into the
                # exception, and from there into a traceback in public logs.
                raise RuntimeError(
                    f"file download failed: HTTP {response.status_code}"
                )
            return response.content, file_path

    async def set_webhook(self, url: str, secret: str) -> dict:
        return await self._call(
            "setWebhook",
            url=url,
            secret_token=secret,
            allowed_updates=["message", "edited_message"],
            drop_pending_updates=True,
        )

    async def delete_webhook(self) -> dict:
        return await self._call("deleteWebhook", drop_pending_updates=True)

    async def get_webhook_info(self) -> dict:
        return await self._call("getWebhookInfo")

    async def set_my_commands(self, commands: list[dict[str, str]]) -> dict:
        return await self._call("setMyCommands", commands=commands)


def _strip_tags(text: str) -> str:
    return re.sub(r"<[^>]+>", "", text)


def redact(text: str, token: str) -> str:
    """Remove a bot token from text before it reaches a log.

    The Bot API puts the token in the URL path, so any httpx exception string
    carries it. Space logs are world-readable on a public Space, so anything
    derived from a failed request must pass through here first.
    """
    text = str(text)
    if token:
        text = text.replace(token, "***")
    # Belt and braces: catch any /bot<token>/ shape, including a stale token.
    return re.sub(r"/bot\d+:[A-Za-z0-9_-]+", "/bot***", text)