Add new_announcement command

11 files changed, 296 insertions, 5 deletions

diff --git a/handlers/admin/__init__.py b/handlers/admin/__init__.py
index d0a5587..2f7c74f 100644
--- a/handlers/admin/__init__.py
+++ b/handlers/admin/__init__.py
@@ -1,6 +1,14 @@
 from aiogram import F, Router
 from aiogram.filters import MagicData
 
+# isort: off
+from . import new_announcement
+# isort: on
+
 router = Router(name="admin")
 router.message.filter(MagicData(F.user.is_admin()))
 router.callback_query.filter(MagicData(F.user.is_admin()))
+
+router.include_routers(
+    new_announcement.router,
+)
diff --git a/handlers/admin/new_announcement.py b/handlers/admin/new_announcement.py
new file mode 100644
index 0000000..79cf8d4
--- /dev/null
+++ b/handlers/admin/new_announcement.py
@@ -0,0 +1,108 @@
+from datetime import UTC, datetime
+
+from aiogram import Bot, F, Router
+from aiogram.exceptions import TelegramAPIError
+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,
+    ReplyKeyboardRemove,
+)
+from pydantic.main import BaseModel
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from libs.fsm import get_data, set_data
+from libs.msg import publish
+from models import Announcement, RichText, User
+
+router = Router(name="new_announcement")
+
+
+class NewAnnouncementStates(StatesGroup):
+    message = State()
+
+
+class NewAnnouncementData(BaseModel):
+    rich_text: RichText | None = None
+
+
+SEND_BUTTON = "Опубликовать"
+CANCEL_BUTTON = "Отменить создание"
+
+
+@router.message(Command("new_announcement"))
+async def new_announcement_command(msg: Message, state: FSMContext) -> None:
+    await msg.answer(
+        "Укажите сообщение для анонса.",
+        reply_markup=ReplyKeyboardMarkup(
+            keyboard=[
+                [
+                    KeyboardButton(text=SEND_BUTTON),
+                    KeyboardButton(text=CANCEL_BUTTON),
+                ]
+            ],
+            resize_keyboard=True,
+        ),
+    )
+    await state.set_state(NewAnnouncementStates.message)
+
+
+@router.message(NewAnnouncementStates.message, F.text == SEND_BUTTON)
+async def announcement_send(
+    msg: Message,
+    bot: Bot,
+    state: FSMContext,
+    session: AsyncSession,
+) -> None:
+    users = await session.scalars(select(User.id).where(User.id != msg.chat.id))
+    data = await get_data(state, NewAnnouncementData)
+
+    if data.rich_text is None:
+        await msg.answer("Для публикации анонса укажите текст сообщения.")
+        return
+
+    status_template = "Публикация анонса...\nОпубликовано: {}"
+    status_msg = await msg.answer(status_template.format(0))
+
+    async for n in publish(bot, users, data.rich_text):
+        try:
+            await status_msg.edit_text(status_template.format(n))
+        except TelegramAPIError:
+            pass
+
+    announcement = Announcement(message=data, datetime=datetime.now(UTC))
+    session.add(announcement)
+
+    await status_msg.delete()
+    await msg.answer(
+        "Анонс отправлен всем пользователям",
+        reply_markup=ReplyKeyboardRemove(),
+    )
+    await state.clear()
+
+
+@router.message(NewAnnouncementStates.message, F.text == CANCEL_BUTTON)
+async def announcement_cancel(msg: Message, state: FSMContext) -> None:
+    await msg.answer("Создание анонса отменено", reply_markup=ReplyKeyboardRemove())
+    await state.clear()
+
+
+@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:
+        await msg.answer(
+            "Неверный тип сообщения.\n"
+            "Бот поддерживает отправку только текстовых анонсов."
+        )
+        return
+
+    rich_text = RichText.from_message(msg)
+    await set_data(state, NewAnnouncementData(rich_text=rich_text))
+
+    msg_rich_text = RichText.from_text("Сообщение вашего анонса:\n", rich_text)
+    await msg_rich_text.send(bot, msg.chat.id)
diff --git a/libs/__init__.py b/libs/__init__.py
new file mode 100644
index 0000000..d8ed122
--- /dev/null
+++ b/libs/__init__.py
@@ -0,0 +1,7 @@
+from . import fsm, msg, storage
+
+__all__ = [
+    "storage",
+    "fsm",
+    "msg",
+]
diff --git a/libs/fsm.py b/libs/fsm.py
new file mode 100644
index 0000000..00ccacb
--- /dev/null
+++ b/libs/fsm.py
@@ -0,0 +1,23 @@
+from contextlib import asynccontextmanager
+
+from aiogram.fsm.context import FSMContext
+from pydantic.main import BaseModel
+from typing_extensions import AsyncGenerator
+
+
+async def set_data(state: FSMContext, model: BaseModel) -> None:
+    await state.set_data(model.model_dump())
+
+
+async def get_data[T: BaseModel](state: FSMContext, model_type: type[T]) -> T:
+    return model_type.model_validate(await state.get_data())
+
+
+@asynccontextmanager
+async def edit_data[T: BaseModel](
+    state: FSMContext,
+    model_type: type[T],
+) -> AsyncGenerator[T]:
+    model = await get_data(state, model_type)
+    yield model
+    await set_data(state, model)
diff --git a/libs/msg.py b/libs/msg.py
new file mode 100644
index 0000000..9bcc52a
--- /dev/null
+++ b/libs/msg.py
@@ -0,0 +1,25 @@
+import asyncio
+from typing import AsyncGenerator, Iterable
+
+from aiogram import Bot
+from aiogram.exceptions import TelegramAPIError, TelegramRetryAfter
+
+from models import RichText
+
+
+async def publish(
+    bot: Bot,
+    users: Iterable[int],
+    rich_text: RichText,
+) -> AsyncGenerator[int]:
+    for n, user_id in enumerate(users, start=1):
+        for _ in range(5):
+            try:
+                await rich_text.send(bot, user_id)
+                break
+            except TelegramRetryAfter as e:
+                await asyncio.sleep(e.retry_after + 1)
+            except TelegramAPIError:
+                await asyncio.sleep(5)
+
+        yield n
diff --git a/libs/storage.py b/libs/storage.py
new file mode 100644
index 0000000..6220cfc
--- /dev/null
+++ b/libs/storage.py
@@ -0,0 +1,75 @@
+from pathlib import Path
+from typing import Any, Mapping
+
+from aiofiles import open as open
+from aiogram.fsm.state import State
+from aiogram.fsm.storage.base import (
+    BaseStorage,
+    DefaultKeyBuilder,
+    KeyBuilder,
+    StateType,
+    StorageKey,
+)
+from pydantic import TypeAdapter
+from pydantic.main import BaseModel
+
+
+class Record(BaseModel):
+    data: dict[str, Any] = {}
+    state: str | None = None
+
+
+class JsonStorage(BaseStorage):
+    file_path: Path
+    records: dict[str, Record]
+    records_adapter: TypeAdapter
+    key_builder: KeyBuilder
+
+    def __init__(self, file_path: Path, key_builder: KeyBuilder | None = None) -> None:
+        self.file_path = file_path
+        self.records = {}
+        self.records_adapter = TypeAdapter(dict[str, Record])
+        self.key_builder = DefaultKeyBuilder() if key_builder is None else key_builder
+
+    async def read(self) -> None:
+        async with open(self.file_path, "rb") as file:
+            json = await file.read()
+            self.records = self.records_adapter.validate_json(json)
+
+    async def flush(self) -> None:
+        async with open(self.file_path, "wb") as file:
+            json = self.records_adapter.dump_json(self.records)
+            await file.write(json)
+
+    async def get_record(self, key: StorageKey) -> Record:
+        await self.read()
+        record_key = self.key_builder.build(key)
+        if record_key not in self.records:
+            self.records[record_key] = Record()
+        return self.records[record_key]
+
+    async def set_state(self, key: StorageKey, state: StateType = None) -> None:
+        record = await self.get_record(key)
+        record.state = state.state if isinstance(state, State) else state
+        await self.flush()
+
+    async def get_state(self, key: StorageKey) -> str | None:
+        record = await self.get_record(key)
+        return record.state
+
+    async def set_data(self, key: StorageKey, data: Mapping[str, Any]) -> None:
+        if not isinstance(data, dict):
+            raise TypeError(
+                f"Data must be a dict or dict-like object, got {type(data).__name__}",
+                data,
+            )
+        record = await self.get_record(key)
+        record.data = data.copy()
+        await self.flush()
+
+    async def get_data(self, key: StorageKey) -> dict[str, Any]:
+        record = await self.get_record(key)
+        return record.data
+
+    async def close(self) -> None:
+        await self.flush()
diff --git a/models/rich_text.py b/models/rich_text.py
index 2b433ec..e5441b6 100644
--- a/models/rich_text.py
+++ b/models/rich_text.py
@@ -1,7 +1,43 @@
-from aiogram.types import MessageEntity
+from typing import Self
+
+from aiogram import Bot
+from aiogram.types import Message, MessageEntity
 from pydantic import BaseModel
 
 
 class RichText(BaseModel):
     text: str
     entities: list[MessageEntity] = []
