diff --git a/mybot/__init__.py b/mybot/__init__.py index 460229a..e04db7a 100644 --- a/mybot/__init__.py +++ b/mybot/__init__.py @@ -27,7 +27,8 @@ def create_bot(config: Config, i18n: I18N, engine): if config.use_webhook: bot.set_webhook(config.webhook.url, drop_pending_updates=config.webhook.drop_pending_updates, - max_connections=config.webhook.max_connections) + max_connections=config.webhook.max_connections, + secret_token=config.webhook.secret_token) return bot diff --git a/mybot/config.py b/mybot/config.py index 398684f..776943b 100644 --- a/mybot/config.py +++ b/mybot/config.py @@ -1,4 +1,5 @@ import os +import secrets from dataclasses import dataclass @@ -27,6 +28,12 @@ class WebhookConfig: url_path: str max_connections: int drop_pending_updates: bool + use_secret_token: bool + secret_token: str + + def __post_init__(self): + if self.use_secret_token and not self.secret_token: + self.secret_token = secrets.token_hex() @property def url(self): @@ -35,9 +42,11 @@ class WebhookConfig: @classmethod def from_env(cls): return cls(os.getenv("WEBHOOK_DOMAIN"), - os.getenv("WEBHOOK_URL_PATH"), + os.getenv("WEBHOOK_URL_PATH", "/"), int(os.getenv("WEBHOOK_MAX_CONNECTIONS", 40)), - bool(int(os.getenv("WEBHOOK_DROP_PENDING", True)))) + bool(int(os.getenv("WEBHOOK_DROP_PENDING", True))), + bool(int(os.getenv("WEBHOOK_USE_SECRET_TOKEN", True))), + os.getenv("WEBHOOK_SECRET_TOKEN")) @dataclass diff --git a/mybot/webhook/__init__.py b/mybot/webhook/__init__.py index 74f7b09..b10e3e8 100644 --- a/mybot/webhook/__init__.py +++ b/mybot/webhook/__init__.py @@ -1,35 +1,35 @@ -from flask import Flask, Blueprint, request, abort, g +from flask import Flask, request, abort, g from telebot import TeleBot from telebot.types import Update from ..config import Config -bot_bp = Blueprint("bot", __name__) - - -@bot_bp.route("/", methods=["GET", "POST"]) def handle_updates(): if request.method == "GET": - abort(404) + abort(404) # safer to 404 + if g.config.webhook.use_secret_token: + if request.headers.get("X-Telegram-Bot-Api-Secret-Token") != g.config.webhook.secret_token: + abort(404) if request.headers.get("content-type") == "application/json": update = Update.de_json(request.get_json()) g.bot.process_new_updates([update]) return "" else: - abort(403) + abort(404) # safer to 404 -def inject_g(bot: TeleBot, config: Config): +def inject_g(**kwargs): def inner(): - g.bot = bot - g.config = config + for k, v in kwargs.items(): + setattr(g, k, v) return inner def create_app(bot: TeleBot, config: Config): app = Flask(__name__) - app.register_blueprint(bot_bp, url_prefix=f"{config.webhook.url_path}") - app.before_request(inject_g(bot, config)) - + app.add_url_rule(config.webhook.url_path, + view_func=handle_updates, + methods=["GET", "POST"]) + app.before_request(inject_g(bot=bot, config=config)) return app