Добавлены миддлвари на примере ThrottlingMiddleware
This commit is contained in:
7
middlewares/__init__.py
Normal file
7
middlewares/__init__.py
Normal file
@@ -0,0 +1,7 @@
|
||||
from aiogram import Dispatcher
|
||||
|
||||
from .throttling import ThrottlingMiddleware
|
||||
|
||||
|
||||
def setup(dp: Dispatcher):
|
||||
dp.middleware.setup(ThrottlingMiddleware)
|
||||
49
middlewares/throttling.py
Normal file
49
middlewares/throttling.py
Normal file
@@ -0,0 +1,49 @@
|
||||
import asyncio
|
||||
|
||||
from aiogram import types, Dispatcher
|
||||
from aiogram.dispatcher import DEFAULT_RATE_LIMIT
|
||||
from aiogram.dispatcher.handler import CancelHandler, current_handler
|
||||
from aiogram.dispatcher.middlewares import BaseMiddleware
|
||||
from aiogram.utils.exceptions import Throttled
|
||||
|
||||
|
||||
class ThrottlingMiddleware(BaseMiddleware):
|
||||
"""
|
||||
Simple middleware
|
||||
"""
|
||||
|
||||
def __init__(self, limit=DEFAULT_RATE_LIMIT, key_prefix='antiflood_'):
|
||||
self.rate_limit = limit
|
||||
self.prefix = key_prefix
|
||||
super(ThrottlingMiddleware, self).__init__()
|
||||
|
||||
# noinspection PyUnusedLocal
|
||||
async def on_process_message(self, message: types.Message, data: dict):
|
||||
handler = current_handler.get()
|
||||
dispatcher = Dispatcher.get_current()
|
||||
if handler:
|
||||
limit = getattr(handler, 'throttling_rate_limit', self.rate_limit)
|
||||
key = getattr(handler, 'throttling_key', f"{self.prefix}_{handler.__name__}")
|
||||
else:
|
||||
limit = self.rate_limit
|
||||
key = f"{self.prefix}_message"
|
||||
try:
|
||||
await dispatcher.throttle(key, rate=limit)
|
||||
except Throttled as t:
|
||||
await self.message_throttled(message, t)
|
||||
raise CancelHandler()
|
||||
|
||||
async def message_throttled(self, message: types.Message, throttled: Throttled):
|
||||
handler = current_handler.get()
|
||||
dispatcher = Dispatcher.get_current()
|
||||
if handler:
|
||||
key = getattr(handler, 'throttling_key', f"{self.prefix}_{handler.__name__}")
|
||||
else:
|
||||
key = f"{self.prefix}_message"
|
||||
delta = throttled.rate - throttled.delta
|
||||
if throttled.exceeded_count <= 2:
|
||||
await message.reply('Too many requests! ')
|
||||
await asyncio.sleep(delta)
|
||||
thr = await dispatcher.check_key(key)
|
||||
if thr.exceeded_count == throttled.exceeded_count:
|
||||
await message.reply('Unlocked.')
|
||||
Reference in New Issue
Block a user