From dc49f5d74f8f152dd05caf82ebe975bc99dc42cb Mon Sep 17 00:00:00 2001 From: Whispering Wind <151555003+ReturnFI@users.noreply.github.com> Date: Sat, 10 May 2025 14:06:55 +0330 Subject: [PATCH] feat: Merge traffic collection and user kicking --- core/cli.py | 2 +- core/cli_api.py | 8 +- core/scripts/hysteria2/install.sh | 6 +- core/traffic.py | 187 ++++++++++++++++++++++++++++-- 4 files changed, 185 insertions(+), 18 deletions(-) diff --git a/core/cli.py b/core/cli.py index ea964d1..8c684f6 100644 --- a/core/cli.py +++ b/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) diff --git a/core/cli_api.py b/core/cli_api.py index 219f050..4bbc44e 100644 --- a/core/cli_api.py +++ b/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.''' - data = traffic.traffic_status(no_gui=True if not display_output else no_gui) + 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 diff --git a/core/scripts/hysteria2/install.sh b/core/scripts/hysteria2/install.sh index 767f1cd..761ed4d 100644 --- a/core/scripts/hysteria2/install.sh +++ b/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 "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 "*/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 */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 - } diff --git a/core/traffic.py b/core/traffic.py index 251b14c..3ec1363 100644 --- a/core/traffic.py +++ b/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: - print(f"Error: Failed to read secret from {CONFIG_FILE}. Details: {e}") - return + if not no_gui: + print(f"Error: Failed to read secret from {CONFIG_FILE}. Details: {e}") + return None if not secret: - print("Error: Secret not found in config.json") - return + if not no_gui: + print("Error: Secret not found in config.json") + 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: - print(f"Error communicating with Hysteria2 API: {e}") - return + if not no_gui: + print(f"Error communicating with Hysteria2 API: {e}") + 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: - print("Error: Failed to parse existing users data JSON file.") - return + if not no_gui: + print("Error: Failed to parse existing users data JSON file.") + 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) \ No newline at end of file