From 549a40e8afe579d2448809ef566e91f84e1440f9 Mon Sep 17 00:00:00 2001 From: xtaodada Date: Sun, 23 Jun 2024 00:27:16 +0800 Subject: [PATCH] :sparkles: Support webhook --- application.py | 61 +++++++++++++++++++++++++++++-------------- config.py | 32 ++++++++++++++++------- dependence/mtproto.py | 2 +- ratelimiter.py | 3 +++ 4 files changed, 68 insertions(+), 30 deletions(-) diff --git a/application.py b/application.py index 5764f83..1691480 100644 --- a/application.py +++ b/application.py @@ -10,6 +10,8 @@ from typing import Callable, List, Optional, TYPE_CHECKING, TypeVar, Union import pytz import uvicorn from fastapi import FastAPI +from starlette.requests import Request +from starlette.responses import Response from telegram import Bot, Update from telegram.error import NetworkError, TelegramError, TimedOut from telegram.ext import ( @@ -32,8 +34,8 @@ from utils.models.signal import Singleton if TYPE_CHECKING: from asyncio import Task + from telegram import Bot from types import FrameType - from uvicorn._types import ASGIApplication from gram_core.ratelimiter import T_CalledAPIFunc from gram_core.handler.hookhandler import T_PreprocessorsFunc @@ -73,9 +75,9 @@ class Application(Singleton): .get_updates_connect_timeout(application_config.update_connect_timeout) .get_updates_pool_timeout(application_config.update_pool_timeout) .defaults(Defaults(tzinfo=pytz.timezone("Asia/Shanghai"), allow_sending_without_reply=True)) - .token(application_config.bot_token) - .base_url(application_config.bot_base_url) - .base_file_url(application_config.bot_base_file_url) + .token(application_config.bot.token) + .base_url(application_config.bot.base_url) + .base_file_url(application_config.bot.base_file_url) .request( HTTPXRequest( connection_pool_size=application_config.connection_pool_size, @@ -109,7 +111,7 @@ class Application(Singleton): return self._running @property - def web_app(self) -> Union["ASGIApplication", Callable, str]: + def web_app(self) -> Union["FastAPI", Callable, str]: """fastapi app""" return self.web_server.config.app @@ -147,10 +149,6 @@ class Application(Singleton): """启动 BOT""" logger.info("正在启动 BOT 中...") - def error_callback(exc: TelegramError) -> None: - """错误信息回调""" - self.telegram.create_task(self.telegram.process_error(error=exc, update=None)) - await self.telegram.initialize() logger.info("[blue]Telegram[/] 初始化成功", extra={"markup": True}) @@ -175,11 +173,36 @@ class Application(Singleton): self._web_server_task = asyncio.create_task(self.web_server.main_loop()) + await self.start_bot() + + await self.initialize() + logger.success("BOT 初始化成功") + logger.debug("BOT 开始启动") + + await self._on_startup() + await self.telegram.start() + self._running = True + logger.success("BOT 启动成功") + + async def start_bot(self): + """启动 BOT""" + + def error_callback(exc: TelegramError) -> None: + """错误信息回调""" + self.telegram.create_task(self.telegram.process_error(error=exc, update=None)) + + if application_config.bot.is_webhook and application_config.webserver.enable: + self.register_bot_route() + await self.bot.set_webhook(application_config.bot.webhook_url) for _ in range(5): # 连接至 telegram 服务器 try: - await self.telegram.updater.start_polling( - error_callback=error_callback, allowed_updates=Update.ALL_TYPES - ) + if application_config.bot.is_webhook and application_config.webserver.enable: + await self.bot.set_webhook(application_config.bot.webhook_url) + else: + await self.bot.delete_webhook() + await self.telegram.updater.start_polling( + error_callback=error_callback, allowed_updates=Update.ALL_TYPES + ) break except TimedOut: logger.warning("连接至 [blue]telegram[/] 服务器失败，正在重试", extra={"markup": True}) @@ -192,14 +215,14 @@ class Application(Singleton): logger.error("网络连接出现问题, 请检查您的网络状况.") raise SystemExit from e - await self.initialize() - logger.success("BOT 初始化成功") - logger.debug("BOT 开始启动") + def register_bot_route(self): + """注册 webhook 路由""" - await self._on_startup() - await self.telegram.start() - self._running = True - logger.success("BOT 启动成功") + @self.web_app.post("/telegram") + async def telegram(request: Request) -> Response: + """Handle incoming Telegram updates by putting them into the `update_queue`""" + await self.telegram.updater.update_queue.put(Update.de_json(data=await request.json(), bot=self.bot)) + return Response() def stop_signal_handler(self, signum: int): """终止信号处理""" diff --git a/config.py b/config.py index 5c33d3e..da8039d 100644 --- a/config.py +++ b/config.py @@ -1,6 +1,6 @@ from enum import Enum from pathlib import Path -from typing import List, Optional, Union, Set +from typing import List, Optional, Set import dotenv from pydantic import AnyUrl, Field @@ -121,6 +121,26 @@ class NoticeConfig(Settings): env_prefix = "notice_" +class BotConfig(Settings): + """Bot 基础设置""" + + token: str = "" + """BOT的token""" + base_url: str = "https://api.telegram.org/bot" + """Telegram API URL""" + base_file_url: str = "https://api.telegram.org/file/bot" + """Telegram API File URL""" + official: List[str] = ["PaimonMasterBot", "HonkaiStarRail_ZH_Bot"] + """PaiGramTeam Bot""" + is_webhook: bool = False + """接收更新类型 1 为 webhook 0 为 轮询 pull""" + webhook_url: str = "http://127.0.0.1:8080/telegram" + """webhook url""" + + class Config(Settings.Config): + env_prefix = "bot_" + + class ApplicationConfig(Settings): debug: bool = False """debug 开关""" @@ -132,15 +152,6 @@ class ApplicationConfig(Settings): proxy_url: Optional[AnyUrl] = None """代理链接""" - bot_token: str = "" - """BOT的token""" - bot_base_url: str = "https://api.telegram.org/bot" - """Telegram API URL""" - bot_base_file_url: str = "https://api.telegram.org/file/bot" - """Telegram API File URL""" - bot_official: List[str] = ["PaimonMasterBot", "HonkaiStarRail_ZH_Bot"] - """PaiGramTeam Bot""" - owner: Optional[int] = None channels: List[int] = [] @@ -180,6 +191,7 @@ class ApplicationConfig(Settings): mtproto: MTProtoConfig = MTProtoConfig() error: ErrorConfig = ErrorConfig() notice: NoticeConfig = NoticeConfig() + bot: BotConfig = BotConfig() ApplicationConfig.update_forward_refs() diff --git a/dependence/mtproto.py b/dependence/mtproto.py index eadbb29..f580fff 100644 --- a/dependence/mtproto.py +++ b/dependence/mtproto.py @@ -57,7 +57,7 @@ class MTProto(BaseService.Dependence): api_id=bot_config.mtproto.api_id, api_hash=bot_config.mtproto.api_hash, name=self.name, - bot_token=bot_config.bot_token, + bot_token=bot_config.bot.token, proxy=self.proxy, ) await self.client.start() diff --git a/ratelimiter.py b/ratelimiter.py index eaf599d..3822dc1 100644 --- a/ratelimiter.py +++ b/ratelimiter.py @@ -60,6 +60,9 @@ class RateLimiter(BaseRateLimiter[int]): await self._on_called_api(endpoint, data, result) return result except RetryAfter as exc: + if endpoint == "setWebhook" and exc.retry_after == 1: + # webhook 已被正确设置 + return True logger.warning("chat_id[%s] 触发洪水限制 当前被服务器限制 retry_after[%s]秒", chat_id, exc.retry_after) self._limiter_info[chat_id] = time + (exc.retry_after * 2) sleep = exc.retry_after + 0.1