add trial limits for create clusters_mode

2025-05-25 23:54:43 +03:00
parent d79deb307e
commit d55fbf0fed
5 changed files with 43 additions and 29 deletions
@@ -13,7 +13,7 @@ from aiogram.types import (
 from aiogram.utils.keyboard import InlineKeyboardBuilder

 from bot import bot
-from config import CONNECT_PHONE_BUTTON, SUPPORT_CHAT_URL
+from config import CONNECT_PHONE_BUTTON, SUPPORT_CHAT_URL, TRIAL_CONFIG
 from database import (
    get_key_details,
    get_tariff_by_id,
@@ -73,8 +73,16 @@ async def key_cluster_mode(
    expiry_timestamp = int(expiry_time.timestamp() * 1000)

    try:
+        data = await state.get_data() if state else {}
+        is_trial = data.get("is_trial", False)
+
        device_limit = 0
-        if plan:
+        traffic_limit_bytes = None
+
+        if is_trial:
+            device_limit = TRIAL_CONFIG.get("hwid_limit", 1)
+            traffic_limit_bytes = int(TRIAL_CONFIG.get("traffic_limit_gb", 100) * 1024**3)
+        elif plan:
            tariff = await get_tariff_by_id(session, plan)
            if tariff and tariff.get("device_limit") is not None:
                device_limit = int(tariff["device_limit"])
@@ -89,6 +97,7 @@ async def key_cluster_mode(
            plan=plan,
            session=session,
            hwid_limit=device_limit,
+            traffic_limit_bytes=traffic_limit_bytes,
        )

        logger.info(
@@ -103,9 +112,7 @@ async def key_cluster_mode(
        remnawave_link = key_record.get("remnawave_link")
        final_link = public_link or remnawave_link or ""

-        data = await state.get_data() if state else {}
-
-        if data.get("is_trial"):
+        if is_trial:
            trial_status = await get_trial(session, tg_id)
            if trial_status in [0, -1]:
                await update_trial(session, tg_id, 1)
@@ -9,7 +9,7 @@ from aiogram.utils.keyboard import InlineKeyboardBuilder

 from config import (
    NOTIFY_EXTRA_DAYS,
-    TRIAL_TIME,
+    TRIAL_CONFIG,
    TRIAL_TIME_DISABLE,
    USE_COUNTRY_SELECTION,
    USE_NEW_PAYMENT_FLOW,
@@ -66,20 +66,23 @@ async def handle_key_creation(
    if not TRIAL_TIME_DISABLE:
        trial_status = await get_trial(session, tg_id)
        if trial_status in [0, -1]:
+            base_days = TRIAL_CONFIG["duration_days"]
            extra_days = NOTIFY_EXTRA_DAYS if trial_status == -1 else 0
-            expiry_time = current_time + timedelta(days=TRIAL_TIME + extra_days)
-            logger.info(
-                f"Доступен {TRIAL_TIME + extra_days}-дневный пробный период пользователю {tg_id}."
-            )
+            total_days = base_days + extra_days
+            expiry_time = current_time + timedelta(days=total_days)
+
+            logger.info(f"[Trial] Доступен {total_days}-дневный триал для пользователя {tg_id}")
+
            await edit_or_send_message(
                target_message=(
-                    message_or_query
-                    if isinstance(message_or_query, Message)
-                    else message_or_query.message
+                    message_or_query.message
+                    if isinstance(message_or_query, CallbackQuery)
+                    else message_or_query
                ),
                text=CREATING_CONNECTION_MSG,
                reply_markup=None,
            )
+
            await state.update_data(is_trial=True)
            await create_key(tg_id, expiry_time, state, session, message_or_query)
            return
@@ -90,9 +93,9 @@ async def handle_key_creation(
    if not tariffs:
        await edit_or_send_message(
            target_message=(
-                message_or_query
-                if isinstance(message_or_query, Message)
-                else message_or_query.message
+                message_or_query.message
+                if isinstance(message_or_query, CallbackQuery)
+                else message_or_query
            ),
            text="❌ Нет доступных тарифов для выбранного кластера.",
            reply_markup=None,
@@ -114,6 +117,7 @@ async def handle_key_creation(
        if isinstance(message_or_query, CallbackQuery)
        else message_or_query
    )
+
    await edit_or_send_message(
        target_message=target_message,
        text=SELECT_TARIFF_PLAN_MSG,
@@ -32,6 +32,7 @@ async def create_key_on_cluster(
    session: AsyncSession = None,
    remnawave_link: str = None,
    hwid_limit: int = None,
+    traffic_limit_bytes: int = None, 
 ):
    try:
        servers = await get_servers(session, include_enabled=True)
@@ -59,19 +60,15 @@ async def create_key_on_cluster(
            )
            return

-        traffic_limit_bytes = None
-        if plan is not None:
+        if plan is not None and traffic_limit_bytes is None:
            tariff = await get_tariff_by_id(session, plan)
            if not tariff:
                raise ValueError(f"Тариф с id={plan} не найден.")
            traffic_limit_bytes = (
                int(tariff["traffic_limit"]) if tariff["traffic_limit"] else None
            )
-            hwid_limit = (
-                int(tariff["device_limit"])
-                if tariff["device_limit"] is not None
-                else None
-            )
+            if hwid_limit is None and tariff.get("device_limit") is not None:
+                hwid_limit = int(tariff["device_limit"])

        remnawave_servers = [
            s
@@ -10,7 +10,7 @@ from config import (
    NOTIFY_INACTIVE,
    NOTIFY_INACTIVE_TRAFFIC,
    SUPPORT_CHAT_URL,
-    TRIAL_TIME,
+    TRIAL_CONFIG,
 )
 from database import (
    add_notification,
@@ -40,6 +40,8 @@ async def notify_inactive_trial_users(bot: Bot, session: AsyncSession):
    logger.info(f"Найдено {len(users)} неактивных пользователей для уведомления.")
    messages = []

+    trial_days = TRIAL_CONFIG["duration_days"]
+
    for user in users:
        tg_id = user["tg_id"]
        username = user["username"]
@@ -59,7 +61,7 @@ async def notify_inactive_trial_users(bot: Bot, session: AsyncSession):
        trial_extended = user["last_notification_time"] is not None

        if trial_extended:
-            total_days = NOTIFY_EXTRA_DAYS + TRIAL_TIME
+            total_days = NOTIFY_EXTRA_DAYS + trial_days
            message = TRIAL_INACTIVE_BONUS_MSG.format(
                display_name=display_name,
                extra_days_formatted=format_days(NOTIFY_EXTRA_DAYS),
@@ -68,7 +70,8 @@ async def notify_inactive_trial_users(bot: Bot, session: AsyncSession):
            await mark_trial_extended(tg_id, session)
        else:
            message = TRIAL_INACTIVE_FIRST_MSG.format(
-                display_name=display_name, trial_time_formatted=format_days(TRIAL_TIME)
+                display_name=display_name,
+                trial_time_formatted=format_days(trial_days),
            )

        messages.append(
@@ -18,7 +18,7 @@ from sqlalchemy import desc, func, select
 from sqlalchemy.ext.asyncio import AsyncSession

 from bot import bot
-from config import ADMIN_ID, INLINE_MODE, TOP_REFERRAL_BUTTON, TRIAL_TIME, USERNAME_BOT
+from config import ADMIN_ID, INLINE_MODE, TOP_REFERRAL_BUTTON, TRIAL_CONFIG, USERNAME_BOT
 from database import (
    add_referral,
    add_user,
@@ -86,13 +86,16 @@ async def inline_referral_handler(inline_query: InlineQuery):
    referral_link = (
        f"https://t.me/{USERNAME_BOT}?start=referral_{inline_query.from_user.id}"
    )
-    trial_time_formatted = format_days(TRIAL_TIME)
+    trial_days = TRIAL_CONFIG["duration_days"]
+    trial_time_formatted = format_days(trial_days)
+
    results: list[InlineQueryResultArticle] = []

    for index, offer in enumerate(REFERRAL_OFFERS):
        description = offer["description"][:64]
        message_text = offer["message"].format(
-            trial_time=TRIAL_TIME, trial_time_formatted=trial_time_formatted
+            trial_time=trial_days,
+            trial_time_formatted=trial_time_formatted
        )[:4096]

        builder = InlineKeyboardBuilder()