重构机器人

bot/middlewares/__init__.py

@@ -0,0 +1,7 @@
+"""
+Bot 中间件模块
+"""
+
+from bot.middlewares.auth import require_auth, AuthMiddleware
+
+__all__ = ["require_auth", "AuthMiddleware"]

bot/middlewares/auth.py

@@ -0,0 +1,111 @@
+"""
+用户鉴权中间件
+
+检查用户是否在白名单中
+"""
+
+from functools import wraps
+from typing import Callable, Any
+
+from telegram import Update
+from telegram.ext import ContextTypes
+
+
+class AuthMiddleware:
+    """
+    鉴权中间件类
+    """
+
+    def __init__(self, allowed_users: set, admin_users: set = None):
+        """
+        初始化鉴权中间件
+
+        参数:
+            allowed_users: 允许使用的用户 ID 集合
+            admin_users: 管理员用户 ID 集合
+        """
+        self.allowed_users = allowed_users or set()
+        self.admin_users = admin_users or set()
+
+    def is_allowed(self, user_id: int) -> bool:
+        """检查用户是否允许使用"""
+        # 如果白名单为空，允许所有人
+        if not self.allowed_users:
+            return True
+        return user_id in self.allowed_users or user_id in self.admin_users
+
+    def is_admin(self, user_id: int) -> bool:
+        """检查用户是否是管理员"""
+        return user_id in self.admin_users
+
+
+def require_auth(func: Callable) -> Callable:
+    """
+    鉴权装饰器
+
+    用于命令处理函数，检查用户是否在白名单中
+
+    用法:
+        @require_auth
+        async def my_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
+            ...
+    """
+
+    @wraps(func)
+    async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE, *args, **kwargs) -> Any:
+        user = update.effective_user
+        if not user:
+            return
+
+        user_id = user.id
+        allowed_users = context.bot_data.get("allowed_users", set())
+        admin_users = context.bot_data.get("admin_users", set())
+
+        # 如果白名单为空，允许所有人使用
+        if not allowed_users and not admin_users:
+            return await func(update, context, *args, **kwargs)
+
+        # 检查是否在白名单或管理员列表中
+        if user_id in allowed_users or user_id in admin_users:
+            return await func(update, context, *args, **kwargs)
+
+        # 未授权
+        await update.message.reply_text(
+            f"⛔ 你没有权限使用此 Bot\n\n"
+            f"你的用户 ID: `{user_id}`\n\n"
+            f"请联系管理员将你的 ID 添加到白名单",
+            parse_mode="Markdown"
+        )
+        return None
+
+    return wrapper
+
+
+def require_admin(func: Callable) -> Callable:
+    """
+    管理员权限装饰器
+
+    用于需要管理员权限的命令
+
+    用法:
+        @require_admin
+        async def admin_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
+            ...
+    """
+
+    @wraps(func)
+    async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE, *args, **kwargs) -> Any:
+        user = update.effective_user
+        if not user:
+            return
+
+        user_id = user.id
+        admin_users = context.bot_data.get("admin_users", set())
+
+        if user_id in admin_users:
+            return await func(update, context, *args, **kwargs)
+
+        await update.message.reply_text("⛔ 此命令需要管理员权限")
+        return None
+
+    return wrapper