mirror of
https://github.com/PaiGramTeam/PaiGram.git
synced 2024-11-23 00:01:45 +00:00
8f424bf0d4
♻️ 重构插件系统 ⚙️ 重写插件 🎨 改进代码结构 📝 完善文档 Co-authored-by: zhxy-CN <admin@owo.cab> Co-authored-by: 洛水居室 <luoshuijs@outlook.com> Co-authored-by: xtaodada <xtao@xtaolink.cn> Co-authored-by: Li Chuangbo <im@chuangbo.li>
93 lines
3.4 KiB
Python
93 lines
3.4 KiB
Python
import time
|
||
from functools import wraps
|
||
from typing import Callable, cast
|
||
|
||
from telegram import Update
|
||
from telegram.error import TelegramError
|
||
from telegram.ext import filters, CallbackContext
|
||
|
||
from utils.log import logger
|
||
|
||
|
||
def restricts(filters_chat: filters = filters.ALL, return_data=None, try_delete_message: bool = False,
|
||
restricts_time: int = 5):
|
||
"""用于装饰在指定函数预防洪水攻击的装饰器
|
||
|
||
被修饰的函数生声明必须为
|
||
|
||
async def command_func(update, context)
|
||
或
|
||
async def command_func(self, update, context
|
||
|
||
如果修饰的函数属于
|
||
ConversationHandler
|
||
参数
|
||
return_data
|
||
必须传入
|
||
ConversationHandler.END
|
||
|
||
我真™是服了某些闲着没事干的群友了
|
||
|
||
:param filters_chat: 要限制的群
|
||
:param return_data:
|
||
:param try_delete_message:
|
||
:param restricts_time:
|
||
:return: return_data
|
||
"""
|
||
|
||
def decorator(func: Callable):
|
||
@wraps(func)
|
||
async def restricts_func(*args, **kwargs):
|
||
if len(args) == 3:
|
||
# self update context
|
||
_, update, context = args
|
||
elif len(args) == 2:
|
||
# update context
|
||
update, context = args
|
||
else:
|
||
return await func(*args, **kwargs)
|
||
update = cast(Update, update)
|
||
context = cast(CallbackContext, context)
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if filters_chat.filter(message):
|
||
command_time = context.user_data.get("command_time", 0)
|
||
count = context.user_data.get("usage_count", 0)
|
||
restrict_since = context.user_data.get("restrict_since", 0)
|
||
# 洪水防御
|
||
if restrict_since:
|
||
if (time.time() - restrict_since) >= 60 * 5:
|
||
del context.user_data["restrict_since"]
|
||
del context.user_data["usage_count"]
|
||
else:
|
||
return return_data
|
||
else:
|
||
if count == 5:
|
||
context.user_data["restrict_since"] = time.time()
|
||
await message.reply_text("你已经触发洪水防御，请等待5分钟")
|
||
logger.warning(f"用户 {user.full_name}[{user.id}] 触发洪水限制 已被限制5分钟")
|
||
return return_data
|
||
# 单次使用限制
|
||
if command_time:
|
||
if (time.time() - command_time) <= restricts_time:
|
||
context.user_data["usage_count"] = count + 1
|
||
if filters.ChatType.GROUPS.filter(message):
|
||
if try_delete_message:
|
||
try:
|
||
await message.delete()
|
||
except TelegramError as exc:
|
||
logger.warning("删除消息失败")
|
||
logger.exception(exc)
|
||
return return_data
|
||
else:
|
||
if count >= 1:
|
||
context.user_data["usage_count"] = count - 1
|
||
|
||
context.user_data["command_time"] = time.time()
|
||
|
||
return await func(*args, **kwargs)
|
||
|
||
return restricts_func
|
||
|
||
return decorator
|