From f585b82cd60004454ceaa08f3b3233d5e3fec9bf Mon Sep 17 00:00:00 2001
From: Zakhar Izmaylov <ptdev@kedruss.ru>
Date: Sun, 17 Nov 2024 13:11:49 +0300
Subject: [PATCH] Add admin middleware

---
 bot.py                        |  4 ++++
 handlers/admin/admin_panel.py |  2 +-
 handlers/profile.py           |  6 +++++-
 handlers/start.py             | 10 +++++++---
 middlewares/admin.py          | 24 ++++++++++++++++++++++++
 requirements.txt              |  3 ++-
 6 files changed, 43 insertions(+), 6 deletions(-)
 create mode 100644 middlewares/admin.py

diff --git a/bot.py b/bot.py
index 63264af4..ac37e8bf 100644
--- a/bot.py
+++ b/bot.py
@@ -2,6 +2,7 @@ from aiogram import Bot, Dispatcher, Router
 from aiogram.fsm.storage.memory import MemoryStorage
 
 from config import API_TOKEN, CRYPTO_BOT_ENABLE, FREEKASSA_ENABLE, ROBOKASSA_ENABLE, STARS_ENABLE, YOOKASSA_ENABLE
+from middlewares.admin import AdminMiddleware
 from middlewares.database import DatabaseMiddleware
 from middlewares.logging import LoggingMiddleware
 
@@ -47,3 +48,6 @@ dp.callback_query.middleware(LoggingMiddleware())
 
 dp.message.middleware(DatabaseMiddleware())
 dp.callback_query.middleware(DatabaseMiddleware())
+
+dp.message.middleware(AdminMiddleware())
+dp.callback_query.middleware(AdminMiddleware())
diff --git a/handlers/admin/admin_panel.py b/handlers/admin/admin_panel.py
index 9c29a794..787418f9 100644
--- a/handlers/admin/admin_panel.py
+++ b/handlers/admin/admin_panel.py
@@ -23,7 +23,7 @@ class UserEditorState(StatesGroup):
     displaying_user_info = State()
 
 
-@router.message(Command("admin"),F.data == "admin", IsAdminFilter())
+@router.message(Command("admin"), F.data == "admin", IsAdminFilter())
 async def handle_admin_command(message: types.Message):
     builder = InlineKeyboardBuilder()
     builder.row(
diff --git a/handlers/profile.py b/handlers/profile.py
index 7958a5b5..b756ef6d 100644
--- a/handlers/profile.py
+++ b/handlers/profile.py
@@ -15,7 +15,7 @@ router = Router()
 
 
 async def process_callback_view_profile(
-    callback_query: types.CallbackQuery, state: FSMContext
+    callback_query: types.CallbackQuery, state: FSMContext, admin: bool
 ):
     chat_id = callback_query.from_user.id
     username = callback_query.from_user.full_name
@@ -52,6 +52,10 @@ async def process_callback_view_profile(
         builder.row(
             InlineKeyboardButton(text="💰 Поддержать проект", callback_data="donate")
         )
+        if admin:
+            builder.row(
+                InlineKeyboardButton(text="🔧 Администратор", callback_data="admin")
+            )
         builder.row(
             InlineKeyboardButton(text="⬅️ Главное меню", callback_data="back_to_menu")
         )
diff --git a/handlers/start.py b/handlers/start.py
index 24833c99..814e9cae 100644
--- a/handlers/start.py
+++ b/handlers/start.py
@@ -15,7 +15,7 @@ from logger import logger
 router = Router()
 
 
-async def send_welcome_message(chat_id: int, trial_status: int):
+async def send_welcome_message(chat_id: int, trial_status: int, admin: bool):
     image_path = os.path.join(os.path.dirname(__file__), "pic.jpg")
 
     builder = InlineKeyboardBuilder()
@@ -26,6 +26,10 @@ async def send_welcome_message(chat_id: int, trial_status: int):
     builder.row(
         InlineKeyboardButton(text="👤 Личный кабинет", callback_data="view_profile")
     )
+    if admin:
+        builder.row(
+            InlineKeyboardButton(text="🔧 Администратор", callback_data="admin")
+        )
     builder.row(
         InlineKeyboardButton(text="📞 Техническая поддержка", url=SUPPORT_CHAT_URL),
     )
@@ -53,7 +57,7 @@ async def send_welcome_message(chat_id: int, trial_status: int):
 
 
 @router.message(Command("start"))
-async def start_command(message: Message):
+async def start_command(message: Message, admin: bool):
     logger.info(f"Received start command with text: {message.text}")
     if "referral_" in message.text:
         referrer_tg_id = int(message.text.split("referral_")[1])
@@ -66,7 +70,7 @@ async def start_command(message: Message):
             await message.answer("Вы уже зарегистрированы в системе!")
 
     trial_status = await get_trial(message.from_user.id)
-    await send_welcome_message(message.chat.id, trial_status)
+    await send_welcome_message(message.chat.id, trial_status, admin)
 
 
 @router.callback_query(F.data == "connect_vpn")
diff --git a/middlewares/admin.py b/middlewares/admin.py
new file mode 100644
index 00000000..1e532d7c
--- /dev/null
+++ b/middlewares/admin.py
@@ -0,0 +1,24 @@
+from typing import Any, Awaitable, Callable, Dict
+
+from aiogram import BaseMiddleware
+from aiogram.types import TelegramObject
+
+from config import ADMIN_ID
+
+
+class AdminMiddleware(BaseMiddleware):
+    async def __call__(
+        self,
+        handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]],
+        event: TelegramObject,
+        data: Dict[str, Any],
+    ) -> Any:
+        if isinstance(ADMIN_ID, list):
+            if event.from_user.id in ADMIN_ID:
+                data["admin"] = True
+        elif isinstance(ADMIN_ID, int):
+            if event.from_user.id == ADMIN_ID:
+                data["admin"] = True
+        else:
+            data["admin"] = False
+        return await handler(event, data)
diff --git a/requirements.txt b/requirements.txt
index 262fc7fb..e3e0f098 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -28,4 +28,5 @@ yookassa==3.3.0
 loguru
 aiocryptopay
 py3xui
-sqlalchemy
\ No newline at end of file
+sqlalchemy
+robokassa
\ No newline at end of file