diff --git a/main.py b/main.py index d4535e6..a566da1 100644 --- a/main.py +++ b/main.py @@ -12,7 +12,7 @@ from jobs import stop_scheduler, start_scheduler from middlewares.auth import CheckAdminAccess from routers import setup_routers -from utils import EnvSettings, logger, storage +from utils import EnvSettings, logger, Storage async def on_startup() -> None: @@ -52,7 +52,7 @@ async def main() -> None: ) # Set up Dispatcher - dp = Dispatcher(storage=storage) + dp = Dispatcher(storage=Storage) # Set routers to dp dp.include_router(setup_routers()) diff --git a/middlewares/auth.py b/middlewares/auth.py index cdfde15..8c79f4d 100644 --- a/middlewares/auth.py +++ b/middlewares/auth.py @@ -6,7 +6,7 @@ from aiogram import BaseMiddleware from aiogram.types import Update -from utils import logger, storage, EnvSettings +from utils import logger, Storage, EnvSettings # pylint: disable=too-few-public-methods @@ -26,7 +26,7 @@ async def __call__( user = None if event.message: user = event.message.from_user - await storage.add_log_message(user.id, event.message.message_id) + await Storage.add_log_message(user.id, event.message.message_id) elif event.callback_query: user = event.callback_query.from_user elif event.inline_query: diff --git a/routers/base.py b/routers/base.py index d085d0e..c78b81b 100644 --- a/routers/base.py +++ b/routers/base.py @@ -8,7 +8,7 @@ from aiogram.filters.command import CommandStart, Command from aiogram.fsm.context import FSMContext -from utils import MessageTexts, storage, BotKeyboards +from utils import MessageTexts, Storage, BotKeyboards from models import PagesCallbacks, PagesActions router = Router() @@ -24,7 +24,7 @@ async def start(message: Message, state: FSMContext): new_message = await message.answer( text=MessageTexts.START, reply_markup=BotKeyboards.home() ) - return await storage.clear_and_add_message(new_message) + return await Storage.clear_and_add_message(new_message) @router.callback_query(PagesCallbacks.filter(F.page.is_(PagesActions.HOME))) @@ -37,7 +37,7 @@ async def home(callback: CallbackQuery, state: FSMContext): new_message = await callback.message.answer( text=MessageTexts.START, reply_markup=BotKeyboards.home() ) - return await storage.clear_and_add_message(new_message) + return await Storage.clear_and_add_message(new_message) @router.message(Command(commands=["version", "v"])) @@ -46,4 +46,4 @@ async def version(message: Message): Handler for the '/version' and '/v' commands. Sends the version information of the bot. """ new_message = await message.answer(text=MessageTexts.VERSION) - return await storage.clear_and_add_message(new_message) + return await Storage.clear_and_add_message(new_message) diff --git a/routers/user.py b/routers/user.py index 6aaf497..7cee878 100644 --- a/routers/user.py +++ b/routers/user.py @@ -15,7 +15,7 @@ text_info, helpers, MessageTexts, - storage, + Storage, EnvSettings, BotKeyboards, ) @@ -50,50 +50,55 @@ async def user_create_base_username(message: Message, state: FSMContext): Handles the input for the base username in the user creation process. """ await state.update_data(base_username=message.text) - await state.set_state(UserCreateForm.start_number) + await state.set_state(UserCreateForm.how_much) new_message = await message.answer( - text=MessageTexts.ASK_CREATE_USER_START_NUMBER, + text=MessageTexts.ASK_CREATE_USER_HOW_MUCH, reply_markup=BotKeyboards.cancel(), ) - return await storage.clear_and_add_message(new_message) + return await Storage.clear_and_add_message(new_message) -@router.message(StateFilter(UserCreateForm.start_number)) -async def user_create_start_number(message: Message, state: FSMContext): +@router.message(StateFilter(UserCreateForm.how_much)) +async def user_create_how_much(message: Message, state: FSMContext): """ - Handles the input for the starting number in the user creation process. + Handles the input for the 'how much' field in the user creation process. """ if not message.text.isdigit(): new_message = await message.answer(text=MessageTexts.JUST_NUMBER) - return await storage.add_log_message( + return await Storage.add_log_message( message.from_user.id, new_message.message_id ) - await state.update_data(start_number=int(message.text)) - await state.set_state(UserCreateForm.how_much) - new_message = await message.answer( - text=MessageTexts.ASK_CREATE_USER_HOW_MUCH, reply_markup=BotKeyboards.cancel() - ) - return await storage.clear_and_add_message(new_message) + await state.update_data(how_much=int(message.text)) + if int(message.text) == 1: + await state.set_state(UserCreateForm.data_limit) + text = MessageTexts.ASK_CREATE_USER_DATA_LIMIT + else: + await state.set_state(UserCreateForm.start_number) + text = MessageTexts.ASK_CREATE_USER_START_NUMBER -@router.message(StateFilter(UserCreateForm.how_much)) -async def user_create_how_much(message: Message, state: FSMContext): + new_message = await message.answer(text=text, reply_markup=BotKeyboards.cancel()) + return await Storage.clear_and_add_message(new_message) + + +@router.message(StateFilter(UserCreateForm.start_number)) +async def user_create_start_number(message: Message, state: FSMContext): """ - Handles the input for the 'how much' field in the user creation process. + Handles the input for the starting number in the user creation process. """ if not message.text.isdigit(): new_message = await message.answer(text=MessageTexts.JUST_NUMBER) - return await storage.add_log_message( + return await Storage.add_log_message( message.from_user.id, new_message.message_id ) - await state.update_data(how_much=int(message.text)) - await state.set_state(UserCreateForm.data_limit) + await state.update_data(start_number=int(message.text)) + await state.set_state(UserCreateForm.date_limit) new_message = await message.answer( text=MessageTexts.ASK_CREATE_USER_DATA_LIMIT, reply_markup=BotKeyboards.cancel() ) - return await storage.clear_and_add_message(new_message) + return await Storage.clear_and_add_message(new_message) @router.message(StateFilter(UserCreateForm.data_limit)) @@ -103,7 +108,7 @@ async def user_create_data_limit(message: Message, state: FSMContext): """ if not message.text.isdigit(): new_message = await message.answer(text=MessageTexts.JUST_NUMBER) - return await storage.add_log_message( + return await Storage.add_log_message( message.from_user.id, new_message.message_id ) @@ -112,7 +117,7 @@ async def user_create_data_limit(message: Message, state: FSMContext): new_message = await message.answer( text=MessageTexts.ASK_CREATE_USER_DATE_LIMIT, reply_markup=BotKeyboards.cancel() ) - return await storage.clear_and_add_message(new_message) + return await Storage.clear_and_add_message(new_message) @router.message(StateFilter(UserCreateForm.date_limit)) @@ -122,7 +127,7 @@ async def user_create_date_limit(message: Message, state: FSMContext): """ if not message.text.isdigit(): new_message = await message.answer(text=MessageTexts.JUST_NUMBER) - return await storage.add_log_message( + return await Storage.add_log_message( message.from_user.id, new_message.message_id ) @@ -131,7 +136,7 @@ async def user_create_date_limit(message: Message, state: FSMContext): text=MessageTexts.ASK_CREATE_USER_STATUS, reply_markup=BotKeyboards.user_status(AdminActions.ADD), ) - return await storage.clear_and_add_message(new_message) + return await Storage.clear_and_add_message(new_message) @router.callback_query(UserStatusCallbacks.filter(F.action.is_(AdminActions.ADD))) @@ -239,7 +244,10 @@ async def user_create_inbounds_save(callback: CallbackQuery, state: FSMContext): inbounds_dict = {k: v for k, v in inbounds_dict.items() if v} for i in range(int(data["how_much"])): - username = f"{data['base_username']}{int(data['start_number']) + i}" + if int(data["how_much"]) == 1: + username = data["base_username"] + else: + username = f"{data['base_username']}{int(data['start_number']) + i}" new_user = await panel.create_user( username=username, status=data["status"], @@ -268,4 +276,4 @@ async def user_create_inbounds_save(callback: CallbackQuery, state: FSMContext): new_message = await callback.message.answer( text=MessageTexts.START, reply_markup=BotKeyboards.home() ) - return await storage.clear_and_add_message(new_message) + return await Storage.clear_and_add_message(new_message) diff --git a/utils/__init__.py b/utils/__init__.py index 78ec732..690ec8d 100644 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -12,7 +12,7 @@ EnvSettings = EnvFile() MessageTexts = MessageTextsFile() KeyboardTexts = KeyboardTextsFile() -storage = SQLAlchemyStorage() +Storage = SQLAlchemyStorage() -__all__ = ["storage", "logger", "KeyboardTexts", "MessageTexts", "BotKeyboards"] +__all__ = ["Storage", "logger", "KeyboardTexts", "MessageTexts", "BotKeyboards"]