Add add_user command

4 files changed, 109 insertions, 8 deletions

diff --git a/handlers/admin/__init__.py b/handlers/admin/__init__.py
index 98b127f..02f9e6e 100644
--- a/handlers/admin/__init__.py
+++ b/handlers/admin/__init__.py
@@ -4,6 +4,7 @@ from aiogram.filters import MagicData
 # isort: off
 from . import new_announcement
 from . import new_invoice
+from . import add_user
 # isort: on
 
 router = Router(name="admin")
@@ -13,4 +14,5 @@ router.callback_query.filter(MagicData(F.user.is_admin()))
 router.include_routers(
     new_announcement.router,
     new_invoice.router,
+    add_user.router,
 )
diff --git a/handlers/admin/add_user.py b/handlers/admin/add_user.py
new file mode 100644
index 0000000..b9b8604
--- /dev/null
+++ b/handlers/admin/add_user.py
@@ -0,0 +1,99 @@
+from datetime import UTC, datetime
+
+from aiogram import F, Router
+from aiogram.enums.button_style import ButtonStyle
+from aiogram.filters import Command
+from aiogram.fsm.context import FSMContext
+from aiogram.fsm.state import State, StatesGroup
+from aiogram.types import (
+    KeyboardButton,
+    KeyboardButtonRequestUsers,
+    Message,
+    ReplyKeyboardMarkup,
+    ReplyKeyboardRemove,
+)
+from pydantic import BaseModel
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from libs.fsm import edit_data, get_data, set_data
+from models import User
+
+router = Router(name="add_user")
+
+
+class AddUserStates(StatesGroup):
+    user_id = State()
+    vpn_link = State()
+
+
+class AddUserData(BaseModel):
+    user_id: int | None = None
+
+
+CANCEL_BUTTON = "Отменить добавление"
+
+
+@router.message(Command("add_user"))
+async def add_user_command(msg: Message, state: FSMContext) -> None:
+    await msg.answer(
+        "Выберете пользователя которого хотите добавить.",
+        reply_markup=ReplyKeyboardMarkup(
+            keyboard=[
+                [
+                    KeyboardButton(
+                        text="Выбрать пользователя",
+                        style=ButtonStyle.PRIMARY,
+                        request_users=KeyboardButtonRequestUsers(request_id=0),
+                    ),
+                ],
+                [
+                    KeyboardButton(text=CANCEL_BUTTON, style=ButtonStyle.DANGER),
+                ],
+            ],
+            resize_keyboard=True,
+        ),
+    )
+    await set_data(state, AddUserData(user_id=None))
+    await state.set_state(AddUserStates.user_id)
+
+
+@router.message(AddUserStates(), F.text == CANCEL_BUTTON)
+async def add_user_cancel(msg: Message, state: FSMContext) -> None:
+    await msg.answer(
+        "Добавление пользователей отменено",
+        reply_markup=ReplyKeyboardRemove(),
+    )
+    await state.clear()
+
+
+@router.message(AddUserStates.user_id)
+async def add_user_user_id(msg: Message, state: FSMContext) -> None:
+    if msg.users_shared is None:
+        await msg.answer("Вы должны воспользоваться кнопкой ниже.")
+        return
+
+    async with edit_data(state, AddUserData) as data:
+        data.user_id = msg.users_shared.users[0].user_id
+
+    await msg.answer("Укажите ссылку для доступа к VPN")
+    await state.set_state(AddUserStates.vpn_link)
+
+
+@router.message(AddUserStates.vpn_link)
+async def add_user_vpn_link(
+    msg: Message,
+    state: FSMContext,
+    session: AsyncSession,
+) -> None:
+    if msg.text is None:
+        await msg.answer("Вы должны указать ссылку отправив текстовое сообщение.")
+        return
+
+    data = await get_data(state, AddUserData)
+    assert data.user_id is not None
+
+    session.add(User(id=data.user_id, vpn_link=msg.text, datetime=datetime.now(UTC)))
+    await session.flush()
+
+    await msg.answer("Пользователь добавлен.")
+    await state.clear()
diff --git a/handlers/admin/new_announcement.py b/handlers/admin/new_announcement.py
index 0920c47..6aa099e 100644
--- a/handlers/admin/new_announcement.py
+++ b/handlers/admin/new_announcement.py
@@ -7,7 +7,6 @@ from aiogram.filters import Command
 from aiogram.fsm.context import FSMContext
 from aiogram.fsm.state import State, StatesGroup
 from aiogram.types import (
-    ContentType,
     KeyboardButton,
     Message,
     ReplyKeyboardMarkup,
@@ -67,6 +66,10 @@ async def announcement_send(
         await msg.answer("Для публикации анонса укажите текст сообщения.")
         return
 
+    announcement = Announcement(message=data.rich_text, datetime=datetime.now(UTC))
+    session.add(announcement)
+    await session.flush()
+
     status_template = "Публикация анонса...\nОпубликовано: {}"
     status_msg = await msg.answer(status_template.format(0))
 
@@ -76,9 +79,6 @@ async def announcement_send(
         except TelegramAPIError:
             pass
 
-    announcement = Announcement(message=data, datetime=datetime.now(UTC))
-    session.add(announcement)
-
     await status_msg.delete()
     await msg.answer(
         "Анонс отправлен всем пользователям",
@@ -95,7 +95,7 @@ async def announcement_cancel(msg: Message, state: FSMContext) -> None:
 
 @router.message(NewAnnouncementStates.message)
 async def announcement_message(msg: Message, bot: Bot, state: FSMContext) -> None:
-    if msg.content_type != ContentType.TEXT or msg.text is None:
+    if msg.text is None:
         await msg.answer(
             "Неверный тип сообщения.\n"
             "Бот поддерживает отправку только текстовых анонсов."
diff --git a/handlers/admin/new_invoice.py b/handlers/admin/new_invoice.py
index 7e1a64d..0427d90 100644
--- a/handlers/admin/new_invoice.py
+++ b/handlers/admin/new_invoice.py
@@ -1,7 +1,7 @@
 from datetime import UTC, datetime
 
 from aiogram import Bot, F, Router
-from aiogram.enums import ButtonStyle, ContentType
+from aiogram.enums import ButtonStyle
 from aiogram.exceptions import TelegramAPIError
 from aiogram.filters import Command
 from aiogram.fsm.context import FSMContext
@@ -69,7 +69,7 @@ async def invoice_send(
     status_template = "Рассылка счёта...\nОтправлено: {}"
     status_msg = await msg.answer(status_template.format(0))
 
-    invoice = Invoice(message=data, datetime=datetime.now(UTC))
+    invoice = Invoice(message=data.rich_text, datetime=datetime.now(UTC))
     session.add(invoice)
     await session.flush()
 
@@ -95,7 +95,7 @@ async def invoice_cancel(msg: Message, state: FSMContext) -> None:
 
 @router.message(NewInvoiceStates.message)
 async def invoice_message(msg: Message, bot: Bot, state: FSMContext) -> None:
-    if msg.content_type != ContentType.TEXT or msg.text is None:
+    if msg.text is None:
         await msg.answer(
             "Неверный тип сообщения.\n"
             "Бот поддерживает отправку только текстовых счетов."