Reimage update_subscription callback to work as function

2025-01-17 21:44:54 +03:00
parent 327c37bfc9
commit 466c2dcf31
4 changed files with 120 additions and 74 deletions
@@ -6,20 +6,20 @@ from aiogram import F, Router, types
 from aiogram.fsm.context import FSMContext
 from aiogram.fsm.state import State, StatesGroup
 from aiogram.types import CallbackQuery
-from filters.admin import IsAdminFilter

 from config import TOTAL_GB
 from database import delete_user_data, get_client_id_by_email, get_servers_from_db, restore_trial, update_key_expiry
+from filters.admin import IsAdminFilter
 from handlers.keys.key_utils import (
    delete_key_from_cluster,
    delete_key_from_db,
-    renew_key_in_cluster,
+    renew_key_in_cluster, update_subscription,
 )
 from handlers.utils import sanitize_key_name
 from keyboards.admin.panel_kb import AdminPanelCallback, build_admin_back_kb
 from keyboards.admin.users_kb import build_user_edit_kb, build_key_edit_kb, build_key_delete_kb, \
    build_user_delete_kb, AdminUserEditorCallback, build_editor_kb, build_users_balance_kb, \
-    build_users_balance_change_kb
+    build_users_balance_change_kb, build_user_key_kb
 from logger import logger

 router = Router()
@@ -203,6 +203,7 @@ async def handle_trial_restore(
        session: Any
 ):
    tg_id = callback_data.tg_id
