scripts/telegram_webhook.py

"""Safely activate, register, inspect, or remove the Telegram webhook.

The ``activate`` action is intended for a trusted workstation.  It prompts
with :func:`getpass.getpass` when the bot token is not inherited by the
process, and never writes that token or generated secrets to disk.
"""
from __future__ import annotations

import argparse
import getpass
import json
import os
import secrets
import sys
import time
from typing import Any

import requests

REPO_ID = "Really-amin/SimpleChatbot"
PUBLIC_BASE = "https://really-amin-simplechatbot.hf.space"
WEBHOOK_PATH = "/api/telegram/webhook"
ALLOWED_UPDATES = ["message", "callback_query"]


def _token() -> str:
    value = os.environ.get("TELEGRAM_BOT_TOKEN", "").strip()
    if value:
        return value
    return getpass.getpass("Telegram Bot Token: ").strip()


def _call(token: str, method: str, payload: dict[str, Any] | None = None) -> dict[str, Any]:
    """Call Telegram without ever including the token in diagnostics."""
    try:
        response = requests.post(
            f"https://api.telegram.org/bot{token}/{method}",
            json=payload or {},
            timeout=(5, 15),
        )
        try:
            body = response.json()
        except ValueError:
            body = {"ok": False, "description": "invalid Telegram response"}
        if not isinstance(body, dict):
            return {"ok": False, "description": "invalid Telegram response"}
        body["_status"] = response.status_code
        return body
    except requests.RequestException:
        return {"ok": False, "description": "Telegram request failed"}


def _safe_info(body: dict[str, Any]) -> dict[str, Any]:
    result = body.get("result") if isinstance(body.get("result"), dict) else {}
    return {
        "ok": bool(body.get("ok")),
        "status": body.get("_status"),
        "url_configured": bool(result.get("url")),
        "url_path": (str(result.get("url", "")).split(".hf.space", 1)[-1] if result.get("url") else ""),
        "pending_update_count": result.get("pending_update_count", 0),
        "allowed_updates": result.get("allowed_updates", []),
        "last_error": bool(result.get("last_error_message")),
        "description": body.get("description", ""),
    }


def _hf_api():
    from huggingface_hub import HfApi
    token = os.environ.get("HF_TOKEN", "").strip() or None
    return HfApi(token=token)


def _wait_running(api: Any, timeout: int = 900) -> bool:
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            runtime = api.get_space_runtime(REPO_ID)
            stage = str(getattr(runtime, "stage", "") or "").upper()
            if stage == "RUNNING":
                return True
            if stage in {"RUNTIME_ERROR", "BUILD_ERROR", "ERROR"}:
                return False
        except Exception:
            pass
        time.sleep(10)
    return False


def _configure_space(api: Any, webhook_secret: str, bootstrap_secret: str, bot_token: str) -> None:
    # The token is updated in the Space Secret but is never printed or saved.
    api.add_space_secret(REPO_ID, "TELEGRAM_BOT_TOKEN", bot_token)
    api.add_space_secret(REPO_ID, "TELEGRAM_WEBHOOK_SECRET", webhook_secret)
    api.add_space_secret(REPO_ID, "TELEGRAM_BOOTSTRAP_SECRET", bootstrap_secret)
    for key, value in {
        "TELEGRAM_ENABLED": "true",
        "TELEGRAM_MODE": "webhook",
        "TELEGRAM_PUBLIC_BASE_URL": PUBLIC_BASE,
        "TELEGRAM_WEBHOOK_PATH": WEBHOOK_PATH,
    }.items():
        api.add_space_variable(REPO_ID, key, value)


