From 80c21ccd070aaf7f429ef479a36aedb5dedfa2c4 Mon Sep 17 00:00:00 2001 From: Whispering Wind <151555003+ReturnFI@users.noreply.github.com> Date: Sat, 6 Sep 2025 22:38:32 +0330 Subject: [PATCH] feat: migrate user management from json to mongodb --- core/cli_api.py | 57 +++-- core/scripts/db/database.py | 41 ++++ core/scripts/db/install.sh | 54 +++++ core/scripts/hysteria2/add_user.py | 97 +++----- core/scripts/hysteria2/backup.py | 68 +++++- core/scripts/hysteria2/bulk_users.py | 122 +++++----- core/scripts/hysteria2/edit_user.py | 124 ++++++++++ core/scripts/hysteria2/edit_user.sh | 268 --------------------- core/scripts/hysteria2/get_user.py | 43 ++-- core/scripts/hysteria2/kick.py | 204 +++++++--------- core/scripts/hysteria2/list_users.py | 46 ++-- core/scripts/hysteria2/list_users.sh | 6 - core/scripts/hysteria2/remove_user.py | 38 +-- core/scripts/hysteria2/reset_user.py | 55 +++-- core/scripts/hysteria2/restore.py | 234 ++++++++---------- core/scripts/hysteria2/wrapper_uri.py | 30 ++- core/traffic.py | 331 ++++++++++++-------------- 17 files changed, 824 insertions(+), 994 deletions(-) create mode 100644 core/scripts/db/database.py create mode 100644 core/scripts/db/install.sh create mode 100644 core/scripts/hysteria2/edit_user.py delete mode 100644 core/scripts/hysteria2/edit_user.sh delete mode 100644 core/scripts/hysteria2/list_users.sh diff --git a/core/cli_api.py b/core/cli_api.py index 07e30cb..e73117b 100644 --- a/core/cli_api.py +++ b/core/cli_api.py @@ -28,7 +28,7 @@ class Command(Enum): GET_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'get_user.py') ADD_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'add_user.py') BULK_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'bulk_users.py') - EDIT_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'edit_user.sh') + EDIT_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'edit_user.py') RESET_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'reset_user.py') REMOVE_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'remove_user.py') SHOW_USER_URI = os.path.join(SCRIPT_DIR, 'hysteria2', 'show_user_uri.py') @@ -307,43 +307,40 @@ def bulk_user_add(traffic_gb: float, expiration_days: int, count: int, prefix: s def edit_user(username: str, new_username: str | None, new_traffic_limit: int | None, new_expiration_days: int | None, renew_password: bool, renew_creation_date: bool, blocked: bool | None, unlimited_ip: bool | None): ''' - Edits an existing user's details. + Edits an existing user's details by calling the new edit_user.py script with named flags. ''' if not username: raise InvalidInputError('Error: username is required') - if new_traffic_limit is not None and new_traffic_limit < 0: - raise InvalidInputError('Error: traffic limit must be a non-negative number.') - if new_expiration_days is not None and new_expiration_days < 0: - raise InvalidInputError('Error: expiration days must be a non-negative number.') + command_args = ['python3', Command.EDIT_USER.value, username] - password = generate_password() if renew_password else '' - creation_date = datetime.now().strftime('%Y-%m-%d') if renew_creation_date else '' + if new_username: + command_args.extend(['--new-username', new_username]) - blocked_str = '' - if blocked is True: - blocked_str = 'true' - elif blocked is False: - blocked_str = 'false' + if new_traffic_limit is not None: + if new_traffic_limit < 0: + raise InvalidInputError('Error: traffic limit must be a non-negative number.') + command_args.extend(['--traffic-gb', str(new_traffic_limit)]) - unlimited_str = '' - if unlimited_ip is True: - unlimited_str = 'true' - elif unlimited_ip is False: - unlimited_str = 'false' + if new_expiration_days is not None: + if new_expiration_days < 0: + raise InvalidInputError('Error: expiration days must be a non-negative number.') + command_args.extend(['--expiration-days', str(new_expiration_days)]) + + if renew_password: + password = generate_password() + command_args.extend(['--password', password]) + + if renew_creation_date: + creation_date = datetime.now().strftime('%Y-%m-%d') + command_args.extend(['--creation-date', creation_date]) + + if blocked is not None: + command_args.extend(['--blocked', 'true' if blocked else 'false']) + + if unlimited_ip is not None: + command_args.extend(['--unlimited', 'true' if unlimited_ip else 'false']) - command_args = [ - 'bash', - Command.EDIT_USER.value, - username, - new_username or '', - str(new_traffic_limit) if new_traffic_limit is not None else '', - str(new_expiration_days) if new_expiration_days is not None else '', - password, - creation_date, - blocked_str, - unlimited_str - ] run_cmd(command_args) diff --git a/core/scripts/db/database.py b/core/scripts/db/database.py new file mode 100644 index 0000000..30bb403 --- /dev/null +++ b/core/scripts/db/database.py @@ -0,0 +1,41 @@ +import pymongo +from bson.objectid import ObjectId + +class Database: + def __init__(self, db_name="blitz_panel", collection_name="users"): + try: + self.client = pymongo.MongoClient("mongodb://localhost:27017/") + self.db = self.client[db_name] + self.collection = self.db[collection_name] + self.client.server_info() + except pymongo.errors.ConnectionFailure as e: + print(f"Could not connect to MongoDB: {e}") + raise + + def add_user(self, user_data): + username = user_data.pop('username', None) + if not username: + raise ValueError("Username is required") + + if self.collection.find_one({"_id": username.lower()}): + return None + + user_data['_id'] = username.lower() + return self.collection.insert_one(user_data) + + def get_user(self, username): + return self.collection.find_one({"_id": username.lower()}) + + def get_all_users(self): + return list(self.collection.find({})) + + def update_user(self, username, updates): + return self.collection.update_one({"_id": username.lower()}, {"$set": updates}) + + def delete_user(self, username): + return self.collection.delete_one({"_id": username.lower()}) + +try: + db = Database() +except pymongo.errors.ConnectionFailure: + db = None \ No newline at end of file diff --git a/core/scripts/db/install.sh b/core/scripts/db/install.sh new file mode 100644 index 0000000..4b65a87 --- /dev/null +++ b/core/scripts/db/install.sh @@ -0,0 +1,54 @@ +#!/bin/bash + +set -e + +echo "Updating package list..." +sudo apt-get update -y + +if ! command -v mongod &> /dev/null +then + echo "MongoDB not found. Installing from official repository..." + + echo "Installing prerequisites..." + sudo apt-get install -y gnupg curl + + echo "Importing the MongoDB public GPG key..." + curl -fsSL https://www.mongodb.org/static/pgp/server-7.0.asc | \ + sudo gpg -o /usr/share/keyrings/mongodb-server-7.0.gpg --dearmor + + echo "Creating a list file for MongoDB..." + echo "deb [ arch=amd64,arm64 signed-by=/usr/share/keyrings/mongodb-server-7.0.gpg ] https://repo.mongodb.org/apt/debian bookworm/mongodb-org/7.0 main" | \ + sudo tee /etc/apt/sources.list.d/mongodb-org-7.0.list + + echo "Reloading local package database..." + sudo apt-get update -y + + echo "Installing the MongoDB packages..." + sudo apt-get install -y mongodb-org +else + echo "MongoDB is already installed." +fi + +echo "Starting and enabling MongoDB service..." +sudo systemctl start mongod +sudo systemctl enable mongod + +echo "Installing pymongo for Python..." +if python3 -m pip --version &> /dev/null; then + python3 -m pip install pymongo +elif pip --version &> /dev/null; then + pip install pymongo +else + echo "pip is not installed. Please install python3-pip" + exit 1 +fi + +# echo "Installing MongoDB driver for Go..." +# if command -v go &> /dev/null +# then +# go get go.mongodb.org/mongo-driver/mongo +# else +# echo "Go is not installed. Skipping Go driver installation." +# fi + +echo "Database setup is complete." \ No newline at end of file diff --git a/core/scripts/hysteria2/add_user.py b/core/scripts/hysteria2/add_user.py index d9853e7..f66301c 100644 --- a/core/scripts/hysteria2/add_user.py +++ b/core/scripts/hysteria2/add_user.py @@ -1,33 +1,22 @@ #!/usr/bin/env python3 -import json import sys import os import subprocess import re from datetime import datetime -from init_paths import * -from paths import * +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +from db.database import db def add_user(username, traffic_gb, expiration_days, password=None, creation_date=None, unlimited_user=False): - """ - Adds a new user to the USERS_FILE. - - Args: - username (str): The username to add. - traffic_gb (str): The traffic limit in GB. - expiration_days (str): The number of days until the account expires. - password (str, optional): The user's password. If None, a random one is generated. - creation_date (str, optional): The account creation date in YYYY-MM-DD format. Defaults to None. - unlimited_user (bool, optional): If True, user is exempt from IP limits. Defaults to False. - - Returns: - int: 0 on success, 1 on failure. - """ if not username or not traffic_gb or not expiration_days: print(f"Usage: {sys.argv[0]} [password] [creation_date] [unlimited_user (true/false)]") return 1 + if db is None: + print("Error: Database connection failed. Please ensure MongoDB is running and configured.") + return 1 + try: traffic_bytes = int(float(traffic_gb) * 1073741824) expiration_days = int(expiration_days) @@ -48,61 +37,45 @@ def add_user(username, traffic_gb, expiration_days, password=None, creation_date print("Error: Failed to generate password. Please install 'pwgen' or ensure /proc access.") return 1 - if creation_date: - if not re.match(r"^[0-9]{4}-[0-9]{2}-[0-9]{2}$", creation_date): - print("Invalid date format. Expected YYYY-MM-DD.") - return 1 - try: - datetime.strptime(creation_date, "%Y-%m-%d") - except ValueError: - print("Invalid date. Please provide a valid date in YYYY-MM-DD format.") - return 1 - else: - creation_date = None - if not re.match(r"^[a-zA-Z0-9_]+$", username): print("Error: Username can only contain letters, numbers, and underscores.") return 1 - if not os.path.isfile(USERS_FILE): - try: - with open(USERS_FILE, 'w') as f: - json.dump({}, f) - except IOError: - print(f"Error: Could not create {USERS_FILE}.") + try: + if db.get_user(username_lower): + print("User already exists.") return 1 - try: - with open(USERS_FILE, 'r+') as f: + user_data = { + "username": username_lower, + "password": password, + "max_download_bytes": traffic_bytes, + "expiration_days": expiration_days, + "blocked": False, + "unlimited_user": unlimited_user + } + + if creation_date: + if not re.match(r"^[0-9]{4}-[0-9]{2}-[0-9]{2}$", creation_date): + print("Invalid date format. Expected YYYY-MM-DD.") + return 1 try: - users_data = json.load(f) - except json.JSONDecodeError: - print(f"Error: {USERS_FILE} contains invalid JSON.") + datetime.strptime(creation_date, "%Y-%m-%d") + user_data["account_creation_date"] = creation_date + except ValueError: + print("Invalid date. Please provide a valid date in YYYY-MM-DD format.") return 1 - for existing_username in users_data: - if existing_username.lower() == username_lower: - print("User already exists.") - return 1 + result = db.add_user(user_data) + if result: + print(f"User {username} added successfully.") + return 0 + else: + print(f"Error: Failed to add user {username}.") + return 1 - users_data[username_lower] = { - "password": password, - "max_download_bytes": traffic_bytes, - "expiration_days": expiration_days, - "account_creation_date": creation_date, - "blocked": False, - "unlimited_user": unlimited_user - } - - f.seek(0) - json.dump(users_data, f, indent=4) - f.truncate() - - print(f"User {username} added successfully.") - return 0 - - except IOError: - print(f"Error: Could not write to {USERS_FILE}.") + except Exception as e: + print(f"An error occurred: {e}") return 1 if __name__ == "__main__": diff --git a/core/scripts/hysteria2/backup.py b/core/scripts/hysteria2/backup.py index 01dd60f..fa419bb 100644 --- a/core/scripts/hysteria2/backup.py +++ b/core/scripts/hysteria2/backup.py @@ -1,27 +1,69 @@ #!/usr/bin/env python3 import zipfile +import subprocess +import shutil from pathlib import Path from datetime import datetime -backup_dir = Path("/opt/hysbackup") -backup_file = backup_dir / f"hysteria_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip" +# --- Configuration --- +DB_NAME = "blitz_panel" +BACKUP_ROOT_DIR = Path("/opt/hysbackup") +TIMESTAMP = datetime.now().strftime('%Y%m%d_%H%M%S') +BACKUP_FILENAME = BACKUP_ROOT_DIR / f"hysteria_backup_{TIMESTAMP}.zip" +TEMP_DUMP_DIR = BACKUP_ROOT_DIR / f"mongodump_{TIMESTAMP}" -files_to_backup = [ +FILES_TO_BACKUP = [ Path("/etc/hysteria/ca.key"), Path("/etc/hysteria/ca.crt"), - Path("/etc/hysteria/users.json"), Path("/etc/hysteria/config.json"), Path("/etc/hysteria/.configs.env"), ] -backup_dir.mkdir(parents=True, exist_ok=True) +def create_backup(): + """Dumps the MongoDB database and zips it with config files.""" + try: + BACKUP_ROOT_DIR.mkdir(parents=True, exist_ok=True) + TEMP_DUMP_DIR.mkdir(parents=True) -try: - with zipfile.ZipFile(backup_file, 'w') as zipf: - for file_path in files_to_backup: - if file_path.exists(): - zipf.write(file_path, arcname=file_path.name) - print("Backup successfully created") -except Exception as e: - print("Backup failed!", str(e)) + print(f"Dumping database '{DB_NAME}'...") + mongodump_cmd = [ + "mongodump", + f"--db={DB_NAME}", + f"--out={TEMP_DUMP_DIR}" + ] + subprocess.run(mongodump_cmd, check=True, capture_output=True) + print("Database dump successful.") + + print(f"Creating backup archive: {BACKUP_FILENAME}") + with zipfile.ZipFile(BACKUP_FILENAME, 'w', zipfile.ZIP_DEFLATED) as zipf: + for file_path in FILES_TO_BACKUP: + if file_path.exists() and file_path.is_file(): + zipf.write(file_path, arcname=file_path.name) + print(f" - Added {file_path.name}") + else: + print(f" - Warning: Skipping missing file {file_path}") + + dump_content_path = TEMP_DUMP_DIR / DB_NAME + if dump_content_path.exists(): + for file_path in dump_content_path.rglob('*'): + arcname = file_path.relative_to(TEMP_DUMP_DIR) + zipf.write(file_path, arcname=arcname) + print(f" - Added database dump for '{DB_NAME}'") + + print("\nBackup successfully created.") + + except FileNotFoundError: + print("\nBackup failed! 'mongodump' command not found. Is MongoDB installed and in your PATH?") + except subprocess.CalledProcessError as e: + print("\nBackup failed! Error during mongodump.") + print(f" - Stderr: {e.stderr.decode().strip()}") + except Exception as e: + print(f"\nBackup failed! An unexpected error occurred: {e}") + finally: + if TEMP_DUMP_DIR.exists(): + shutil.rmtree(TEMP_DUMP_DIR) + print("Temporary dump directory cleaned up.") + +if __name__ == "__main__": + create_backup() \ No newline at end of file diff --git a/core/scripts/hysteria2/bulk_users.py b/core/scripts/hysteria2/bulk_users.py index 1fdd23f..592c1a4 100644 --- a/core/scripts/hysteria2/bulk_users.py +++ b/core/scripts/hysteria2/bulk_users.py @@ -1,97 +1,83 @@ #!/usr/bin/env python3 -import json import sys import os import subprocess import argparse import re -from datetime import datetime -from init_paths import * -from paths import * + +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +from db.database import db def add_bulk_users(traffic_gb, expiration_days, count, prefix, start_number, unlimited_user): + if db is None: + print("Error: Database connection failed. Please ensure MongoDB is running.") + return 1 + try: traffic_bytes = int(float(traffic_gb) * 1073741824) except ValueError: print("Error: Traffic limit must be a numeric value.") return 1 - if not os.path.isfile(USERS_FILE): - try: - with open(USERS_FILE, 'w') as f: - json.dump({}, f) - except IOError: - print(f"Error: Could not create {USERS_FILE}.") + potential_usernames = [] + for i in range(count): + username = f"{prefix}{start_number + i}" + if not re.match(r"^[a-zA-Z0-9_]+$", username): + print(f"Error: Generated username '{username}' contains invalid characters. Aborting.") return 1 + potential_usernames.append(username.lower()) try: - with open(USERS_FILE, 'r+') as f: - try: - users_data = json.load(f) - except json.JSONDecodeError: - print(f"Error: {USERS_FILE} contains invalid JSON.") - return 1 + existing_docs = db.collection.find({"_id": {"$in": potential_usernames}}, {"_id": 1}) + existing_users_set = {doc['_id'] for doc in existing_docs} + except Exception as e: + print(f"Error querying database for existing users: {e}") + return 1 + + new_usernames = [u for u in potential_usernames if u not in existing_users_set] + new_users_count = len(new_usernames) - existing_users_lower = {u.lower() for u in users_data} - new_users_to_add = {} - creation_date = None - - try: - password_process = subprocess.run(['pwgen', '-s', '32', str(count)], capture_output=True, text=True, check=True) - passwords = password_process.stdout.strip().split('\n') - except (FileNotFoundError, subprocess.CalledProcessError): - print("Warning: 'pwgen' not found or failed. Falling back to UUID for password generation.") - passwords = [subprocess.check_output(['cat', '/proc/sys/kernel/random/uuid'], text=True).strip() for _ in range(count)] - - if len(passwords) < count: - print("Error: Could not generate enough passwords.") - return 1 - - for i in range(count): - username = f"{prefix}{start_number + i}" - username_lower = username.lower() - - if not re.match(r"^[a-zA-Z0-9_]+$", username_lower): - print(f"Error: Generated username '{username}' contains invalid characters. Use only letters, numbers, and underscores.") - continue - - if username_lower in existing_users_lower or username_lower in new_users_to_add: - print(f"Warning: User '{username}' already exists. Skipping.") - continue - - new_users_to_add[username_lower] = { - "password": passwords[i], - "max_download_bytes": traffic_bytes, - "expiration_days": expiration_days, - "account_creation_date": creation_date, - "blocked": False, - "unlimited_user": unlimited_user - } - # print(f"Preparing to add user: {username}") - - if not new_users_to_add: - print("No new users to add.") - return 0 - - users_data.update(new_users_to_add) - - f.seek(0) - json.dump(users_data, f, indent=4) - f.truncate() - - print(f"\nSuccessfully added {len(new_users_to_add)} users.") + if new_users_count == 0: + print("No new users to add. All generated usernames already exist.") return 0 - except IOError: - print(f"Error: Could not read or write to {USERS_FILE}.") + if count > new_users_count: + print(f"Warning: {count - new_users_count} user(s) already exist. Skipping them.") + + try: + password_process = subprocess.run(['pwgen', '-s', '32', str(new_users_count)], capture_output=True, text=True, check=True) + passwords = password_process.stdout.strip().split('\n') + except (FileNotFoundError, subprocess.CalledProcessError): + print("Warning: 'pwgen' not found or failed. Falling back to UUID for password generation.") + passwords = [subprocess.check_output(['cat', '/proc/sys/kernel/random/uuid'], text=True).strip() for _ in range(new_users_count)] + + if len(passwords) < new_users_count: + print("Error: Could not generate enough passwords.") return 1 + + users_to_insert = [] + for i, username in enumerate(new_usernames): + user_doc = { + "_id": username, + "password": passwords[i], + "max_download_bytes": traffic_bytes, + "expiration_days": expiration_days, + "blocked": False, + "unlimited_user": unlimited_user + } + users_to_insert.append(user_doc) + + try: + db.collection.insert_many(users_to_insert, ordered=False) + print(f"\nSuccessfully added {len(users_to_insert)} new users.") + return 0 except Exception as e: - print(f"An unexpected error occurred: {e}") + print(f"An unexpected error occurred during database insert: {e}") return 1 if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Add bulk users to Hysteria2.") + parser = argparse.ArgumentParser(description="Add bulk users to Hysteria2 via database.") parser.add_argument("-t", "--traffic-gb", dest="traffic_gb", type=float, required=True, help="Traffic limit for each user in GB.") parser.add_argument("-e", "--expiration-days", dest="expiration_days", type=int, required=True, help="Expiration duration for each user in days.") parser.add_argument("-c", "--count", type=int, required=True, help="Number of users to create.") diff --git a/core/scripts/hysteria2/edit_user.py b/core/scripts/hysteria2/edit_user.py new file mode 100644 index 0000000..8b27f40 --- /dev/null +++ b/core/scripts/hysteria2/edit_user.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 + +import sys +import os +import argparse +import re +from datetime import datetime + +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +from db.database import db + +def edit_user(username, new_username=None, new_password=None, traffic_gb=None, expiration_days=None, creation_date=None, blocked=None, unlimited_user=None): + if db is None: + print("Error: Database connection failed.", file=sys.stderr) + return 1 + + username_lower = username.lower() + try: + user_data = db.get_user(username_lower) + if not user_data: + print(f"Error: User '{username}' not found.", file=sys.stderr) + return 1 + except Exception as e: + print(f"Error fetching user data: {e}", file=sys.stderr) + return 1 + + updates = {} + + if new_password: + updates['password'] = new_password + + if traffic_gb is not None: + updates['max_download_bytes'] = int(float(traffic_gb) * 1073741824) + + if expiration_days is not None: + updates['expiration_days'] = int(expiration_days) + + if creation_date is not None: + if creation_date.lower() == 'null': + updates['account_creation_date'] = None + else: + updates['account_creation_date'] = creation_date + + if blocked is not None: + updates['blocked'] = blocked + + if unlimited_user is not None: + updates['unlimited_user'] = unlimited_user + + try: + if updates: + db.update_user(username_lower, updates) + print(f"User '{username}' attributes updated successfully.") + + if new_username and new_username.lower() != username_lower: + new_username_lower = new_username.lower() + if db.get_user(new_username_lower): + print(f"Error: Target username '{new_username}' already exists.", file=sys.stderr) + return 1 + + updated_user_data = db.get_user(username_lower) + + updated_user_data.pop('_id') + updated_user_data['_id'] = new_username_lower + + db.collection.insert_one(updated_user_data) + db.delete_user(username_lower) + print(f"User '{username}' successfully renamed to '{new_username}'.") + + elif not updates and not (new_username and new_username.lower() != username_lower): + print("No changes specified.") + + except Exception as e: + print(f"An error occurred during update: {e}", file=sys.stderr) + return 1 + + return 0 + + +def str_to_bool(val): + if val.lower() in ('true', 'y', '1'): + return True + elif val.lower() in ('false', 'n', '0'): + return False + raise argparse.ArgumentTypeError('Boolean value expected (true/false, y/n, 1/0).') + +def validate_date(date_str): + if date_str.lower() == 'null': + return date_str + try: + datetime.strptime(date_str, "%Y-%m-%d") + return date_str + except ValueError: + raise argparse.ArgumentTypeError("Invalid date format. Expected YYYY-MM-DD or 'null'.") + +def validate_username(username): + if not re.match(r"^[a-zA-Z0-9_]+$", username): + raise argparse.ArgumentTypeError("Username can only contain letters, numbers, and underscores.") + return username + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Edit a Hysteria2 user's details in the database.") + parser.add_argument("username", type=validate_username, help="The current username of the user to edit.") + parser.add_argument("--new-username", dest="new_username", type=validate_username, help="New username for the user.") + parser.add_argument("--password", dest="new_password", help="New password for the user.") + parser.add_argument("--traffic-gb", dest="traffic_gb", type=float, help="New traffic limit in GB (e.g., 50). Use 0 for unlimited.") + parser.add_argument("--expiration-days", dest="expiration_days", type=int, help="New expiration in days from creation date (e.g., 30). Use 0 for unlimited.") + parser.add_argument("--creation-date", dest="creation_date", type=validate_date, help="New creation date in YYYY-MM-DD format, or 'null' to reset to On-hold.") + parser.add_argument("--blocked", type=str_to_bool, help="Set blocked status (true/false).") + parser.add_argument("--unlimited", dest="unlimited_user", type=str_to_bool, help="Set unlimited user status for IP limits (true/false).") + + args = parser.parse_args() + + sys.exit(edit_user( + username=args.username, + new_username=args.new_username, + new_password=args.new_password, + traffic_gb=args.traffic_gb, + expiration_days=args.expiration_days, + creation_date=args.creation_date, + blocked=args.blocked, + unlimited_user=args.unlimited_user + )) \ No newline at end of file diff --git a/core/scripts/hysteria2/edit_user.sh b/core/scripts/hysteria2/edit_user.sh deleted file mode 100644 index c006bfd..0000000 --- a/core/scripts/hysteria2/edit_user.sh +++ /dev/null @@ -1,268 +0,0 @@ -#!/bin/bash - -source /etc/hysteria/core/scripts/utils.sh -source /etc/hysteria/core/scripts/path.sh - -readonly GB_TO_BYTES=$((1024 * 1024 * 1024)) - -validate_username() { - local username=$1 - if [ -z "$username" ]; then - return 0 - fi - if ! [[ "$username" =~ ^[a-zA-Z0-9]+$ ]]; then - echo "Username can only contain letters and numbers." - return 1 - fi - return 0 -} - - -validate_traffic_limit() { - local traffic_limit=$1 - if [ -z "$traffic_limit" ]; then - return 0 - fi - if ! [[ "$traffic_limit" =~ ^[0-9]+$ ]]; then - echo "Error: Traffic limit must be a valid non-negative number (use 0 for unlimited)." - return 1 - fi - return 0 -} - -validate_expiration_days() { - local expiration_days=$1 - if [ -z "$expiration_days" ]; then - return 0 - fi - if ! [[ "$expiration_days" =~ ^[0-9]+$ ]]; then - echo "Error: Expiration days must be a valid non-negative number (use 0 for unlimited)." - return 1 - fi - return 0 -} - -validate_date() { - local date_str=$1 - if [ -z "$date_str" ]; then - return 0 - fi - if [ "$date_str" == "null" ]; then - return 0 - fi - if ! [[ "$date_str" =~ ^[0-9]{4}-[0-9]{2}-[0-9]{2}$ ]]; then - echo "Invalid date format. Expected YYYY-MM-DD." - return 1 - elif ! date -d "$date_str" >/dev/null 2>&1; then - echo "Invalid date. Please provide a valid date in YYYY-MM-DD format." - return 1 - fi - return 0 -} - -validate_blocked_status() { - local status=$1 - if [ -z "$status" ]; then - return 0 - fi - if [ "$status" != "true" ] && [ "$status" != "false" ]; then - echo "Blocked status must be 'true' or 'false'." - return 1 - fi - return 0 -} - -validate_unlimited_status() { - local status=$1 - if [ -z "$status" ]; then - return 0 - fi - if [ "$status" != "true" ] && [ "$status" != "false" ]; then - echo "Unlimited status must be 'true' or 'false'." - return 1 - fi - return 0 -} - - -convert_boolean_status() { - local status=$1 - case "$status" in - y|Y) echo "true" ;; - n|N) echo "false" ;; - true|false) echo "$status" ;; - *) echo "false" ;; - esac -} - -get_user_info() { - local username=$1 - python3 $CLI_PATH get-user -u "$username" -} - -update_user_info() { - local old_username=$1 - local new_username=$2 - local new_password=$3 - local new_max_download_bytes=$4 - local new_expiration_days=$5 - local new_account_creation_date=$6 - local new_blocked=$7 - local new_unlimited=$8 - - if [ ! -f "$USERS_FILE" ]; then - echo -e "${red}Error:${NC} File '$USERS_FILE' not found." - return 1 - fi - - - user_exists=$(jq -e --arg username "$old_username" '.[$username]' "$USERS_FILE" >/dev/null 2>&1 ) - - if [ $? -ne 0 ]; then - echo -e "${red}Error:${NC} User '$old_username' not found." - return 1 - fi - - existing_user_data=$(jq --arg username "$old_username" '.[$username]' "$USERS_FILE") - upload_bytes=$(echo "$existing_user_data" | jq -r '.upload_bytes // 0') - download_bytes=$(echo "$existing_user_data" | jq -r '.download_bytes // 0') - status=$(echo "$existing_user_data" | jq -r '.status // "Offline"') - - - echo "Updating user:" - echo "Username: $new_username" - echo "Password: ${new_password:-(not changed)}" - echo "Max Download Bytes: ${new_max_download_bytes:-(not changed)}" - echo "Expiration Days: ${new_expiration_days:-(not changed)}" - echo "Creation Date: ${new_account_creation_date:-(not changed)}" - echo "Blocked: $new_blocked" - echo "Unlimited IP: $new_unlimited" - - - jq \ - --arg old_username "$old_username" \ - --arg new_username "$new_username" \ - --arg password "${new_password:-null}" \ - --argjson max_download_bytes "${new_max_download_bytes:-null}" \ - --argjson expiration_days "${new_expiration_days:-null}" \ - --arg account_creation_date "${new_account_creation_date:-null}" \ - --argjson blocked "$new_blocked" \ - --argjson unlimited "$new_unlimited" \ - --argjson upload_bytes "$upload_bytes" \ - --argjson download_bytes "$download_bytes" \ - --arg status "$status" \ - ' - .[$new_username] = .[$old_username] | - (if $old_username != $new_username then del(.[$old_username]) else . end) | - .[$new_username] |= ( - .password = ($password // .password) | - .max_download_bytes = ($max_download_bytes // .max_download_bytes) | - .expiration_days = ($expiration_days // .expiration_days) | - .account_creation_date = (if $account_creation_date == "null" then null else $account_creation_date end) | - .blocked = $blocked | - .unlimited_user = $unlimited | - .status = $status - ) | - (if .[$new_username].status != "On-hold" then - .[$new_username] |= ( - .upload_bytes = $upload_bytes | - .download_bytes = $download_bytes - ) - else . end - ) - ' "$USERS_FILE" > tmp.$$.json && mv tmp.$$.json "$USERS_FILE" - - if [ $? -ne 0 ]; then - echo "Failed to update user '$old_username' in '$USERS_FILE'." - return 1 - fi - - return 0 -} - - -edit_user() { - local username=$1 - local new_username=$2 - local new_traffic_limit=$3 - local new_expiration_days=$4 - local new_password=$5 - local new_creation_date=$6 - local new_blocked=$7 - local new_unlimited=$8 - - - local user_info=$(get_user_info "$username") - - if [ $? -ne 0 ] || [ -z "$user_info" ]; then - echo "User '$username' not found." - return 1 - fi - - local password=$(echo "$user_info" | jq -r '.password') - local traffic_limit=$(echo "$user_info" | jq -r '.max_download_bytes') - local expiration_days=$(echo "$user_info" | jq -r '.expiration_days') - local creation_date=$(echo "$user_info" | jq -r '.account_creation_date') - local blocked=$(echo "$user_info" | jq -r '.blocked') - local unlimited_user=$(echo "$user_info" | jq -r '.unlimited_user // false') - - if ! validate_username "$new_username"; then - echo "Invalid username: $new_username" - return 1 - fi - - if ! validate_traffic_limit "$new_traffic_limit"; then - echo "Invalid traffic limit: $new_traffic_limit" - return 1 - fi - - - if ! validate_expiration_days "$new_expiration_days"; then - echo "Invalid expiration days: $new_expiration_days" - return 1 - fi - - - if ! validate_date "$new_creation_date"; then - echo "Invalid creation date: $new_creation_date" - return 1 - fi - - if ! validate_blocked_status "$new_blocked"; then - echo "Invalid blocked status: $new_blocked" - return 1 - fi - - if ! validate_unlimited_status "$new_unlimited"; then - echo "Invalid unlimited status: $new_unlimited" - return 1 - fi - - - new_username=${new_username:-$username} - new_password=${new_password:-$password} - - - if [ -n "$new_traffic_limit" ]; then - new_traffic_limit=$((new_traffic_limit * GB_TO_BYTES)) - else - new_traffic_limit=$traffic_limit - fi - - new_expiration_days=${new_expiration_days:-$expiration_days} - new_creation_date=${new_creation_date:-$creation_date} - new_blocked=$(convert_boolean_status "${new_blocked:-$blocked}") - new_unlimited=$(convert_boolean_status "${new_unlimited:-$unlimited_user}") - - - if ! update_user_info "$username" "$new_username" "$new_password" "$new_traffic_limit" "$new_expiration_days" "$new_creation_date" "$new_blocked" "$new_unlimited"; then - return 1 - fi - - echo "User updated successfully." - return 0 - -} - - -edit_user "$1" "$2" "$3" "$4" "$5" "$6" "$7" "$8" \ No newline at end of file diff --git a/core/scripts/hysteria2/get_user.py b/core/scripts/hysteria2/get_user.py index 1273675..80968d9 100644 --- a/core/scripts/hysteria2/get_user.py +++ b/core/scripts/hysteria2/get_user.py @@ -4,12 +4,13 @@ import json import sys import os import getopt -from init_paths import * -from paths import * + +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +from db.database import db def get_user_info(username): """ - Retrieves and prints information for a specific user from the USERS_FILE. + Retrieves and prints information for a specific user from the database. Args: username (str): The username to look up. @@ -17,30 +18,20 @@ def get_user_info(username): Returns: int: 0 on success, 1 on failure. """ - if not os.path.isfile(USERS_FILE): - print(f"users.json file not found at {USERS_FILE}!") + if db is None: + print("Error: Database connection failed. Please ensure MongoDB is running.") return 1 - + try: - with open(USERS_FILE, 'r') as f: - users_data = json.load(f) - except json.JSONDecodeError: - print(f"Error: {USERS_FILE} contains invalid JSON.") - return 1 - - if username in users_data: - user_info = users_data[username] - print(json.dumps(user_info, indent=4)) # Print with indentation for readability - # upload_bytes = user_info.get('upload_bytes', "No upload data available") - # download_bytes = user_info.get('download_bytes', "No download data available") - # status = user_info.get('status', "Status unavailable") - # You can choose to print these individually as well, if needed - # print(f"Upload Bytes: {upload_bytes}") - # print(f"Download Bytes: {download_bytes}") - # print(f"Status: {status}") - return 0 - else: - print(f"User '{username}' not found in {USERS_FILE}.") + user_info = db.get_user(username) + if user_info: + print(json.dumps(user_info, indent=4)) + return 0 + else: + print(f"User '{username}' not found in the database.") + return 1 + except Exception as e: + print(f"An error occurred while fetching user data: {e}") return 1 if __name__ == "__main__": @@ -54,7 +45,7 @@ if __name__ == "__main__": for opt, arg in opts: if opt in ("-u", "--username"): - username = arg + username = arg.lower() if not username: print(f"Usage: {sys.argv[0]} -u ") diff --git a/core/scripts/hysteria2/kick.py b/core/scripts/hysteria2/kick.py index ade55f7..9fb9c80 100644 --- a/core/scripts/hysteria2/kick.py +++ b/core/scripts/hysteria2/kick.py @@ -3,16 +3,16 @@ import os import sys import json -import time import fcntl -import shutil import datetime -from concurrent.futures import ThreadPoolExecutor -from init_paths import * -from paths import * -from hysteria2_api import Hysteria2Client - import logging +from concurrent.futures import ThreadPoolExecutor + +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +from db.database import db +from hysteria2_api import Hysteria2Client +from paths import CONFIG_FILE + logging.basicConfig( stream=sys.stdout, level=logging.INFO, @@ -22,8 +22,8 @@ logging.basicConfig( logger = logging.getLogger() LOCKFILE = "/tmp/kick.lock" -BACKUP_FILE = f"{USERS_FILE}.bak" MAX_WORKERS = 8 +API_BASE_URL = 'http://127.0.0.1:25413' def acquire_lock(): try: @@ -34,144 +34,104 @@ def acquire_lock(): logger.warning("Another instance is already running. Exiting.") sys.exit(1) -def kick_users(usernames, secret): +def get_secret(): try: - client = Hysteria2Client( - base_url="http://127.0.0.1:25413", - secret=secret - ) - - client.kick_clients(usernames) - logger.info(f"Successfully kicked {len(usernames)} users: {', '.join(usernames)}") - return True - except Exception as e: - logger.error(f"Error kicking users: {str(e)}") - return False + with open(CONFIG_FILE, 'r') as f: + config = json.load(f) + return config.get('trafficStats', {}).get('secret') + except (json.JSONDecodeError, FileNotFoundError): + return None -def process_user(username, user_data, config_secret, users_data): - blocked = user_data.get('blocked', False) - - if blocked: - logger.info(f"Skipping {username} as they are already 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: - logger.info(f"Skipping {username} due to missing account creation date.") - return None - +def kick_users_api(usernames, secret): 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: - logger.info(f"Setting blocked=True for user {username}") - users_data[username]['blocked'] = True - return username - else: - logger.info(f"Skipping {username} due to invalid or missing data.") - return None - + client = Hysteria2Client(base_url=API_BASE_URL, secret=secret) + client.kick_clients(usernames) + logger.info(f"Successfully sent kick command for users: {', '.join(usernames)}") except Exception as e: - logger.error(f"Error processing user {username}: {str(e)}") + logger.error(f"Error kicking users via API: {e}") + +def process_user(user_doc): + username = user_doc.get('_id') + + if not username or user_doc.get('blocked', False): return None + + account_creation_date = user_doc.get('account_creation_date') + if not account_creation_date: + return None + + should_block = False + + try: + expiration_days = user_doc.get('expiration_days', 0) + if expiration_days > 0: + creation_date = datetime.datetime.strptime(account_creation_date, "%Y-%m-%d") + expiration_date = creation_date + datetime.timedelta(days=expiration_days) + if datetime.datetime.now() >= expiration_date: + should_block = True + logger.info(f"User {username} is expired.") + + if not should_block: + max_download_bytes = user_doc.get('max_download_bytes', 0) + if max_download_bytes > 0: + total_bytes = user_doc.get('download_bytes', 0) + user_doc.get('upload_bytes', 0) + if total_bytes >= max_download_bytes: + should_block = True + logger.info(f"User {username} has exceeded their traffic limit.") + + if should_block: + return username + + except (ValueError, TypeError) as e: + logger.error(f"Error processing user {username} due to invalid data: {e}") return None def main(): lock_file = acquire_lock() - try: - shutil.copy2(USERS_FILE, BACKUP_FILE) - logger.info(f"Created backup of users file at {BACKUP_FILE}") - - try: - with open(CONFIG_FILE, 'r') as f: - config = json.load(f) - secret = config.get('trafficStats', {}).get('secret', '') - if not secret: - logger.error("No secret found in config file") - sys.exit(1) - except Exception as e: - logger.error(f"Failed to load config file: {str(e)}") - shutil.copy2(BACKUP_FILE, USERS_FILE) + if db is None: + logger.error("Database connection failed. Exiting.") + sys.exit(1) + + secret = get_secret() + if not secret: + logger.error(f"Could not find secret in {CONFIG_FILE}. Exiting.") sys.exit(1) - try: - with open(USERS_FILE, 'r') as f: - users_data = json.load(f) - logger.info(f"Loaded data for {len(users_data)} users") - except json.JSONDecodeError: - logger.error("Invalid users.json. Restoring backup.") - shutil.copy2(BACKUP_FILE, USERS_FILE) - sys.exit(1) - except Exception as e: - logger.error(f"Failed to load users file: {str(e)}") - shutil.copy2(BACKUP_FILE, USERS_FILE) - sys.exit(1) + all_users = db.get_all_users() + logger.info(f"Loaded {len(all_users)} users from the database for processing.") - users_to_kick = [] - logger.info(f"Processing {len(users_data)} users in parallel with {MAX_WORKERS} workers") + users_to_block = [] 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() - } - + future_to_user = {executor.submit(process_user, user_doc): user_doc for user_doc in all_users} for future in future_to_user: - username = future.result() - if username: - users_to_kick.append(username) - logger.info(f"User {username} added to kick list") + result = future.result() + if result: + users_to_block.append(result) - if users_to_kick: - logger.info(f"Saving changes to users file for {len(users_to_kick)} blocked users") - for retry in range(3): - try: - with open(USERS_FILE, 'w') as f: - json.dump(users_data, f, indent=2) - break - except Exception as e: - logger.error(f"Failed to save users file (attempt {retry+1}): {str(e)}") - time.sleep(1) - if retry == 2: - raise + if not users_to_block: + logger.info("No users to block or kick.") + return + + logger.info(f"Found {len(users_to_block)} users to block: {', '.join(users_to_block)}") - if users_to_kick: - logger.info(f"Kicking {len(users_to_kick)} users") - batch_size = 50 - for i in range(0, len(users_to_kick), batch_size): - batch = users_to_kick[i:i+batch_size] - logger.info(f"Processing batch of {len(batch)} users") - kick_users(batch, secret) - for username in batch: - logger.info(f"Blocked and kicked user {username}") - else: - logger.info("No users to kick") + for username in users_to_block: + db.update_user(username, {'blocked': True}) + logger.info("Successfully updated user statuses to 'blocked' in the database.") + + batch_size = 50 + for i in range(0, len(users_to_block), batch_size): + batch = users_to_block[i:i + batch_size] + kick_users_api(batch, secret) except Exception as e: - logger.error(f"An error occurred: {str(e)}") - logger.info("Restoring users file from backup") - shutil.copy2(BACKUP_FILE, USERS_FILE) + logger.error(f"An unexpected error occurred in main execution: {e}", exc_info=True) sys.exit(1) finally: fcntl.flock(lock_file, fcntl.LOCK_UN) lock_file.close() - logger.info("Script completed") + logger.info("Script finished.") if __name__ == "__main__": diff --git a/core/scripts/hysteria2/list_users.py b/core/scripts/hysteria2/list_users.py index a1ce151..61e28e5 100644 --- a/core/scripts/hysteria2/list_users.py +++ b/core/scripts/hysteria2/list_users.py @@ -7,8 +7,9 @@ try: except ImportError: sys.exit("Error: hysteria2_api library not found. Please install it.") -sys.path.append(str(Path(__file__).parent.parent)) -from paths import USERS_FILE, CONFIG_FILE, API_BASE_URL +sys.path.append(str(Path(__file__).resolve().parent.parent)) +from db.database import db +from paths import CONFIG_FILE, API_BASE_URL def get_secret() -> str | None: if not CONFIG_FILE.exists(): @@ -20,34 +21,45 @@ def get_secret() -> str | None: except (json.JSONDecodeError, IOError): return None -def get_users() -> dict: - if not USERS_FILE.exists(): - return {} +def get_users_from_db() -> list: + if db is None: + print("Error: Database connection failed.", file=sys.stderr) + return [] try: - with USERS_FILE.open('r') as f: - return json.load(f) - except (json.JSONDecodeError, IOError): - return {} + users = db.get_all_users() + for user in users: + user['username'] = user.pop('_id') + return users + except Exception as e: + print(f"Error retrieving users from database: {e}", file=sys.stderr) + return [] def main(): - users_dict = get_users() + users_list = get_users_from_db() + if not users_list: + print(json.dumps([], indent=2)) + return + secret = get_secret() - if secret and users_dict: + if secret: try: client = Hysteria2Client(base_url=API_BASE_URL, secret=secret) online_clients = client.get_online_clients() - + + users_dict = {user['username']: user for user in users_list} for username, status in online_clients.items(): if status.is_online and username in users_dict: users_dict[username]['online_count'] = status.connections - except Exception: + + users_list = list(users_dict.values()) + + except Exception as e: + print(f"Warning: Could not connect to Hysteria2 API to get online status. {e}", file=sys.stderr) pass - users_list = [ - {**user_data, 'username': username, 'online_count': user_data.get('online_count', 0)} - for username, user_data in users_dict.items() - ] + for user in users_list: + user.setdefault('online_count', 0) print(json.dumps(users_list, indent=2)) diff --git a/core/scripts/hysteria2/list_users.sh b/core/scripts/hysteria2/list_users.sh deleted file mode 100644 index ca76d18..0000000 --- a/core/scripts/hysteria2/list_users.sh +++ /dev/null @@ -1,6 +0,0 @@ -#!/bin/bash - -source /etc/hysteria/core/scripts/path.sh - - -cat "$USERS_FILE" \ No newline at end of file diff --git a/core/scripts/hysteria2/remove_user.py b/core/scripts/hysteria2/remove_user.py index 0663eb4..94195b9 100644 --- a/core/scripts/hysteria2/remove_user.py +++ b/core/scripts/hysteria2/remove_user.py @@ -1,46 +1,34 @@ #!/usr/bin/env python3 -import json import sys import os -import asyncio -from init_paths import * -from paths import * -def sync_remove_user(username): - if not os.path.isfile(USERS_FILE): - return 1, f"Error: Config file {USERS_FILE} not found." +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +from db.database import db + +def remove_user(username): + if db is None: + return 1, "Error: Database connection failed. Please ensure MongoDB is running." try: - with open(USERS_FILE, 'r') as f: - try: - users_data = json.load(f) - except json.JSONDecodeError: - return 1, f"Error: {USERS_FILE} contains invalid JSON." - - if username in users_data: - del users_data[username] - with open(USERS_FILE, 'w') as f: - json.dump(users_data, f, indent=4) + result = db.delete_user(username) + if result.deleted_count > 0: return 0, f"User {username} removed successfully." else: return 1, f"Error: User {username} not found." except Exception as e: - return 1, f"Error: {str(e)}" + return 1, f"An error occurred while removing the user: {e}" -async def remove_user(username): - return await asyncio.to_thread(sync_remove_user, username) - -async def main(): +def main(): if len(sys.argv) != 2: print(f"Usage: {sys.argv[0]} ") sys.exit(1) - username = sys.argv[1] - exit_code, message = await remove_user(username) + username = sys.argv[1].lower() + exit_code, message = remove_user(username) print(message) sys.exit(exit_code) if __name__ == "__main__": - asyncio.run(main()) + main() \ No newline at end of file diff --git a/core/scripts/hysteria2/reset_user.py b/core/scripts/hysteria2/reset_user.py index 98598eb..8274b77 100644 --- a/core/scripts/hysteria2/reset_user.py +++ b/core/scripts/hysteria2/reset_user.py @@ -1,15 +1,15 @@ #!/usr/bin/env python3 -import json import sys import os from datetime import date -from init_paths import * -from paths import * + +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +from db.database import db def reset_user(username): """ - Resets the data usage, status, and creation date of a user in the USERS_FILE. + Resets the data usage, status, and creation date of a user in the database. Args: username (str): The username to reset. @@ -17,35 +17,34 @@ def reset_user(username): Returns: int: 0 on success, 1 on failure. """ - if not os.path.isfile(USERS_FILE): - print(f"Error: File '{USERS_FILE}' not found.") + if db is None: + print("Error: Database connection failed. Please ensure MongoDB is running.") return 1 try: - with open(USERS_FILE, 'r') as f: - users_data = json.load(f) - except json.JSONDecodeError: - print(f"Error: {USERS_FILE} contains invalid JSON.") - return 1 + user = db.get_user(username) + if not user: + print(f"Error: User '{username}' not found in the database.") + return 1 - if username not in users_data: - print(f"Error: User '{username}' not found in '{USERS_FILE}'.") - return 1 + updates = { + 'upload_bytes': 0, + 'download_bytes': 0, + 'status': 'Offline', + 'account_creation_date': date.today().strftime("%Y-%m-%d"), + 'blocked': False + } - today = date.today().strftime("%Y-%m-%d") - users_data[username]['upload_bytes'] = 0 - users_data[username]['download_bytes'] = 0 - users_data[username]['status'] = "Offline" - users_data[username]['account_creation_date'] = today - users_data[username]['blocked'] = False + result = db.update_user(username, updates) + if result.modified_count > 0: + print(f"User '{username}' has been reset successfully.") + return 0 + else: + print(f"User '{username}' data was already in a reset state. No changes made.") + return 0 - try: - with open(USERS_FILE, 'w') as f: - json.dump(users_data, f, indent=4) - print(f"User '{username}' has been reset successfully.") - return 0 - except IOError: - print(f"Error: Failed to reset user '{username}' in '{USERS_FILE}'.") + except Exception as e: + print(f"An error occurred while resetting the user: {e}") return 1 if __name__ == "__main__": @@ -53,6 +52,6 @@ if __name__ == "__main__": print(f"Usage: {sys.argv[0]} ") sys.exit(1) - username_to_reset = sys.argv[1] + username_to_reset = sys.argv[1].lower() exit_code = reset_user(username_to_reset) sys.exit(exit_code) \ No newline at end of file diff --git a/core/scripts/hysteria2/restore.py b/core/scripts/hysteria2/restore.py index a8a85b7..387560d 100644 --- a/core/scripts/hysteria2/restore.py +++ b/core/scripts/hysteria2/restore.py @@ -7,158 +7,122 @@ import shutil import zipfile import tempfile import subprocess -import datetime from pathlib import Path -from init_paths import * -from paths import * -def run_command(command, capture_output=True, check=False): - """Run a shell command and return its output""" - result = subprocess.run( - command, - shell=True, - capture_output=capture_output, - text=True, - check=check - ) - if capture_output: - return result.returncode, result.stdout.strip() - return result.returncode, None +# --- Configuration --- +DB_NAME = "blitz_panel" +HYSTERIA_CONFIG_DIR = Path("/etc/hysteria") +CLI_PATH = Path("/etc/hysteria/core/cli.py") + +def run_command(command, check=False): + """Run a shell command.""" + try: + return subprocess.run( + command, + shell=True, + capture_output=True, + text=True, + check=check + ) + except subprocess.CalledProcessError as e: + print(f"Error executing command: '{e.cmd}'", file=sys.stderr) + print(f"Stderr: {e.stderr}", file=sys.stderr) + raise def main(): if len(sys.argv) < 2: - print("Error: Backup file path is required.") + print("Error: Backup file path is required.", file=sys.stderr) return 1 - backup_zip_file = sys.argv[1] + backup_zip_file = Path(sys.argv[1]) - if not os.path.isfile(backup_zip_file): - print(f"Error: Backup file not found: {backup_zip_file}") + if not backup_zip_file.is_file(): + print(f"Error: Backup file not found: {backup_zip_file}", file=sys.stderr) return 1 - if not backup_zip_file.lower().endswith('.zip'): - print("Error: Backup file must be a .zip file.") + if backup_zip_file.suffix.lower() != '.zip': + print("Error: Backup file must be a .zip file.", file=sys.stderr) return 1 - timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - restore_dir = f"/tmp/hysteria_restore_{timestamp}" - target_dir = "/etc/hysteria" - try: - os.makedirs(restore_dir, exist_ok=True) - - try: - with zipfile.ZipFile(backup_zip_file) as zf: - zf.testzip() - zf.extractall(restore_dir) - except zipfile.BadZipFile: - print("Error: Invalid ZIP file.") - return 1 - except Exception as e: - print(f"Error: Could not extract the ZIP file: {e}") - return 1 - - expected_files = [ - "ca.key", - "ca.crt", - "users.json", - "config.json", - ".configs.env" - ] - - for file in expected_files: - file_path = os.path.join(restore_dir, file) - if not os.path.isfile(file_path): - print(f"Error: Required file '{file}' is missing from the backup.") - return 1 - - existing_backup_dir = f"/opt/hysbackup/restore_pre_backup_{timestamp}" - os.makedirs(existing_backup_dir, exist_ok=True) - - for file in expected_files: - source_file = os.path.join(target_dir, file) - dest_file = os.path.join(existing_backup_dir, file) - - if os.path.isfile(source_file): - try: - shutil.copy2(source_file, dest_file) - except Exception as e: - print(f"Error creating backup file before restore from '{source_file}': {e}") - return 1 - - for file in expected_files: - source_file = os.path.join(restore_dir, file) - dest_file = os.path.join(target_dir, file) - + with tempfile.TemporaryDirectory() as temp_dir_str: + temp_dir = Path(temp_dir_str) + print(f"Extracting backup to temporary directory: {temp_dir}") + try: - shutil.copy2(source_file, dest_file) - except Exception as e: - print(f"Error: replace Configuration Files '{file}': {e}") - shutil.rmtree(existing_backup_dir, ignore_errors=True) + with zipfile.ZipFile(backup_zip_file) as zf: + zf.extractall(temp_dir) + except zipfile.BadZipFile: + print("Error: Invalid or corrupt ZIP file.", file=sys.stderr) + return 1 + + dump_dir = temp_dir / DB_NAME + if not dump_dir.is_dir(): + print("Error: Backup is in an old format or is missing the database dump.", file=sys.stderr) + print("Please use a backup created with the new MongoDB-aware script.", file=sys.stderr) return 1 - - config_file = os.path.join(target_dir, "config.json") - - if os.path.isfile(config_file): - print("Checking and adjusting config.json based on system state...") - ret_code, networkdef = run_command("ip route | grep '^default' | awk '{print $5}'") - networkdef = networkdef.strip() - - if networkdef: - with open(config_file, 'r') as f: - config = json.load(f) - - for outbound in config.get('outbounds', []): - if outbound.get('name') == 'v4' and 'direct' in outbound: - current_v4_device = outbound['direct'].get('bindDevice', '') - - if current_v4_device != networkdef: - print(f"Updating v4 outbound bindDevice from '{current_v4_device}' to '{networkdef}'...") - outbound['direct']['bindDevice'] = networkdef - - with open(config_file, 'w') as f: - json.dump(config, f, indent=2) - - ret_code, _ = run_command("systemctl is-active --quiet wg-quick@wgcf.service", capture_output=False) - - if ret_code != 0: - print("wgcf service is NOT active. Removing warps outbound and any ACL rules...") - - with open(config_file, 'r') as f: - config = json.load(f) - - config['outbounds'] = [outbound for outbound in config.get('outbounds', []) - if outbound.get('name') != 'warps'] - - if 'acl' in config and 'inline' in config['acl']: - config['acl']['inline'] = [rule for rule in config['acl']['inline'] - if not rule.startswith('warps(')] - - with open(config_file, 'w') as f: - json.dump(config, f, indent=2) - - run_command("chown hysteria:hysteria /etc/hysteria/ca.key /etc/hysteria/ca.crt", - capture_output=False) - run_command("chmod 640 /etc/hysteria/ca.key /etc/hysteria/ca.crt", - capture_output=False) - - ret_code, _ = run_command(f"python3 {CLI_PATH} restart-hysteria2", capture_output=False) - - if ret_code != 0: - print("Error: Restart service failed.") - return 1 - - print("Hysteria configuration restored and updated successfully.") - return 0 - - except Exception as e: - print(f"An unexpected error occurred: {e}") + print("Restoring MongoDB database... (This will drop the current user data)") + run_command(f"mongorestore --db={DB_NAME} --drop --dir='{temp_dir}'", check=True) + print("Database restored successfully.") + + files_to_copy = ["config.json", ".configs.env", "ca.key", "ca.crt"] + print("Restoring configuration files...") + for filename in files_to_copy: + src = temp_dir / filename + if src.exists(): + shutil.copy2(src, HYSTERIA_CONFIG_DIR / filename) + print(f" - Restored {filename}") + + adjust_config_file() + print("Setting permissions...") + run_command(f"chown hysteria:hysteria {HYSTERIA_CONFIG_DIR / 'ca.key'} {HYSTERIA_CONFIG_DIR / 'ca.crt'}") + run_command(f"chmod 640 {HYSTERIA_CONFIG_DIR / 'ca.key'} {HYSTERIA_CONFIG_DIR / 'ca.crt'}") + + print("Restarting Hysteria service...") + run_command(f"python3 {CLI_PATH} restart-hysteria2", check=True) + + print("\nRestore completed successfully.") + return 0 + + except subprocess.CalledProcessError: + print("\nRestore failed due to a command execution error.", file=sys.stderr) return 1 - finally: - shutil.rmtree(restore_dir, ignore_errors=True) - if 'existing_backup_dir' in locals(): - shutil.rmtree(existing_backup_dir, ignore_errors=True) + except Exception as e: + print(f"\nAn unexpected error occurred during restore: {e}", file=sys.stderr) + return 1 + +def adjust_config_file(): + """Performs system-specific adjustments to the restored config.json.""" + config_file = HYSTERIA_CONFIG_DIR / "config.json" + if not config_file.exists(): + return + + print("Adjusting config.json based on current system state...") + try: + with open(config_file, 'r') as f: + config = json.load(f) + + result = run_command("ip route | grep '^default' | awk '{print $5}'") + network_device = result.stdout.strip() + if network_device: + for outbound in config.get('outbounds', []): + if outbound.get('name') == 'v4' and 'direct' in outbound: + outbound['direct']['bindDevice'] = network_device + + result = run_command("systemctl is-active --quiet wg-quick@wgcf.service") + if result.returncode != 0: + print(" - WARP service is inactive, removing related configuration.") + config['outbounds'] = [o for o in config.get('outbounds', []) if o.get('name') != 'warps'] + if 'acl' in config and 'inline' in config['acl']: + config['acl']['inline'] = [r for r in config['acl']['inline'] if not r.startswith('warps(')] + + with open(config_file, 'w') as f: + json.dump(config, f, indent=2) + + except Exception as e: + print(f"Warning: Could not adjust config.json. {e}", file=sys.stderr) + if __name__ == "__main__": sys.exit(main()) \ No newline at end of file diff --git a/core/scripts/hysteria2/wrapper_uri.py b/core/scripts/hysteria2/wrapper_uri.py index 57a0d4a..1590ee6 100644 --- a/core/scripts/hysteria2/wrapper_uri.py +++ b/core/scripts/hysteria2/wrapper_uri.py @@ -7,7 +7,8 @@ import argparse from functools import lru_cache from typing import Dict, List, Any -from init_paths import * +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +from db.database import db from paths import * @lru_cache(maxsize=None) @@ -42,10 +43,12 @@ def generate_uri(username: str, auth_password: str, ip: str, port: str, def process_users(target_usernames: List[str]) -> List[Dict[str, Any]]: config = load_json_file(CONFIG_FILE) - all_users = load_json_file(USERS_FILE) - - if not config or not all_users: - print("Error: Could not load Hysteria2 configuration or user files.", file=sys.stderr) + if not config: + print("Error: Could not load Hysteria2 configuration file.", file=sys.stderr) + sys.exit(1) + + if db is None: + print("Error: Database connection failed.", file=sys.stderr) sys.exit(1) nodes = load_json_file(NODES_JSON_PATH) or [] @@ -73,7 +76,7 @@ def process_users(target_usernames: List[str]) -> List[Dict[str, Any]]: results = [] for username in target_usernames: - user_data = all_users.get(username) + user_data = db.get_user(username) if not user_data or "password" not in user_data: results.append({"username": username, "error": "User not found or password not set"}) continue @@ -104,17 +107,20 @@ def process_users(target_usernames: List[str]) -> List[Dict[str, Any]]: def main(): parser = argparse.ArgumentParser(description="Efficiently generate Hysteria2 URIs for multiple users.") parser.add_argument('usernames', nargs='*', help="A list of usernames to process.") - parser.add_argument('--all', action='store_true', help="Process all users from users.json.") + parser.add_argument('--all', action='store_true', help="Process all users from the database.") args = parser.parse_args() target_usernames = args.usernames if args.all: - all_users = load_json_file(USERS_FILE) - if all_users: - target_usernames = list(all_users.keys()) - else: - print("Error: Could not load users.json to process all users.", file=sys.stderr) + if db is None: + print("Error: Database connection failed.", file=sys.stderr) + sys.exit(1) + try: + all_users_docs = db.get_all_users() + target_usernames = [user['_id'] for user in all_users_docs] + except Exception as e: + print(f"Error retrieving all users from database: {e}", file=sys.stderr) sys.exit(1) if not target_usernames: diff --git a/core/traffic.py b/core/traffic.py index f27244a..2c4d0cf 100644 --- a/core/traffic.py +++ b/core/traffic.py @@ -1,23 +1,22 @@ #!/usr/bin/env python3 + import json import os import sys -import time import fcntl -import shutil import datetime -from concurrent.futures import ThreadPoolExecutor from hysteria2_api import Hysteria2Client +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, os.path.join(SCRIPT_DIR, 'scripts')) + +from db.database import db + 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 +LOCKFILE = "/tmp/hysteria_traffic.lock" 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) @@ -25,92 +24,22 @@ def acquire_lock(): except IOError: sys.exit(1) -def traffic_status(no_gui=False): - """Updates and retrieves traffic statistics for all users.""" - green = '\033[0;32m' - cyan = '\033[0;36m' - NC = '\033[0m' - +def get_secret(): try: - with open(CONFIG_FILE, 'r') as config_file: - 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}") + with open(CONFIG_FILE, 'r') as f: + config = json.load(f) + return config.get('trafficStats', {}).get('secret') + except (json.JSONDecodeError, FileNotFoundError): return None - if not secret: - 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: - if not no_gui: - print(f"Error communicating with Hysteria2 API: {e}") - return None - - users_data = {} - if os.path.exists(USERS_FILE): - try: - 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 None - - for user in users_data: - users_data[user]["status"] = "Offline" - - for user_id, status in online_status.items(): - if user_id in users_data: - users_data[user_id]["status"] = "Online" if status.is_online else "Offline" - else: - users_data[user_id] = { - "upload_bytes": 0, "download_bytes": 0, - "status": "Online" if status.is_online else "Offline" - } - - for user_id, stats in traffic_stats.items(): - if user_id in users_data: - users_data[user_id]["upload_bytes"] = users_data[user_id].get("upload_bytes", 0) + stats.upload_bytes - users_data[user_id]["download_bytes"] = users_data[user_id].get("download_bytes", 0) + stats.download_bytes - else: - online = user_id in online_status and online_status[user_id].is_online - users_data[user_id] = { - "upload_bytes": stats.upload_bytes, "download_bytes": stats.download_bytes, - "status": "Online" if online else "Offline" - } - - today_date = datetime.datetime.now().strftime("%Y-%m-%d") - for username, user_data in users_data.items(): - is_on_hold = not user_data.get("account_creation_date") - - if is_on_hold: - is_online = user_data.get("status") == "Online" - has_traffic = user_data.get("download_bytes", 0) > 0 or user_data.get("upload_bytes", 0) > 0 - - if is_online or has_traffic: - user_data["account_creation_date"] = today_date - else: - user_data["status"] = "On-hold" - - with open(USERS_FILE, 'w') as users_file: - json.dump(users_data, users_file, indent=4) - - if not no_gui: - display_traffic_data(users_data, green, cyan, NC) - - return users_data +def format_bytes(bytes_val): + if bytes_val < 1024: return f"{bytes_val}B" + elif bytes_val < 1048576: return f"{bytes_val / 1024:.2f}KB" + elif bytes_val < 1073741824: return f"{bytes_val / 1048576:.2f}MB" + elif bytes_val < 1099511627776: return f"{bytes_val / 1073741824:.2f}GB" + else: return f"{bytes_val / 1099511627776:.2f}TB" 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 @@ -123,113 +52,151 @@ def display_traffic_data(data, green, cyan, NC): for user, entry in data.items(): upload_bytes = entry.get("upload_bytes", 0) download_bytes = entry.get("download_bytes", 0) - status = entry.get("status", "Offline") - + status = entry.get("status", "On-hold") formatted_tx = format_bytes(upload_bytes) formatted_rx = format_bytes(download_bytes) - print(f"{user:<15} {green}{formatted_tx:<15}{NC} {cyan}{formatted_rx:<15}{NC} {status:<10}") print("-------------------------------------------------") -def format_bytes(bytes_val): - """Format bytes as human-readable string""" - if bytes_val < 1024: return f"{bytes_val}B" - elif bytes_val < 1048576: return f"{bytes_val / 1024:.2f}KB" - elif bytes_val < 1073741824: return f"{bytes_val / 1048576:.2f}MB" - elif bytes_val < 1099511627776: return f"{bytes_val / 1073741824:.2f}GB" - else: return f"{bytes_val / 1099511627776:.2f}TB" +def traffic_status(no_gui=False): + green, cyan, NC = '\033[0;32m', '\033[0;36m', '\033[0m' + + if db is None: + if not no_gui: print("Error: Database connection failed.") + return None -def kick_users(usernames, secret): - """Kicks specified users from the server""" + secret = get_secret() + if not secret: + if not no_gui: print(f"Error: Secret not found or failed to read {CONFIG_FILE}.") + 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 None + + try: + all_users = db.get_all_users() + initial_users_data = {user['_id']: user for user in all_users} + except Exception as e: + if not no_gui: print(f"Error fetching users from database: {e}") + return None + + today_date = datetime.datetime.now().strftime("%Y-%m-%d") + users_to_update = {} + + for username, user_data in initial_users_data.items(): + updates = {} + is_online = username in online_status and online_status[username].is_online + + if username in traffic_stats: + new_upload = user_data.get('upload_bytes', 0) + traffic_stats[username].upload_bytes + new_download = user_data.get('download_bytes', 0) + traffic_stats[username].download_bytes + if new_upload != user_data.get('upload_bytes'): updates['upload_bytes'] = new_upload + if new_download != user_data.get('download_bytes'): updates['download_bytes'] = new_download + + is_activated = "account_creation_date" in user_data + + if not is_activated: + current_traffic = traffic_stats.get(username) + has_activity = is_online or (current_traffic and (current_traffic.upload_bytes > 0 or current_traffic.download_bytes > 0)) + + if has_activity: + updates["account_creation_date"] = today_date + updates["status"] = "Online" if is_online else "Offline" + else: + if user_data.get("status") != "On-hold": + updates["status"] = "On-hold" + else: + new_status = "Online" if is_online else "Offline" + if user_data.get("status") != new_status: + updates["status"] = new_status + + if updates: + users_to_update[username] = updates + + if users_to_update: + try: + for username, update_data in users_to_update.items(): + db.update_user(username, update_data) + except Exception as e: + if not no_gui: print(f"Error updating database: {e}") + return None + + if not no_gui: + # For display, merge updates into the initial data + for username, updates in users_to_update.items(): + initial_users_data[username].update(updates) + display_traffic_data(initial_users_data, green, cyan, NC) + + return initial_users_data + +def kick_api_call(usernames, secret): 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, users_data): - """Process a single user to check if they should be kicked""" - if user_data.get('blocked', False): return None - - account_creation_date = user_data.get('account_creation_date') - if not account_creation_date: return None - - max_download_bytes = user_data.get('max_download_bytes', 0) - expiration_days = user_data.get('expiration_days', 0) - total_bytes = user_data.get('download_bytes', 0) + user_data.get('upload_bytes', 0) - - should_block = False - try: - if expiration_days > 0: - creation_date = datetime.datetime.strptime(account_creation_date, "%Y-%m-%d") - expiration_date = creation_date + datetime.timedelta(days=expiration_days) - if datetime.datetime.now() >= expiration_date: - should_block = True - - if not should_block and max_download_bytes > 0 and total_bytes >= max_download_bytes: - should_block = True - - if should_block: - users_data[username]['blocked'] = True - return username - except (ValueError, TypeError): - return None - - return None + except Exception as e: + print(f"Failed to kick users via API: {e}", file=sys.stderr) def kick_expired_users(): - """Kicks users who have exceeded their data limits or whose accounts have expired""" - lock_file = acquire_lock() + if db is None: + print("Error: Database connection failed.", file=sys.stderr) + return + + secret = get_secret() + if not secret: + print(f"Error: Secret not found or failed to read {CONFIG_FILE}.", file=sys.stderr) + return - try: - if not os.path.exists(USERS_FILE): return - shutil.copy2(USERS_FILE, BACKUP_FILE) - + all_users = db.get_all_users() + users_to_kick, users_to_block = [], [] + + for user in all_users: + username = user.get('_id') + if not username or user.get('blocked', False) or not user.get('account_creation_date'): + continue + + total_bytes = user.get('download_bytes', 0) + user.get('upload_bytes', 0) + should_block = False 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: - sys.exit(1) + if user.get('expiration_days', 0) > 0: + creation_date = datetime.datetime.strptime(user['account_creation_date'], "%Y-%m-%d") + if datetime.datetime.now() >= creation_date + datetime.timedelta(days=user['expiration_days']): + should_block = True - with open(USERS_FILE, 'r') as f: - users_data = json.load(f) - - users_to_kick = [] - with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor: - futures = [executor.submit(process_user, u, d, users_data) for u, d in users_data.items()] - for future in futures: - result = future.result() - if result: - users_to_kick.append(result) - - if users_to_kick: - with open(USERS_FILE, 'w') as f: - json.dump(users_data, f, indent=4) - - for i in range(0, len(users_to_kick), 50): - batch = users_to_kick[i:i+50] - kick_users(batch, secret) - - except Exception: - if os.path.exists(BACKUP_FILE): - shutil.copy2(BACKUP_FILE, USERS_FILE) - sys.exit(1) - finally: - fcntl.flock(lock_file, fcntl.LOCK_UN) - lock_file.close() + if not should_block and user.get('max_download_bytes', 0) > 0 and total_bytes >= user['max_download_bytes']: + should_block = True + + if should_block: + users_to_kick.append(username) + users_to_block.append(username) + except (ValueError, TypeError): + continue + + if users_to_block: + for username in users_to_block: + db.update_user(username, {'blocked': True}) + + if users_to_kick: + for i in range(0, len(users_to_kick), 50): + kick_api_call(users_to_kick[i:i+50], secret) if __name__ == "__main__": - 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() + lock_file = acquire_lock() + try: + 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"Usage: python {sys.argv[0]} [kick|--no-gui]") 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 + traffic_status(no_gui=False) + finally: + fcntl.flock(lock_file, fcntl.LOCK_UN) + lock_file.close() \ No newline at end of file