ForcePilot/backend/package/yuxi/channel/extensions/telegram/webhook.py

from __future__ import annotations

import asyncio
import hmac
import logging
import secrets

from fastapi import APIRouter, Header, Request

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/webhook/telegram", tags=["telegram"])

WEBHOOK_MAX_BODY_BYTES = 1_048_576
WEBHOOK_READ_TIMEOUT = 30.0

_webhook_app: object | None = None
_webhook_queue: asyncio.Queue | None = None
_webhook_account_id: str = "default"
_webhook_secret: str = ""
_webhook_token: str = ""

_webhook_registry: dict[str, dict] = {}


def _get_webhook_context(account_id: str) -> dict | None:
    return _webhook_registry.get(account_id)


@router.post("")
async def telegram_webhook_receive(
    request: Request,
    x_telegram_bot_api_secret_token: str | None = Header(None, alias="X-Telegram-Bot-Api-Secret-Token"),
):
    ctx = _resolve_webhook_context(x_telegram_bot_api_secret_token)
    if ctx is None:
        logger.warning("Telegram webhook: no matching account found")
        return {"status": "unknown-account"}

    expected_secret = ctx.get("secret", "")
    if expected_secret and not hmac.compare_digest(x_telegram_bot_api_secret_token or "", expected_secret):
        logger.warning("Telegram webhook secret mismatch")
        return {"status": "unauthorized"}

    try:
        body = await request.body()
    except Exception:
        logger.warning("Telegram webhook body read failed")
        return {"status": "ok"}

    if len(body) > WEBHOOK_MAX_BODY_BYTES:
        logger.warning("Telegram webhook body too large: %d bytes", len(body))
        return {"status": "ok"}

    import json

    try:
        update = json.loads(body)
    except json.JSONDecodeError:
        logger.warning("Telegram webhook invalid JSON")
        return {"status": "ok"}

    from yuxi.channel.extensions.telegram.monitor import convert_update_to_unified

    queue = ctx.get("queue")
    account_id = ctx.get("account_id", "default")
    unified = convert_update_to_unified(update, account_id)
    if unified and queue:
        try:
            queue.put_nowait(unified)
        except asyncio.QueueFull:
            logger.warning("Telegram webhook queue full, dropping message")

    if "callback_query" in update:
        cq_id = update["callback_query"].get("id", "")
        if cq_id:
            token = ctx.get("token", "")
            if token:
                from yuxi.channel.extensions.telegram.polling import _answer_callback_quiet

                asyncio.create_task(_answer_callback_quiet(token, cq_id))

    return {"status": "ok"}


def _resolve_webhook_context(secret_token: str | None) -> dict | None:
    if secret_token:
        for ctx in _webhook_registry.values():
            if ctx.get("secret") and hmac.compare_digest(secret_token, ctx["secret"]):
                return ctx

    entries = list(_webhook_registry.values())
    if len(entries) == 1:
        return entries[0]

    if _webhook_queue is not None:
        return {
            "queue": _webhook_queue,
            "account_id": _webhook_account_id,
            "secret": _webhook_secret,
            "token": _webhook_token,
        }

    return None


async def start_webhook(account: dict, queue: asyncio.Queue) -> None:
    global _webhook_app, _webhook_queue, _webhook_account_id, _webhook_secret, _webhook_token

    _webhook_queue = queue
    _webhook_account_id = account.get("account_id", "default")
    _webhook_secret = account.get("webhook_secret", "")
    _webhook_token = account.get("token", "")

    _webhook_registry[_webhook_account_id] = {
        "queue": queue,
        "secret": _webhook_secret,
        "token": _webhook_token,
    }

    token = account.get("token", "")
    webhook_url = account.get("webhook_url", "")

    if not token:
        logger.error("Telegram webhook: no token configured")
        return

    if not webhook_url:
        logger.warning("Telegram webhook: no webhook_url configured, skipping setWebhook")
        return

    import httpx

    payload = {"url": webhook_url}
    if _webhook_secret:
        payload["secret_token"] = _webhook_secret

    try:
        async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client:
            resp = await client.post(
                f"https://api.telegram.org/bot{token}/setWebhook",
                json=payload,
            )
            data = resp.json() if resp.content else {}
            if data.get("ok"):
                logger.info(
                    "Telegram webhook set for account %s at %s",
                    _webhook_account_id, webhook_url,
                )
            else:
                logger.error(
                    "Telegram setWebhook failed for account %s: %s",
                    _webhook_account_id, data.get("description", "unknown"),
                )
    except Exception:
        logger.exception("Telegram setWebhook network error for account %s", _webhook_account_id)


async def stop_webhook(account: dict) -> None:
    global _webhook_app, _webhook_queue, _webhook_account_id, _webhook_secret

    token = account.get("token", "")
    if not token:
        _webhook_queue = None
        _webhook_secret = ""
        return

    import httpx

    try:
        async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client:
            resp = await client.post(
                f"https://api.telegram.org/bot{token}/deleteWebhook",
                json={"drop_pending_updates": False},
            )
            data = resp.json() if resp.content else {}
            if data.get("ok"):
                logger.info("Telegram webhook deleted for account %s", _webhook_account_id)
    except Exception:
        logger.exception("Telegram deleteWebhook error for account %s", _webhook_account_id)

    _webhook_queue = None
    _webhook_secret = ""
    _webhook_token = ""

    _webhook_registry.pop(_webhook_account_id, None)


def verify_webhook_secret(provided: str, expected: str) -> bool:
    if not expected:
        return True
    return hmac.compare_digest(provided or "", expected)


def generate_webhook_secret() -> str:
    return secrets.token_urlsafe(32)


async def get_webhook_info(token: str) -> dict | None:
    import httpx

    try:
        async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client:
            resp = await client.get(
                f"https://api.telegram.org/bot{token}/getWebhookInfo",
            )
            data = resp.json() if resp.content else {}
            if data.get("ok"):
                return data.get("result", {})
    except Exception:
        logger.exception("Telegram getWebhookInfo failed")
    return None