Возможность ограничивать права на бота

This commit is contained in:
mihalin 2021-10-01 19:48:23 +03:00
parent 5fcb5b8900
commit 1aeec0c9d8
5 changed files with 74 additions and 3 deletions

View File

@ -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(
[

View File

@ -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)

View File

@ -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

View File

@ -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()