From f7b7e87cffc9dcb2817b070d7a003ac234c96ec3 Mon Sep 17 00:00:00 2001 From: Tolmachev Igor Date: Mon, 23 Mar 2026 18:07:30 +0300 Subject: Add new_announcement command --- handlers/admin/__init__.py | 8 +++ handlers/admin/new_announcement.py | 108 +++++++++++++++++++++++++++++++++++++ libs/__init__.py | 7 +++ libs/fsm.py | 23 ++++++++ libs/msg.py | 25 +++++++++ libs/storage.py | 75 ++++++++++++++++++++++++++ models/rich_text.py | 38 ++++++++++++- pyproject.toml | 1 + settings.py | 9 +++- shared.py | 5 +- uv.lock | 2 + 11 files changed, 296 insertions(+), 5 deletions(-) create mode 100644 handlers/admin/new_announcement.py create mode 100644 libs/__init__.py create mode 100644 libs/fsm.py create mode 100644 libs/msg.py create mode 100644 libs/storage.py 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" }, -- cgit v1.3