From aaea5b4dbefdc9f58f3dc7ea10f4d7790207edd5 Mon Sep 17 00:00:00 2001
From: Capybara-z <stas89091000@yandex.ru>
Date: Sat, 14 Jun 2025 03:45:47 +0300
Subject: [PATCH] fix: traffic reset / coupon menu display

---
 handlers/admin/clusters/clusters_handler.py |  1 +
 handlers/admin/users/users_handler.py       |  3 ++-
 handlers/coupons.py                         |  3 ++-
 handlers/keys/key_freeze.py                 |  1 +
 handlers/keys/key_utils.py                  |  4 +++-
 handlers/start.py                           | 15 ++++++++++++++-
 6 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/handlers/admin/clusters/clusters_handler.py b/handlers/admin/clusters/clusters_handler.py
index bc7af8a1..ab64cd8f 100644
--- a/handlers/admin/clusters/clusters_handler.py
+++ b/handlers/admin/clusters/clusters_handler.py
@@ -828,6 +828,7 @@ async def handle_days_input(message: Message, state: FSMContext, session: AsyncS
                 new_expiry_time=new_expiry,
                 total_gb=total_gb,
                 session=session,
+                reset_traffic=False,
             )
             await update_key_expiry(session, key.client_id, new_expiry)
 
diff --git a/handlers/admin/users/users_handler.py b/handlers/admin/users/users_handler.py
index 5f1ed39d..b525ebd3 100644
--- a/handlers/admin/users/users_handler.py
+++ b/handlers/admin/users/users_handler.py
@@ -1203,7 +1203,8 @@ async def change_expiry_time(
         new_expiry_time=expiry_time,
         total_gb=traffic_limit,
         session=session,
-        hwid_device_limit=device_limit
+        hwid_device_limit=device_limit,
+        reset_traffic=False,
     )
     
     await update_key_expiry(session, client_id, expiry_time)
diff --git a/handlers/coupons.py b/handlers/coupons.py
index b072d97d..325d47ce 100644
--- a/handlers/coupons.py
+++ b/handlers/coupons.py
@@ -243,8 +243,9 @@ async def handle_key_extension(
             client_id=client_id,
             new_expiry_time=new_expiry,
             total_gb=total_gb,
+            session=session,
             hwid_device_limit=device_limit,
-            session=session
+            reset_traffic=False,
         )
         await update_key_expiry(session, client_id, new_expiry)
         await update_coupon_usage_count(session, coupon.id)
diff --git a/handlers/keys/key_freeze.py b/handlers/keys/key_freeze.py
index 4a8fc2ff..d5f1fa0a 100644
--- a/handlers/keys/key_freeze.py
+++ b/handlers/keys/key_freeze.py
@@ -142,6 +142,7 @@ async def process_callback_unfreeze_subscription_confirm(
             total_gb=total_gb,
             session=session,
             hwid_device_limit=hwid_limit,
+            reset_traffic=False,
         )
 
         text_ok = SUBSCRIPTION_UNFROZEN_MSG
diff --git a/handlers/keys/key_utils.py b/handlers/keys/key_utils.py
index fe6dc3b6..40d4948b 100644
--- a/handlers/keys/key_utils.py
+++ b/handlers/keys/key_utils.py
@@ -287,6 +287,7 @@ async def renew_key_in_cluster(
     total_gb: int,
     session: AsyncSession,
     hwid_device_limit: int = None,
+    reset_traffic: bool = True,
 ):
     try:
         servers = await get_servers(session)
@@ -368,7 +369,8 @@ async def renew_key_in_cluster(
                     )
                     if updated:
                         logger.info(f"Подписка Remnawave {client_id} успешно продлена")
-                        await remna.reset_user_traffic(client_id)
+                        if reset_traffic:
+                            await remna.reset_user_traffic(client_id)
                     else:
                         logger.warning(
                             f"Не удалось продлить подписку Remnawave {client_id}"
diff --git a/handlers/start.py b/handlers/start.py
index cc03bc66..19063ebd 100644
--- a/handlers/start.py
+++ b/handlers/start.py
@@ -19,7 +19,13 @@ from config import (
     SHOW_START_MENU_ONCE,
     SUPPORT_CHAT_URL,
 )
-from database import add_user, check_user_exists, get_trial, get_key_count
+from database import (
+    add_user,
+    check_user_exists,
+    get_trial,
+    get_key_count,
+    get_coupon_by_code,
+)
 from database.models import TrackingSource, User
 from handlers.buttons import (
     ABOUT_VPN,
@@ -175,6 +181,10 @@ async def process_start_logic(
             if "coupons" in part:
                 logger.info(f"Обнаружена ссылка на купон: {part}")
                 coupon_code = part.split("coupons")[1].strip("_")
+                coupon = await get_coupon_by_code(session, coupon_code)
+                if not coupon:
+                    continue
+                    
                 await activate_coupon(
                     message,
                     state,
@@ -183,6 +193,9 @@ async def process_start_logic(
                     admin=admin,
                     user_data=user_data,
                 )
+                
+                if coupon.days:
+                    return
                 continue
 
             if "gift" in part: