Добавлены миддлвари на примере ThrottlingMiddleware

This commit is contained in:
Forden 2020-04-06 00:58:24 +03:00
parent 90791c3091
commit c3da772339
9 changed files with 96 additions and 2 deletions

2
bot.py
View File

@ -5,6 +5,7 @@ from loguru import logger
import filters import filters
import handlers import handlers
import middlewares
from data import config from data import config
bot = Bot(token=config.BOT_TOKEN, parse_mode=types.ParseMode.HTML) bot = Bot(token=config.BOT_TOKEN, parse_mode=types.ParseMode.HTML)
@ -15,6 +16,7 @@ dp = Dispatcher(bot, storage=storage)
# noinspection PyUnusedLocal # noinspection PyUnusedLocal
async def on_startup(web_app: web.Application): async def on_startup(web_app: web.Application):
filters.setup(dp) filters.setup(dp)
middlewares.setup(dp)
handlers.user.setup(dp) handlers.user.setup(dp)
await dp.bot.delete_webhook() await dp.bot.delete_webhook()
await dp.bot.set_webhook(config.WEBHOOK_URL) await dp.bot.set_webhook(config.WEBHOOK_URL)

View File

@ -1,8 +1,10 @@
from aiogram import Dispatcher from aiogram import Dispatcher
from aiogram.dispatcher.filters import CommandStart from aiogram.dispatcher.filters import CommandStart, CommandHelp
from .help import bot_help
from .start import bot_start from .start import bot_start
def setup(dp: Dispatcher): def setup(dp: Dispatcher):
dp.register_message_handler(start, CommandStart()) dp.register_message_handler(bot_start, CommandStart())
dp.register_message_handler(bot_help, CommandHelp())

13
handlers/user/help.py Normal file
View File

@ -0,0 +1,13 @@
from aiogram import types
from utils.misc import rate_limit
@rate_limit(5, 'help')
async def bot_help(msg: types.Message):
text = [
'Список команд: ',
'/start - Начать диалог',
'/help - Получить справку'
]
await msg.answer('\n'.join(text))

7
middlewares/__init__.py Normal file
View 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
View 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.')

View File

@ -0,0 +1 @@
from . import user

View File

@ -0,0 +1,3 @@
from . import db_api
from . import misc
from . import redis

1
utils/misc/__init__.py Normal file
View File

@ -0,0 +1 @@
from .throttling import rate_limit

16
utils/misc/throttling.py Normal file
View File

@ -0,0 +1,16 @@
def rate_limit(limit: int, key=None):
"""
Decorator for configuring rate limit and key in different functions.
:param limit:
:param key:
:return:
"""
def decorator(func):
setattr(func, 'throttling_rate_limit', limit)
if key:
setattr(func, 'throttling_key', key)
return func
return decorator