feat: Merge traffic collection and user kicking

core/cli.py

@@ -227,7 +227,7 @@ def show_user_uri_json(usernames: list[str]):
 
 
 @cli.command('traffic-status')
-@click.option('--no-gui', is_flag=True, help='Retrieve traffic data without displaying output')
+@click.option('--no-gui', is_flag=True, help='Retrieve traffic data without displaying output and kick expired users')
 def traffic_status(no_gui):
     try:
         cli_api.traffic_status(no_gui=no_gui)

core/cli_api.py

@@ -371,8 +371,12 @@ def show_user_uri_json(usernames: list[str]) -> list[dict[str, Any]] | None:
 
 
 def traffic_status(no_gui=False, display_output=True):
-    '''Fetches traffic status.'''
+    if no_gui:
+        data = traffic.traffic_status(no_gui=True)
+        traffic.kick_expired_users()
+    else:
         data = traffic.traffic_status(no_gui=True if not display_output else no_gui)
+    
     return data
 
 

core/scripts/hysteria2/install.sh

@@ -82,10 +82,10 @@ install_hysteria() {
     chmod +x /etc/hysteria/core/scripts/hysteria2/user.sh
     chmod +x /etc/hysteria/core/scripts/hysteria2/kick.py
 
-    (crontab -l ; echo "*/1 * * * * /bin/bash -c 'source /etc/hysteria/hysteria2_venv/bin/activate && python3 /etc/hysteria/core/cli.py traffic-status' >/dev/null 2>&1") | crontab -
+    (crontab -l ; echo "*/1 * * * * /bin/bash -c 'source /etc/hysteria/hysteria2_venv/bin/activate && python3 /etc/hysteria/core/cli.py traffic-status --no-gui'") | crontab -
-    (crontab -l ; echo "0 3 */3 * * /bin/bash -c 'source /etc/hysteria/hysteria2_venv/bin/activate && python3 /etc/hysteria/core/cli.py restart-hysteria2' >/dev/null 2>&1") | crontab -
+    # (crontab -l ; echo "0 3 */3 * * /bin/bash -c 'source /etc/hysteria/hysteria2_venv/bin/activate && python3 /etc/hysteria/core/cli.py restart-hysteria2' >/dev/null 2>&1") | crontab -
     (crontab -l ; echo "0 */6 * * * /bin/bash -c 'source /etc/hysteria/hysteria2_venv/bin/activate && python3 /etc/hysteria/core/cli.py backup-hysteria' >/dev/null 2>&1") | crontab -
-    (crontab -l ; echo "*/1 * * * * /bin/bash -c 'source /etc/hysteria/hysteria2_venv/bin/activate && python3 /etc/hysteria/core/scripts/hysteria2/kick.py' >/dev/null 2>&1") | crontab -
+    # (crontab -l ; echo "*/1 * * * * /bin/bash -c 'source /etc/hysteria/hysteria2_venv/bin/activate && python3 /etc/hysteria/core/scripts/hysteria2/kick.py' >/dev/null 2>&1") | crontab -
 
 }
 

core/traffic.py

@@ -1,15 +1,49 @@
 #!/usr/bin/env python3
 import json
 import os
-import json
+import sys
+import time
+import fcntl
+import shutil
+import datetime
+from concurrent.futures import ThreadPoolExecutor
 from hysteria2_api import Hysteria2Client
 
-# Define static variables for paths and URLs
 CONFIG_FILE = '/etc/hysteria/config.json'
 USERS_FILE = '/etc/hysteria/users.json'
 API_BASE_URL = 'http://127.0.0.1:25413'