+
+    @classmethod
+    def from_message(cls, msg: Message) -> Self:
+        assert msg.text is not None
+        return cls(
+            text=msg.text,
+            entities=[] if msg.entities is None else msg.entities,
+        )
+
+    @classmethod
+    def from_text(cls, *text: RichText | str) -> Self:
+        result = cls(text="", entities=[])
+
+        for t in text:
+            if isinstance(t, RichText):
+                entities = t.entities.copy()
+                for e in entities:
+                    e.offset += len(result.text)
+                result.entities += entities
+
+                result.text += t.text
+            else:
+                result.text += t
+
+        return result
+
+    async def send(self, bot: Bot, chat_id: int) -> Message:
+        return await bot.send_message(
+            chat_id=chat_id,
+            text=self.text,
+            entities=self.entities,
+            parse_mode=None,
+        )
diff --git a/pyproject.toml b/pyproject.toml
index 71a565e..29f1c60 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,6 +4,7 @@ version = "0.0.1"
 description = "Simple bot for vpn management"
 requires-python = ">=3.14"
 dependencies = [
+  "aiofiles>=25.1.0",
   "aiogram>=3.26.0",
   "aiohttp-socks>=0.11.0",
   "aiosqlite>=0.22.1",
diff --git a/settings.py b/settings.py
index 2daff90..04ba431 100644
--- a/settings.py
+++ b/settings.py
@@ -8,12 +8,17 @@ class Env(BaseSettings):
     proxy: str | None = None
 
 
-database_path = Path("./storage/database.db").absolute()
-database_path.parent.mkdir(parents=True, exist_ok=True)
+storage_path = Path("./storage/").absolute()
+storage_path.mkdir(parents=True, exist_ok=True)
+
+database_path = storage_path / "database.db"
 database_url = f"sqlite+aiosqlite:///{database_path}"
 
+json_storage_path = storage_path / "storage.json"
+
 __all__ = [
     "database_path",
     "database_url",
+    "json_storage_path",
     "Env",
 ]
diff --git a/shared.py b/shared.py
index 5c7528c..63d00c3 100644
--- a/shared.py
+++ b/shared.py
@@ -3,7 +3,8 @@ from aiogram.client.default import DefaultBotProperties
 from aiogram.client.session.aiohttp import AiohttpSession
 from aiogram.enums import ParseMode
 
-from settings import Env
+from libs.storage import JsonStorage
+from settings import Env, json_storage_path
 
 env = Env()  # ty:ignore[missing-argument] # pyright: ignore[reportCallIssue]
 bot = Bot(
@@ -11,4 +12,4 @@ bot = Bot(
     session=AiohttpSession(proxy=env.proxy),
     default=DefaultBotProperties(parse_mode=ParseMode.HTML),
 )
-dp = Dispatcher()
+dp = Dispatcher(storage=JsonStorage(json_storage_path))
diff --git a/uv.lock b/uv.lock
index a9b8849..3d84e96 100644
--- a/uv.lock
+++ b/uv.lock
@@ -516,6 +516,7 @@ name = "vpn-manager-bot"
 version = "0.0.1"
 source = { virtual = "." }
 dependencies = [
+    { name = "aiofiles" },
     { name = "aiogram" },
     { name = "aiohttp-socks" },
     { name = "aiosqlite" },
@@ -528,6 +529,7 @@ dependencies = [
 
 [package.metadata]
 requires-dist = [
+    { name = "aiofiles", specifier = ">=25.1.0" },
     { name = "aiogram", specifier = ">=3.26.0" },
     { name = "aiohttp-socks", specifier = ">=0.11.0" },
     { name = "aiosqlite", specifier = ">=0.22.1" },