def _activate(args: argparse.Namespace, bot_token: str) -> int:
    if not bot_token:
        print("A Telegram Bot Token is required.", file=sys.stderr)
        return 2
    webhook_secret = secrets.token_urlsafe(32)
    bootstrap_secret = secrets.token_urlsafe(32)
    try:
        api = _hf_api()
        _configure_space(api, webhook_secret, bootstrap_secret, bot_token)
        print("Space secrets and webhook-mode variables configured.")
        if not _wait_running(api):
            print("Space did not reach RUNNING after configuration.", file=sys.stderr)
            return 1
        me = _call(bot_token, "getMe")
        if not me.get("ok"):
            print("getMe failed; activation stopped.", file=sys.stderr)
            return 1
        user = me.get("result") if isinstance(me.get("result"), dict) else {}
        username = str(user.get("username") or "(username unavailable)")
        webhook = _call(bot_token, "setWebhook", {
            "url": PUBLIC_BASE + WEBHOOK_PATH,
            "secret_token": webhook_secret,
            "allowed_updates": ALLOWED_UPDATES,
            "drop_pending_updates": bool(args.drop_pending),
        })
        if not webhook.get("ok"):
            print("setWebhook failed; activation stopped.", file=sys.stderr)
            return 1
        info = _call(bot_token, "getWebhookInfo")
        print(json.dumps({"bot_username": username, "webhook": _safe_info(info)}, sort_keys=True))
        print(f"One-time owner claim: /claim {bootstrap_secret}")
        print("Send that command in a private chat, then press Enter here.")
        input()
        # The bootstrap secret is consumed by the Space after a successful claim.
        status_url = PUBLIC_BASE + "/api/telegram/bootstrap/status"
        claimed = False
        for _ in range(30):
            try:
                response = requests.get(status_url, headers={"X-Telegram-Bot-Api-Secret-Token": webhook_secret}, timeout=(5, 10))
                data = response.json() if response.headers.get("content-type", "").startswith("application/json") else {}
                claimed = bool(isinstance(data, dict) and data.get("ownerClaimed"))
            except requests.RequestException:
                pass
            if claimed:
                break
            time.sleep(2)
        if claimed:
            try:
                api.delete_space_secret(REPO_ID, "TELEGRAM_BOOTSTRAP_SECRET")
            except Exception:
                print("Owner claimed; remove the bootstrap secret from Space settings manually.")
            print("Owner claim confirmed; one-time bootstrap secret removed.")
        else:
            print("Owner claim was not confirmed before timeout; bootstrap remains available for one valid claim.")
        return 0
    finally:
        # Remove references before returning; no token is written by this tool.
        bot_token = ""
        webhook_secret = ""
        bootstrap_secret = ""


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("action", choices=("activate", "register", "info", "remove"))
    parser.add_argument("--drop-pending", action="store_true")
    args = parser.parse_args()
    if args.action == "activate":
        token = _token()
        if not token:
            print("A Telegram Bot Token is required.", file=sys.stderr)
            return 2
        try:
            return _activate(args, token)
        finally:
            token = ""
    token = _token()
    if not token:
        print("A Telegram Bot Token is required.", file=sys.stderr)
        return 2
    try:
        if args.action == "info":
            print(json.dumps(_safe_info(_call(token, "getWebhookInfo")), sort_keys=True)); return 0
        if args.action == "remove":
            body = _call(token, "deleteWebhook", {"drop_pending_updates": args.drop_pending})
            print(json.dumps({"ok": bool(body.get("ok")), "status": body.get("_status"), "description": body.get("description", "")}, sort_keys=True)); return 0 if body.get("ok") else 1
        public = os.environ.get("TELEGRAM_PUBLIC_BASE_URL", "").rstrip("/") or PUBLIC_BASE
        path = os.environ.get("TELEGRAM_WEBHOOK_PATH", WEBHOOK_PATH)
        secret = os.environ.get("TELEGRAM_WEBHOOK_SECRET", "")
        if not secret:
            print("TELEGRAM_WEBHOOK_SECRET is required for register.", file=sys.stderr); return 2
        body = _call(token, "setWebhook", {"url": public + path, "secret_token": secret, "allowed_updates": ALLOWED_UPDATES, "drop_pending_updates": args.drop_pending})
        print(json.dumps({"ok": bool(body.get("ok")), "status": body.get("_status"), "description": body.get("description", "")}, sort_keys=True)); return 0 if body.get("ok") else 1
    finally:
        token = ""


if __name__ == "__main__":
    raise SystemExit(main())