vpn-telegram-bot/services/tg_webapp.py

"""Валидация Telegram WebApp initData, Login Widget и OIDC id_token."""

from __future__ import annotations

import hashlib
import hmac
import json
import logging
import time
from dataclasses import dataclass
from typing import Any, Mapping
from urllib.parse import parse_qsl

import jwt
from jwt import PyJWKClient

logger = logging.getLogger(__name__)

_JWKS_URL = "https://oauth.telegram.org/.well-known/jwks.json"
_ISSUER = "https://oauth.telegram.org"
_jwks_client: PyJWKClient | None = None


@dataclass(frozen=True, slots=True)
class WebAppUser:
    id: int
    first_name: str
    last_name: str | None = None
    username: str | None = None
    language_code: str | None = None
    is_premium: bool = False

    @property
    def full_name(self) -> str:
        if self.last_name:
            return f"{self.first_name} {self.last_name}"
        return self.first_name


class InitDataError(ValueError):
    pass


def telegram_client_id_from_token(bot_token: str) -> str:
    """Числовой Client ID бота = часть BOT_TOKEN до ':'."""
    raw = (bot_token or "").strip()
    if ":" in raw:
        return raw.split(":", 1)[0].strip()
    return raw


def parse_and_validate_init_data(
    init_data: str,
    bot_token: str,
    *,
    max_age_sec: int = 86400,
) -> WebAppUser:
    if not init_data or not init_data.strip():
        raise InitDataError("empty initData")

    pairs = dict(parse_qsl(init_data, keep_blank_values=True))
    received_hash = pairs.pop("hash", None)
    if not received_hash:
        raise InitDataError("missing hash")

    data_check_string = "\n".join(f"{k}={v}" for k, v in sorted(pairs.items()))
    secret_key = hmac.new(b"WebAppData", bot_token.encode(), hashlib.sha256).digest()
    calculated = hmac.new(secret_key, data_check_string.encode(), hashlib.sha256).hexdigest()
    if not hmac.compare_digest(calculated, received_hash):
        raise InitDataError("invalid hash")

    auth_date = int(pairs.get("auth_date") or 0)
    if auth_date and abs(time.time() - auth_date) > max_age_sec:
        raise InitDataError("initData expired")

    user_raw = pairs.get("user")
    if not user_raw:
        raise InitDataError("missing user")
    user = json.loads(user_raw)
    return WebAppUser(
        id=int(user["id"]),
        first_name=str(user.get("first_name") or "User"),
        last_name=user.get("last_name"),
        username=user.get("username"),
        language_code=user.get("language_code"),
        is_premium=bool(user.get("is_premium")),
    )


def parse_and_validate_login_widget(
    data: Mapping[str, Any],
    bot_token: str,
    *,
    max_age_sec: int = 86400,
) -> WebAppUser:
    """Проверка данных legacy Telegram Login Widget (HMAC)."""
    payload = {str(k): str(v) for k, v in data.items() if v is not None and str(v) != ""}
    received_hash = payload.pop("hash", None)
    if not received_hash:
        raise InitDataError("missing hash")

    data_check_string = "\n".join(f"{k}={payload[k]}" for k in sorted(payload.keys()))
    secret_key = hashlib.sha256(bot_token.encode("utf-8")).digest()
    calculated = hmac.new(
        secret_key,
        data_check_string.encode("utf-8"),
        hashlib.sha256,
    ).hexdigest()
    if not hmac.compare_digest(calculated, received_hash):
        raise InitDataError("invalid hash")

    auth_date = int(payload.get("auth_date") or 0)
    if not auth_date or abs(time.time() - auth_date) > max_age_sec:
        raise InitDataError("login data expired")

    if "id" not in payload:
        raise InitDataError("missing id")

    return WebAppUser(
        id=int(payload["id"]),
        first_name=str(payload.get("first_name") or "User"),
        last_name=payload.get("last_name"),
        username=payload.get("username"),
    )


def _get_jwks_client() -> PyJWKClient:
    global _jwks_client
    if _jwks_client is None:
        _jwks_client = PyJWKClient(_JWKS_URL, cache_keys=True, lifespan=3600)
    return _jwks_client


def parse_and_validate_id_token(
    id_token: str,
    client_id: str,
    *,
    nonce: str | None = None,
) -> WebAppUser:
    """Проверка OIDC id_token нового Telegram Login (JWKS)."""
    token = (id_token or "").strip()
    aud = str(client_id or "").strip()
    if not token:
        raise InitDataError("missing id_token")
    if not aud:
        raise InitDataError("telegram client_id not configured")

    try:
        signing_key = _get_jwks_client().get_signing_key_from_jwt(token)
        payload = jwt.decode(
            token,
            signing_key.key,
            algorithms=["RS256", "RS384", "RS512", "PS256", "ES256", "ES384", "EdDSA"],
            audience=aud,
            issuer=_ISSUER,
            options={"require": ["exp", "iat", "iss", "aud"]},
        )
    except jwt.PyJWTError as exc:
        logger.info("telegram id_token invalid: %s", exc)
        raise InitDataError("invalid id_token") from exc

    if nonce and payload.get("nonce") and str(payload.get("nonce")) != str(nonce):
        raise InitDataError("nonce mismatch")

    tid_raw = payload.get("id")
    if tid_raw is None:
        tid_raw = payload.get("sub")
    try:
        tid = int(tid_raw)
    except (TypeError, ValueError) as exc:
        raise InitDataError("missing user id in id_token") from exc

    given = str(payload.get("given_name") or "").strip()
    family = str(payload.get("family_name") or "").strip() or None
    full = str(payload.get("name") or "").strip()
    if given:
        first_name = given
        last_name = family
    elif full:
        parts = full.split(None, 1)
        first_name = parts[0]
        last_name = parts[1] if len(parts) > 1 else family
    else:
        first_name = "User"
        last_name = family

    username = payload.get("preferred_username") or payload.get("username")
    return WebAppUser(
        id=tid,
        first_name=first_name,
        last_name=last_name,
        username=str(username) if username else None,
    )