+LOCKFILE = "/tmp/kick.lock"
+BACKUP_FILE = f"{USERS_FILE}.bak"
+MAX_WORKERS = 8
+
+# import logging
+# logging.basicConfig(
+#     level=logging.INFO,
+#     format='%(asctime)s: [%(levelname)s] %(message)s',
+#     datefmt='%Y-%m-%d %H:%M:%S'
+# )
+# logger = logging.getLogger()
+# null_handler = logging.NullHandler()
+# logger.handlers = [null_handler]
+
+def acquire_lock():
+    """Acquires a lock file to prevent concurrent execution"""
+    try:
+        lock_file = open(LOCKFILE, 'w')
+        fcntl.flock(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        return lock_file
+    except IOError:
+        sys.exit(1)
 
 def traffic_status(no_gui=False):
+    """Updates and retrieves traffic statistics for all users.
+    
+    Args:
+        no_gui (bool): If True, suppresses output to console
+        
+    Returns:
+        dict: User data including upload/download bytes and status
+    """
     green = '\033[0;32m'
     cyan = '\033[0;36m'
     NC = '\033[0m'
@@ -19,22 +53,24 @@ def traffic_status(no_gui=False):
             config = json.load(config_file)
             secret = config.get('trafficStats', {}).get('secret')
     except (json.JSONDecodeError, FileNotFoundError) as e:
+        if not no_gui:
             print(f"Error: Failed to read secret from {CONFIG_FILE}. Details: {e}")
-        return
+        return None
 
     if not secret:
+        if not no_gui:
             print("Error: Secret not found in config.json")
-        return
+        return None
 
     client = Hysteria2Client(base_url=API_BASE_URL, secret=secret)
 
     try:
         traffic_stats = client.get_traffic_stats(clear=True)
-        
         online_status = client.get_online_clients()
     except Exception as e:
+        if not no_gui:
             print(f"Error communicating with Hysteria2 API: {e}")
-        return
+        return None
 
     users_data = {}
     if os.path.exists(USERS_FILE):
@@ -42,8 +78,9 @@ def traffic_status(no_gui=False):
             with open(USERS_FILE, 'r') as users_file:
                 users_data = json.load(users_file)
         except json.JSONDecodeError:
+            if not no_gui:
                 print("Error: Failed to parse existing users data JSON file.")
-            return
+            return None
 
     for user in users_data:
         users_data[user]["status"] = "Offline"
@@ -79,6 +116,7 @@ def traffic_status(no_gui=False):
     return users_data
 
 def display_traffic_data(data, green, cyan, NC):
+    """Displays traffic data in a formatted table"""
     if not data:
         print("No traffic data to display.")
         return
@@ -100,6 +138,7 @@ def display_traffic_data(data, green, cyan, NC):
         print("-------------------------------------------------")
 
 def format_bytes(bytes):
+    """Format bytes as human-readable string"""
     if bytes < 1024:
         return f"{bytes}B"
     elif bytes < 1048576:
@@ -111,5 +150,129 @@ def format_bytes(bytes):
     else:
         return f"{bytes / 1099511627776:.2f}TB"
 
+def kick_users(usernames, secret):
+    """Kicks specified users from the server"""
+    try:
+        client = Hysteria2Client(
+            base_url=API_BASE_URL,
+            secret=secret
+        )
+        
+        client.kick_clients(usernames)
+        return True
+    except Exception:
+        return False
+
+def process_user(username, user_data, config_secret, users_data):
+    """Process a single user to check if they should be kicked"""
+    blocked = user_data.get('blocked', False)
+    
+    if blocked:
+        return None
+    
+    max_download_bytes = user_data.get('max_download_bytes', 0)
+    expiration_days = user_data.get('expiration_days', 0)
+    account_creation_date = user_data.get('account_creation_date')
+    current_download_bytes = user_data.get('download_bytes', 0)
+    current_upload_bytes = user_data.get('upload_bytes', 0)
+    
+    total_bytes = current_download_bytes + current_upload_bytes
+    
+    if not account_creation_date:
+        return None
+    
+    try:
+        current_date = datetime.datetime.now().timestamp()
+        creation_date = datetime.datetime.fromisoformat(account_creation_date.replace('Z', '+00:00'))
+        expiration_date = (creation_date + datetime.timedelta(days=expiration_days)).timestamp()
+        
+        should_block = False
+        
+        if max_download_bytes > 0 and total_bytes >= 0 and expiration_days > 0:
+            if total_bytes >= max_download_bytes or current_date >= expiration_date:
+                should_block = True
+                
+            if should_block:
+                users_data[username]['blocked'] = True
+                return username
+        
+    except Exception:
+        return None
+    
+    return None
+
+def kick_expired_users():
+    """Kicks users who have exceeded their data limits or whose accounts have expired"""
+    lock_file = acquire_lock()
+    
+    try:
+        shutil.copy2(USERS_FILE, BACKUP_FILE)
+        
+        try:
+            with open(CONFIG_FILE, 'r') as f:
+                config = json.load(f)
+                secret = config.get('trafficStats', {}).get('secret', '')
+                if not secret:
+                    sys.exit(1)
+        except Exception:
+            shutil.copy2(BACKUP_FILE, USERS_FILE)
+            sys.exit(1)
+            
+        try:
+            with open(USERS_FILE, 'r') as f:
+                users_data = json.load(f)
+        except json.JSONDecodeError:
+            shutil.copy2(BACKUP_FILE, USERS_FILE)
+            sys.exit(1)
+        except Exception:
+            shutil.copy2(BACKUP_FILE, USERS_FILE)
+            sys.exit(1)
+            
+        users_to_kick = []
+        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
+            future_to_user = {
+                executor.submit(process_user, username, user_data, secret, users_data): username
+                for username, user_data in users_data.items()
+            }
+            
+            for future in future_to_user:
+                username = future.result()
+                if username:
+                    users_to_kick.append(username)
+        
+        if users_to_kick:
+            for retry in range(3):
+                try:
+                    with open(USERS_FILE, 'w') as f:
+                        json.dump(users_data, f, indent=2)
+                    break 
+                except Exception:
+                    time.sleep(1)
+                    if retry == 2: 
+                        raise
+        
+        if users_to_kick:
+            batch_size = 50 
+            for i in range(0, len(users_to_kick), batch_size):
+                batch = users_to_kick[i:i+batch_size]
+                kick_users(batch, secret)
+                        
+    except Exception:
+        shutil.copy2(BACKUP_FILE, USERS_FILE)
+        sys.exit(1)
+    finally:
+        fcntl.flock(lock_file, fcntl.LOCK_UN)
+        lock_file.close()
+
 if __name__ == "__main__":
-    traffic_status()
+    if len(sys.argv) > 1:
+        if sys.argv[1] == "kick":
+            kick_expired_users()
+        elif sys.argv[1] == "--no-gui":
+            traffic_status(no_gui=True)
+            kick_expired_users()
+        else:
+            print(f"Unknown argument: {sys.argv[1]}")
+            print("Usage: python traffic.py [kick|--no-gui]")
+    else:
+        traffic_status(no_gui=False)