diff --git a/.gitignore b/.gitignore index 49a2a2e..2cbd47f 100644 --- a/.gitignore +++ b/.gitignore @@ -6,4 +6,4 @@ __pycache__ *.pyc config.json docker-compose-release.yaml -docs/build \ No newline at end of file +docs/build diff --git a/main.py b/main.py index b2f26e4..e2ded1a 100644 --- a/main.py +++ b/main.py @@ -3,7 +3,8 @@ import argparse from tortoise import Tortoise from olgram.router import dp -from olgram.settings import TORTOISE_ORM +from olgram.settings import TORTOISE_ORM, OlgramSettings +from olgram.utils.permissions import AccessMiddleware from server.custom import init_redis import olgram.commands.bots # noqa: F401 @@ -22,7 +23,8 @@ async def init_database(): async def init_olgram(): - from olgram.router import bot + from olgram.router import bot, dp + dp.setup_middleware(AccessMiddleware(OlgramSettings.admin_id())) from aiogram.types import BotCommand await bot.set_my_commands( [ diff --git a/olgram/commands/start.py b/olgram/commands/start.py index a8c971b..daf5263 100644 --- a/olgram/commands/start.py +++ b/olgram/commands/start.py @@ -6,11 +6,13 @@ from aiogram import types from aiogram.dispatcher import FSMContext from textwrap import dedent from olgram.settings import OlgramSettings +from olgram.utils.permissions import public from olgram.router import dp @dp.message_handler(commands=["start"], state="*") +@public() async def start(message: types.Message, state: FSMContext): """ Команда /start @@ -33,6 +35,7 @@ async def start(message: types.Message, state: FSMContext): @dp.message_handler(commands=["help"], state="*") +@public() async def help(message: types.Message, state: FSMContext): """ Команда /help @@ -42,3 +45,12 @@ async def help(message: types.Message, state: FSMContext): Техническая поддержка: @civsocit_feedback_bot Версия {OlgramSettings.version()} """)) + + +@dp.message_handler(commands=["chatid"], state="*") +@public() +async def chat_id(message: types.Message, state: FSMContext): + """ + Команда /chatid + """ + await message.answer(message.chat.id) diff --git a/olgram/settings.py b/olgram/settings.py index 629d93e..648a1c6 100644 --- a/olgram/settings.py +++ b/olgram/settings.py @@ -32,6 +32,11 @@ class OlgramSettings(AbstractSettings): def version(cls): return "0.1.0" + @classmethod + @lru_cache + def admin_id(cls): + return cls._get_env("ADMIN_ID", True) + class ServerSettings(AbstractSettings): @classmethod diff --git a/olgram/utils/permissions.py b/olgram/utils/permissions.py new file mode 100644 index 0000000..4755ffd --- /dev/null +++ b/olgram/utils/permissions.py @@ -0,0 +1,52 @@ +import aiogram.types as types +from aiogram.dispatcher.handler import CancelHandler, current_handler +from aiogram.dispatcher.middlewares import BaseMiddleware + + +def public(): + """ + Хендлеры с этим декоратором будут обрабатываться даже если пользователь не является владельцем бота + (например, команда /help) + :return: + """ + + def decorator(func): + setattr(func, "access_public", True) + return func + + return decorator + + +class AccessMiddleware(BaseMiddleware): + def __init__(self, access_chat_id: int): + self._access_chat_id = access_chat_id + super(AccessMiddleware, self).__init__() + + @classmethod + def _is_public_command(cls) -> bool: + handler = current_handler.get() + return handler and getattr(handler, "access_public", False) + + async def on_process_message(self, message: types.Message, data: dict): + admin_id = self._access_chat_id + if not admin_id: + return # Администратор бота вообще не указан + + if self._is_public_command(): # Эта команда разрешена всем пользователям + return + + if message.chat.id != admin_id: + await message.answer("Владелец бота ограничил доступ к этому функционалу 😞") + raise CancelHandler() + + async def on_process_callback_query(self, call: types.CallbackQuery, data: dict): + admin_id = self._access_chat_id + if not admin_id: + return # Администратор бота вообще не указан + + if self._is_public_command(): # Эта команда разрешена всем пользователям + return + + if call.message.chat.id != admin_id: + await call.answer("Владелец бота ограничил доступ к этому функционалу😞") + raise CancelHandler()