diff --git a/mybot/bot.py b/mybot/bot.py index d4f7b51..478c4f6 100644 --- a/mybot/bot.py +++ b/mybot/bot.py @@ -3,6 +3,7 @@ from telebot.storage import StateMemoryStorage, StateRedisStorage from .handlers import register_handlers from .middlewares import setup_middlewares +from .filters import add_custom_filters def create_bot(config, logger, i18n): @@ -29,4 +30,7 @@ def create_bot(config, logger, i18n): logger.debug("Registering handlers") register_handlers(bot) + logger.debug("Adding custom filters") + add_custom_filters(bot, config) + return bot diff --git a/mybot/filters/__init__.py b/mybot/filters/__init__.py index e69de29..377fad7 100644 --- a/mybot/filters/__init__.py +++ b/mybot/filters/__init__.py @@ -0,0 +1,9 @@ +from telebot import TeleBot +from telebot.custom_filters import StateFilter + +from .isowner import IsOwnerFilter + + +def add_custom_filters(bot: TeleBot, config): + bot.add_custom_filter(StateFilter(bot)) + bot.add_custom_filter(IsOwnerFilter(config.OWNER_ID)) diff --git a/mybot/filters/isowner.py b/mybot/filters/isowner.py new file mode 100644 index 0000000..c4f2987 --- /dev/null +++ b/mybot/filters/isowner.py @@ -0,0 +1,12 @@ +from telebot.custom_filters import SimpleCustomFilter +from telebot.types import Message + + +class IsOwnerFilter (SimpleCustomFilter): + key = "is_owner" + + def __init__(self, owner_id: int): + self.owner_id: int = owner_id + + def check(self, message: Message): + return message.from_user.id == self.owner_id