From 4bf5b3a87dafbf1d2f10b7c00c684cb89fef5dde Mon Sep 17 00:00:00 2001
From: Vladless <curov2014@yandex.ru>
Date: Sun, 10 Nov 2024 03:31:19 +0300
Subject: [PATCH] reset traffic

---
 client.py                  | 32 +++++++++++++++++++++++++++++---
 handlers/keys/key_utils.py | 23 ++++++++++++++++++-----
 handlers/keys/keys.py      |  2 +-
 3 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/client.py b/client.py
index 945dd0a1..3aca9b5a 100644
--- a/client.py
+++ b/client.py
@@ -59,6 +59,29 @@ async def add_client(
             return None
 
 
+async def reset_client_traffic(session, server_id: str, email: str) -> bool:
+    """Сбрасывает трафик клиента на сервере"""
+    api_url = SERVERS[server_id]["API_URL"]
+    url = f"{api_url}/panel/api/inbounds/1/resetClientTraffic/{email}"
+    headers = {"Accept": "application/json"}
+
+    try:
+        async with session.post(url, headers=headers) as response:
+            if response.status == 200:
+                logger.info(
+                    f"Трафик клиента {email} успешно сброшен на сервере {server_id}")
+                return True
+            else:
+                logger.error(
+                    f"Ошибка при сбросе трафика клиента {email} на сервере {server_id}: {response.status} - {await response.text()}"
+                )
+                return False
+    except Exception as e:
+        logger.error(
+            f"Ошибка при попытке сброса трафика клиента {email} на сервере {server_id}: {e}")
+        return False
+
+
 async def extend_client_key(
     session, server_id: str, tg_id, client_id, email: str, new_expiry_time: int
 ) -> bool:
@@ -113,7 +136,8 @@ async def extend_client_key(
             ),
         }
 
-        headers = {"Content-Type": "application/json", "Accept": "application/json"}
+        headers = {"Content-Type": "application/json",
+                   "Accept": "application/json"}
 
         try:
             async with session.post(
@@ -122,7 +146,8 @@ async def extend_client_key(
                 headers=headers,
             ) as response:
                 logger.info(f"POST {response.url} Status: {response.status}")
-                logger.info(f"POST Request Data: {json.dumps(payload, indent=2)}")
+                logger.info(
+                    f"POST Request Data: {json.dumps(payload, indent=2)}")
                 response_text = await response.text()
                 logger.info(f"POST Response: {response_text}")
 
@@ -165,7 +190,8 @@ async def extend_client_key_admin(
         ),
     }
 
-    headers = {"Content-Type": "application/json", "Accept": "application/json"}
+    headers = {"Content-Type": "application/json",
+               "Accept": "application/json"}
 
     try:
         async with session.post(
diff --git a/handlers/keys/key_utils.py b/handlers/keys/key_utils.py
index 5545fd07..47b43614 100644
--- a/handlers/keys/key_utils.py
+++ b/handlers/keys/key_utils.py
@@ -2,8 +2,8 @@ import asyncpg
 from loguru import logger
 
 from auth import login_with_credentials
-from client import add_client, delete_client, extend_client_key
-from config import ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_URL
+from client import add_client, delete_client, extend_client_key, reset_client_traffic
+from config import ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_URL, RESET_TRAFFIC
 
 
 async def create_key_on_server(server_id, tg_id, client_id, email, expiry_timestamp):
@@ -34,17 +34,28 @@ async def create_key_on_server(server_id, tg_id, client_id, email, expiry_timest
         logger.error(f"Ошибка на сервере {server_id}: {e}")
 
 
-async def renew_server_key(server_id, tg_id, client_id, email, new_expiry_time):
+async def renew_server_key(
+        server_id,
+        tg_id, client_id,
+        email,
+        new_expiry_time,
+        reset_traffic=RESET_TRAFFIC
+):
     try:
         session = await login_with_credentials(
             server_id, ADMIN_USERNAME, ADMIN_PASSWORD
         )
+
         await extend_client_key(
             session, server_id, tg_id, client_id, email, new_expiry_time
         )
+
+        if reset_traffic:
+            await reset_client_traffic(session, server_id, email)
+
     except Exception as e:
         logger.error(
-            f"Не удалось продлить ключ {client_id} на сервере {server_id}: {e}"
+            f"Не удалось продлить ключ {client_id} и сбросить трафик на сервере {server_id}: {e}"
         )
 
 
@@ -63,7 +74,9 @@ async def delete_key_from_db(client_id):
 async def delete_key_from_server(server_id, client_id):
     """Удаление ключа с сервера"""
     try:
-        async with login_with_credentials(server_id, ADMIN_USERNAME, ADMIN_PASSWORD) as session:
+        async with login_with_credentials(server_id,
+                                          ADMIN_USERNAME,
+                                          ADMIN_PASSWORD) as session:
             success = await delete_client(session, server_id, client_id)
 
             if not success:
diff --git a/handlers/keys/keys.py b/handlers/keys/keys.py
index 4581df38..57b3c00a 100644
--- a/handlers/keys/keys.py
+++ b/handlers/keys/keys.py
@@ -96,7 +96,7 @@ async def process_callback_view_keys(callback_query: types.CallbackQuery):
             else:
                 response_message = NO_KEYS
                 create_key_button = types.InlineKeyboardButton(
-                    text="➕ Создать ключ", callback_data="create_key"
+                    text="➕ Создать подписку", callback_data="create_key"
                 )
                 back_button = types.InlineKeyboardButton(
                     text="🔙 Назад", callback_data="view_profile"