Feat: disabled menu command

main.py

@@ -20,5 +20,5 @@ if __name__ == "__main__":
 
     # Start bot
     bot = NwXrayBot(config.bot_token.get_secret_value())
-    bot.include_routers(HelloHandler(), MenuHandler())
+    bot.include_routers(HelloHandler(), MenuHandler(), AdminHandler())
     uvloop.run(bot.start(skip_updates=True))

nwxraybot/__init__.py

@@ -1,2 +1,3 @@
 from nwxraybot.bot import NwXrayBot
 from nwxraybot.config import Settings
+from nwxraybot.utils import get_code, get_subscription_info

nwxraybot/bot.py

@@ -1,4 +1,5 @@
 from aiogram import Bot, Dispatcher, Router
+from aiogram.utils.callback_answer import CallbackAnswerMiddleware
 
 from nwxraybot.middlewares import UserMiddleware
 
@@ -10,6 +11,7 @@ class NwXrayBot:
         self.bot = Bot(token=token)
         self.dp = Dispatcher()
         self.dp.message.middleware(UserMiddleware())
+        self.dp.message.middleware(CallbackAnswerMiddleware())
 
     def include_routers(self, *routers: Handler):
         for router in routers:

nwxraybot/handlers/__init__.py

@@ -1,2 +1,3 @@
+from nwxraybot.handlers.admin import AdminHandler
 from nwxraybot.handlers.hello import HelloHandler
 from nwxraybot.handlers.menu import MenuHandler

nwxraybot/handlers/admin.py

@@ -0,0 +1,53 @@
+import re
+from datetime import datetime
+
+from aiogram.enums import ParseMode
+from aiogram.filters import Command
+from aiogram.types import Message
+
+from nwxraybot import get_code
+from nwxraybot.meta import Handler
+from nwxraybot.middlewares import AdminMiddleware
+from nwxraybot.models import User
+
+
+class AdminHandler(Handler):
+    def __init__(self) -> None:
+        super().__init__()
+
+        self.router.message.middleware(AdminMiddleware())
+
+        @self.router.message(Command('adduser'))
+        async def add_user(message: Message):
+            mask = r"^(?P<name>[a-zA-Z0-9]+)\s(?P<url>[^\s]+)($|\s(?P<date>[0-9]{2}\.[0-9]{2}\.[0-9]{4})\s(?P<time>[0-9]{2}\:[0-9]{2})$)"
+            text = message.text.replace('/adduser ', '')
+            match = re.match(mask, text)
+            if match is None:
+                await message.reply('Вы ввели команду в неверном формате. Вводите в формате:\n``` /adduser name vless://....  01.01.1970 00:00```', parse_mode=ParseMode.MARKDOWN)
+                return
+            user_dict = match.groupdict()
+            date = None
+            if user_dict['date']:
+                date = datetime.strptime(f"{user_dict['date']} {
+                                         user_dict['time']}", "%d.%m.%Y %H:%M")
+            code = get_code()
+            new_user = User(
+                name=user_dict['name'], url=user_dict['url'], time=date, code=code)
+            new_user.save()
+            await message.answer(f'Пользователь создан. Вот его ссылка для доступа:\n`https://t.me/nwproxybot?start={code}`', parse_mode=ParseMode.MARKDOWN)
+
+        @self.router.message(Command('updateuser'))
+        async def update_user(message: Message):
+            mask = r"^(?P<name>[a-zA-Z0-9]+)\s(?P<date>[0-9]{2}\.[0-9]{2}\.[0-9]{4})\s(?P<time>[0-9]{2}\:[0-9]{2})$"
+            text = message.text.replace('/updateuser ', '')
+            match = re.match(mask, text)
+            if match is None:
+                await message.reply('Вы ввели команду в неверном формате. Вводите в формате:\n``` /updateuser name 01.01.1970 00:00```', parse_mode=ParseMode.MARKDOWN)
+                return
+            user_dict = match.groupdict()
+            date = datetime.strptime(f"{user_dict['date']} {
+                user_dict['time']}", "%d.%m.%Y %H:%M")
+            query = User.update(time=date).where(
+                User.name == user_dict['name'])
+            query.execute()
+            await message.answer('Информация о пользователе обновлена.')

nwxraybot/handlers/hello.py

@@ -1,5 +1,4 @@
-import json
-from datetime import datetime
+from typing import Optional
 
 from aiogram import F, types
 from aiogram.enums import ParseMode
@@ -7,6 +6,7 @@ from aiogram.filters import Command
 from aiogram.types import (CallbackQuery, InlineKeyboardButton,
                            InlineKeyboardMarkup)
 
+from nwxraybot import get_subscription_info
 from nwxraybot.meta import Handler
 from nwxraybot.models import User
 
