112 lines
3.0 KiB
Python
112 lines
3.0 KiB
Python
"""
|
|
用户鉴权中间件
|
|
|
|
检查用户是否在白名单中
|
|
"""
|
|
|
|
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
|