+
    await restore_trial(tg_id, session)
    await callback_query.message.edit_text(
        text="✅ Триал успешно восстановлен!",
@@ -406,9 +407,11 @@ async def handle_change_expiry(
        state: FSMContext
 ):
    email = callback_data.data
+
    await callback_query.message.edit_text(
        text=f"✍️ Введите новое время истечения для ключа <b>{email}</b> в формате <code>YYYY-MM-DD HH:MM:SS</code>:"
    )
+
    await state.update_data(tg_id=callback_data.tg_id, email=email)
    await state.set_state(UserEditorState.waiting_for_expiry_time)

@@ -490,6 +493,29 @@ async def handle_expiry_time_input(
    await state.clear()


+@router.callback_query(
+    AdminUserEditorCallback.filter(F.action == "users_update_key"),
+    IsAdminFilter()
+)
+async def handle_update_key(
+        callback_query: CallbackQuery,
+        callback_data: AdminUserEditorCallback,
+        session: Any
+):
+    tg_id = callback_data.tg_id
+    email = callback_data.data
+
+    try:
+        await update_subscription(tg_id, email, session)
+        await handle_key_edit(callback_query, callback_data, session)
+    except Exception as e:
+        logger.error(f"Ошибка при обновлении ключа {email} администратором: {e}")
+        await callback_query.message.answer(
+            text=f"❗ Произошла ошибка при обновлении ключа: {e}",
+            reply_markup=build_user_key_kb(tg_id, email)
+        )
+
+
@router.callback_query(
    AdminUserEditorCallback.filter(F.action == "users_delete_key"),
    IsAdminFilter()
@@ -1,10 +1,12 @@
 import asyncio
+from typing import Any

 from py3xui import AsyncApi

 from client import add_client, delete_client, extend_client_key
-from config import ADMIN_PASSWORD, ADMIN_USERNAME, LIMIT_IP, SUPERNODE, TOTAL_GB
-from database import get_servers_from_db
+from config import ADMIN_PASSWORD, ADMIN_USERNAME, LIMIT_IP, SUPERNODE, TOTAL_GB, PUBLIC_LINK
+from database import get_servers_from_db, store_key
+from handlers.utils import get_least_loaded_cluster
 from logger import logger


@@ -73,8 +75,6 @@ async def create_key_on_cluster(cluster_id, tg_id, client_id, email, expiry_time
        raise e


-
-
 async def renew_key_in_cluster(cluster_id, email, client_id, new_expiry_time, total_gb):
    try:
        servers = await get_servers_from_db()
@@ -227,3 +227,53 @@ async def update_key_on_cluster(tg_id, client_id, email, expiry_time, cluster_id
            f"Ошибка при обновлении ключа на серверах кластера {cluster_id} для {client_id}: {e}"
        )
        raise e
+
+
+async def update_subscription(tg_id: int, email: str, session: Any) -> None:
+    record = await session.fetchrow(
+        """
+        SELECT k.key, k.expiry_time, k.email, k.server_id, k.client_id
+        FROM keys k
+        WHERE k.tg_id = $1 AND k.email = $2
+        """,
+        tg_id,
+        email,
+    )
+
+    if not record:
+        raise ValueError(f"The key {email} does not exist in database")
+
+    expiry_time = record["expiry_time"]
+    client_id = record["client_id"]
+    public_link = f"{PUBLIC_LINK}{email}/{tg_id}"
+
+    await session.execute(
+        """
+        DELETE FROM keys
+        WHERE tg_id = $1 AND email = $2
+        """,
+        tg_id,
+        email,
+    )
+
+    least_loaded_cluster_id = await get_least_loaded_cluster()
+
+    await asyncio.gather(
+        update_key_on_cluster(
+            tg_id,
+            client_id,
+            email,
+            expiry_time,
+            least_loaded_cluster_id,
+        )
+    )
+
+    await store_key(
+        tg_id,
+        client_id,
+        email,
+        expiry_time,
+        public_link,
+        server_id=least_loaded_cluster_id,
+        session=session,
+    )
@@ -7,6 +7,7 @@ from typing import Any
 from aiogram import F, Router, types
 from aiogram.types import BufferedInputFile, InlineKeyboardButton
 from aiogram.utils.keyboard import InlineKeyboardBuilder
+from handlers.payments.yookassa_pay import process_custom_amount_input

 from bot import bot
 from config import (
@@ -26,7 +27,6 @@ from database import (
    get_balance,
    get_servers_from_db,
    save_temporary_data,
-    store_key,
    update_balance,
    update_key_expiry,
 )
@@ -42,10 +42,9 @@ from handlers.keys.key_utils import (
    delete_key_from_cluster,
    delete_key_from_db,
    renew_key_in_cluster,
-    update_key_on_cluster,
+    update_subscription,
 )
 from handlers.payments.robokassa_pay import handle_custom_amount_input
-from handlers.payments.yookassa_pay import process_custom_amount_input
 from handlers.texts import (
    DISCOUNTS,
    KEY_NOT_FOUND_MSG,
@@ -53,7 +52,7 @@ from handlers.texts import (
    SUCCESS_RENEWAL_MSG,
    key_message,
 )
-from handlers.utils import get_least_loaded_cluster, handle_error
+from handlers.utils import handle_error
 from logger import logger

 locale.setlocale(locale.LC_TIME, "ru_RU.UTF-8")
@@ -64,7 +63,7 @@ router = Router()
@router.callback_query(F.data == "view_keys")
@router.message(F.text == "/subs")
 async def process_callback_or_message_view_keys(
-    callback_query_or_message: types.Message | types.CallbackQuery, session: Any
+        callback_query_or_message: types.Message | types.CallbackQuery, session: Any
 ):
    if isinstance(callback_query_or_message, types.CallbackQuery):
        chat_id = callback_query_or_message.message.chat.id
@@ -124,7 +123,7 @@ def build_keys_response(records):


 async def send_with_optional_image(
-    send_message, send_photo, image_path, text, keyboard
+        send_message, send_photo, image_path, text, keyboard
 ):
    """
    Отправляет сообщение с изображением, если файл существует. В противном случае отправляет только текст.
@@ -268,67 +267,16 @@ async def process_callback_view_key(callback_query: types.CallbackQuery, session

@router.callback_query(F.data.startswith("update_subscription|"))
 async def process_callback_update_subscription(
-    callback_query: types.CallbackQuery, session: Any
+        callback_query: types.CallbackQuery, session: Any
 ):
    tg_id = callback_query.message.chat.id
    email = callback_query.data.split("|")[1]
+
    try:
-        record = await session.fetchrow(
-            """
-            SELECT k.key, k.expiry_time, k.email, k.server_id, k.client_id
-            FROM keys k
-            WHERE k.tg_id = $1 AND k.email = $2
-            """,
-            tg_id,
-            email,
-        )
-
-        if record:
-            expiry_time = record["expiry_time"]
-            client_id = record["client_id"]
-            public_link = f"{PUBLIC_LINK}{email}/{tg_id}"
-
-            try:
-                await session.execute(
-                    """
-                    DELETE FROM keys
-                    WHERE tg_id = $1 AND email = $2
-                    """,
-                    tg_id,
-                    email,
-                )
-            except Exception as delete_error:
-                await callback_query.message.answer(
-                    f"Ошибка при удалении старой подписки: {delete_error}",
-                )
-                return
-
-            least_loaded_cluster_id = await get_least_loaded_cluster()
-
-            await asyncio.gather(
-                update_key_on_cluster(
-                    tg_id,
-                    client_id,
-                    email,
-                    expiry_time,
-                    least_loaded_cluster_id,
-                )
-            )
-
-            await store_key(
-                tg_id,
-                client_id,
-                email,
-                expiry_time,
-                public_link,
-                server_id=least_loaded_cluster_id,
-                session=session,
-            )
-
-            await process_callback_view_key(callback_query, session)
-        else:
-            await callback_query.message.answer("<b>Ключ не найден в базе данных.</b>")
+        await update_subscription(tg_id, email, session)
+        await process_callback_view_key(callback_query, session)
    except Exception as e:
+        logger.error(f"Ошибка при обновлении ключа {email} пользователем: {e}")
        await handle_error(
            tg_id, callback_query, f"Ошибка при обновлении подписки: {e}"
        )
@@ -436,7 +384,7 @@ async def process_callback_renew_key(callback_query: types.CallbackQuery, sessio

@router.callback_query(F.data.startswith("confirm_delete|"))
 async def process_callback_confirm_delete(
-    callback_query: types.CallbackQuery, session: Any
+        callback_query: types.CallbackQuery, session: Any
 ):
    email = callback_query.data.split("|")[1]
    try:
@@ -523,7 +471,8 @@ async def process_callback_renew_plan(callback_query: types.CallbackQuery, sessi
            if balance < cost:
                required_amount = cost - balance

-                logger.info(f"[RENEW] Пользователю {tg_id} не хватает {required_amount}₽. Запуск доплаты через {USE_NEW_PAYMENT_FLOW}")
+                logger.info(
+                    f"[RENEW] Пользователю {tg_id} не хватает {required_amount}₽. Запуск доплаты через {USE_NEW_PAYMENT_FLOW}")

                await save_temporary_data(
                    session,
@@ -604,4 +553,3 @@ async def complete_key_renewal(tg_id, client_id, email, new_expiry_time, total_g
        logger.info(f"[RENEW] Ключ {client_id} успешно продлён на {plan} мес. для пользователя {tg_id}.")

    await renew_key_on_servers()
-
@@ -141,10 +141,24 @@ def build_user_delete_kb(tg_id: int):
    return builder.as_markup()


+def build_user_key_kb(tg_id: int, email: str) -> InlineKeyboardMarkup:
+    builder = InlineKeyboardBuilder()
+    builder.button(
+        text=f"🔙 Назад",
+        callback_data=AdminUserEditorCallback(
+            action="users_key_edit",
+            tg_id=tg_id,
+            data=email
+        ).pack()
+    )
+    builder.adjust(1)
+    return builder.as_markup()
+
+
 def build_key_edit_kb(key_details: dict, email: str) -> InlineKeyboardMarkup:
    builder = InlineKeyboardBuilder()
    builder.button(
-        text="⏳ Изменить время истечения",
+        text="⏳ Время истечения",
        callback_data=AdminUserEditorCallback(
            action="users_change_expiry",
            data=email,
@@ -152,7 +166,15 @@ def build_key_edit_kb(key_details: dict, email: str) -> InlineKeyboardMarkup:
        ).pack()
    )
    builder.button(
-        text="❌ Удалить ключ",
+        text="🔄 Перевыпустить",
+        callback_data=AdminUserEditorCallback(
+            action="users_update_key",
+            data=email,
+            tg_id=key_details["tg_id"]
+        ).pack()
+    )
+    builder.button(
+        text="❌ Удалить",
        callback_data=AdminUserEditorCallback(
            action="users_delete_key",
            data=email,