@@ -14,7 +14,7 @@ from nwxraybot.models import User
 class HelloHandler(Handler):
     def __non_admin_main_menu(self):
         markup = [[
-            InlineKeyboardButton(text="Обновить данные",
+            InlineKeyboardButton(text="Обновить данные🔄",
                                  callback_data='update')
         ]]
         return InlineKeyboardMarkup(
@@ -26,8 +26,23 @@ class HelloHandler(Handler):
 
         @self.router.message(Command("start"))
         async def hello(message: types.Message):
-            await message.reply("Приветствуем в боте NwXray! Здесь вы сможете получить информацию о своем подключении к NwXray", reply_markup=self.__non_admin_main_menu())
+            data = message.text.split()
+            if len(data) == 2:
+                code = data[1]
+                query = User.update(telegram_id=None).where(
+                    User.telegram_id == message.from_user.id)
+                query.execute()
+                user: Optional[User] = User.select().where(
+                    User.code == code).first()
+                if user is None:
+                    await message.answer('Пользователь не найден, обратитесь к администратору за ссылкой!')
+                    return
+                user.telegram_id = message.from_user.id
+                user.save()
+            await message.answer(f"Приветствуем в боте NwXray! Здесь вы сможете получить информацию о своем подключении к NwXray.\n\n{get_subscription_info(message.from_user.id)}",
+                                 reply_markup=self.__non_admin_main_menu(), parse_mode=ParseMode.MARKDOWN)
 
         @self.router.callback_query(F.data == 'update')
         async def update_data(callback: CallbackQuery):
-            await callback.message.delete()
+            await callback.message.edit_text(get_subscription_info(callback.from_user.id),
+                                             reply_markup=self.__non_admin_main_menu(), parse_mode=ParseMode.MARKDOWN)

nwxraybot/handlers/menu.py

@@ -1,5 +1,6 @@
 from aiogram import types
 from aiogram.enums import ParseMode
+from aiogram.filters import Command
 
 from nwxraybot.meta import Handler
 from nwxraybot.models import User
@@ -9,7 +10,8 @@ class MenuHandler(Handler):
     def __init__(self) -> None:
         super().__init__()
 
-        @self.router.message()
+        # TODO: Temporary disabled
+        # @self.router.message(Command('menu'))
         async def menu(message: types.Message) -> None:
             user: User = User.select().where(User.id == message.from_user.id).first()
             if user:

nwxraybot/middlewares/__init__.py

@@ -1 +1,2 @@
+from nwxraybot.middlewares.admin import AdminMiddleware
 from nwxraybot.middlewares.user import UserMiddleware

nwxraybot/middlewares/admin.py

@@ -0,0 +1,19 @@
+from typing import Any, Awaitable, Callable, Dict, Optional
+
+from aiogram import BaseMiddleware
+from aiogram.types import Message
+
+from nwxraybot.models import User
+
+
+class AdminMiddleware(BaseMiddleware):
+    def __init__(self) -> None:
+        pass
+
+    async def __call__(self, handler: Callable[[Message, Dict[str, Any]], Awaitable[Any]], event: Message, data: Dict[str, Any]) -> Any:
+        user: User = User.select().where(
+            User.telegram_id == event.from_user.id).first()
+        if user is None or not user.admin:
+            await event.reply('Вы не обладаете правами администратора для доступа к данной команде.')
+            return None
+        return await handler(event, data)

nwxraybot/middlewares/user.py

@@ -14,8 +14,10 @@ class UserMiddleware(BaseMiddleware):
     async def __call__(self, handler: Callable[[Message, Dict[str, Any]], Awaitable[Any]], event: Message, data: Dict[str, Any]) -> Any:
         if event.chat.type != ChatType.PRIVATE:
             return None
+        if event.text.startswith('/start'):
+            return await handler(event, data)
         user: Optional[User] = User.select().where(
-            User.id == event.from_user.id).first()
+            User.telegram_id == event.from_user.id).first()
         if user is None:
             await event.answer("Вы не зарегистрированы в системе, обратитесь к админам за доступом!")
             return None

nwxraybot/models/user.py

@@ -13,6 +13,8 @@ class User(Model):
     url = CharField()
     time = DateTimeField(null=True)
     admin = BooleanField(default=False)
+    telegram_id = CharField(null=True)
+    code = CharField()
 
     class Meta:
         database = db

nwxraybot/utils.py

@@ -0,0 +1,22 @@
+import logging
+from datetime import datetime
+from secrets import token_urlsafe
+from typing import Optional
+
+from nwxraybot.models import User
+
+
+def get_subscription_info(telegram_id: str) -> str:
+    user: User = User.select().where(User.telegram_id == telegram_id).first()
+    if user is None:
+        logging.error("[get_subscription_info]: User is not found")
+        return "Ошибка\!"
+    date: Optional[datetime] = user.time
+    date_str = "" if date is None else f'До: {
+        date.strftime("%d.%m.%Y %H:%M")} МСК\n'
+    res = f"Информация о подписке:\n{date_str}Ссылка: `{user.url}`"
+    return res
+
+
+def get_code(length: int = 10) -> str:
+    return token_urlsafe(length)[:length]