From 8ba46b9878cc592b766a529cd4f0d1559a3cf144 Mon Sep 17 00:00:00 2001 From: Whispering Wind <151555003+ReturnFI@users.noreply.github.com> Date: Fri, 2 May 2025 13:07:53 +0330 Subject: [PATCH] Refactor: Implement user addition functionality in Python --- core/cli_api.py | 11 ++- core/scripts/hysteria2/add_user.py | 117 +++++++++++++++++++++++++++++ core/scripts/hysteria2/add_user.sh | 62 --------------- 3 files changed, 124 insertions(+), 66 deletions(-) create mode 100644 core/scripts/hysteria2/add_user.py delete mode 100644 core/scripts/hysteria2/add_user.sh diff --git a/core/cli_api.py b/core/cli_api.py index ea611db..d86514e 100644 --- a/core/cli_api.py +++ b/core/cli_api.py @@ -23,7 +23,7 @@ class Command(Enum): CHANGE_PORT_HYSTERIA2 = os.path.join(SCRIPT_DIR, 'hysteria2', 'change_port.py') CHANGE_SNI_HYSTERIA2 = os.path.join(SCRIPT_DIR, 'hysteria2', 'change_sni.sh') GET_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'get_user.py') - ADD_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'add_user.sh') + ADD_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'add_user.py') EDIT_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'edit_user.sh') RESET_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'reset_user.py') REMOVE_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'remove_user.py') @@ -110,8 +110,11 @@ def generate_password() -> str: ''' try: return subprocess.check_output(['pwgen', '-s', '32', '1'], shell=False).decode().strip() - except subprocess.CalledProcessError as e: - raise PasswordGenerationError(f'Failed to generate password: {e}') + except (subprocess.CalledProcessError, FileNotFoundError): + try: + return subprocess.check_output(['cat', '/proc/sys/kernel/random/uuid'], shell=False).decode().strip() + except Exception as e: + raise PasswordGenerationError(f"Failed to generate password: {e}") # endregion @@ -255,7 +258,7 @@ def add_user(username: str, traffic_limit: int, expiration_days: int, password: password = generate_password() if not creation_date: creation_date = datetime.now().strftime('%Y-%m-%d') - run_cmd(['bash', Command.ADD_USER.value, username, str(traffic_limit), str(expiration_days), password, creation_date]) + run_cmd(['python3', Command.ADD_USER.value, username, str(traffic_limit), str(expiration_days), password, creation_date]) 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): diff --git a/core/scripts/hysteria2/add_user.py b/core/scripts/hysteria2/add_user.py new file mode 100644 index 0000000..1d5a9b3 --- /dev/null +++ b/core/scripts/hysteria2/add_user.py @@ -0,0 +1,117 @@ +#!/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 * + +def add_user(username, traffic_gb, expiration_days, password=None, creation_date=None): + """ + 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. If None, the current date is used. + + 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]") + return 1 + + try: + traffic_bytes = int(float(traffic_gb) * 1073741824) + expiration_days = int(expiration_days) + except ValueError: + print("Error: Traffic limit and expiration days must be numeric.") + return 1 + + username_lower = username.lower() + + if not password: + try: + password_process = subprocess.run(['pwgen', '-s', '32', '1'], capture_output=True, text=True, check=True) + password = password_process.stdout.strip() + except FileNotFoundError: + try: + password = subprocess.check_output(['cat', '/proc/sys/kernel/random/uuid'], text=True).strip() + except Exception: + print("Error: Failed to generate password. Please install 'pwgen' or ensure /proc access.") + + if not creation_date: + creation_date = datetime.now().strftime("%Y-%m-%d") + else: + 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 + + if not re.match(r"^[a-zA-Z0-9]+$", username): + print("Error: Username can only contain letters and numbers.") + 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}.") + return 1 + + 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 + + for existing_username in users_data: + if existing_username.lower() == username_lower: + print("User already exists.") + return 1 + + users_data[username_lower] = { + "password": password, + "max_download_bytes": traffic_bytes, + "expiration_days": expiration_days, + "account_creation_date": creation_date, + "blocked": False + } + + 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}.") + return 1 + +if __name__ == "__main__": + if len(sys.argv) not in [4, 6]: + print(f"Usage: {sys.argv[0]} [password] [creation_date]") + sys.exit(1) + + username = sys.argv[1] + traffic_gb = sys.argv[2] + expiration_days = sys.argv[3] + password = sys.argv[4] if len(sys.argv) > 4 else None + creation_date = sys.argv[5] if len(sys.argv) > 5 else None + + exit_code = add_user(username, traffic_gb, expiration_days, password, creation_date) + sys.exit(exit_code) \ No newline at end of file diff --git a/core/scripts/hysteria2/add_user.sh b/core/scripts/hysteria2/add_user.sh deleted file mode 100644 index 925484d..0000000 --- a/core/scripts/hysteria2/add_user.sh +++ /dev/null @@ -1,62 +0,0 @@ -#!/bin/bash - -source /etc/hysteria/core/scripts/path.sh - -add_user() { - if [ $# -ne 3 ] && [ $# -ne 5 ]; then - echo "Usage: $0 [password] [creation_date]" - exit 1 - fi - - username=$1 - traffic_gb=$2 - expiration_days=$3 - password=$4 - creation_date=$5 - - username_lower=$(echo "$username" | tr '[:upper:]' '[:lower:]') - - if [ -z "$password" ]; then - password=$(pwgen -s 32 1) - fi - if [ -z "$creation_date" ]; then - creation_date=$(date +%Y-%m-%d) - else - if ! [[ "$creation_date" =~ ^[0-9]{4}-[0-9]{2}-[0-9]{2}$ ]]; then - echo "Invalid date format. Expected YYYY-MM-DD." - exit 1 - fi - - if ! date -d "$creation_date" >/dev/null 2>&1; then - echo "Invalid date. Please provide a valid date in YYYY-MM-DD format." - exit 1 - fi - fi - - if ! [[ "$username" =~ ^[a-zA-Z0-9]+$ ]]; then - echo -e "${red}Error:${NC} Username can only contain letters and numbers." - exit 1 - fi - - traffic=$(echo "$traffic_gb * 1073741824" | bc) - - if [ ! -f "$USERS_FILE" ]; then - echo "{}" > "$USERS_FILE" - fi - - user_exists=$(jq --arg username "$username_lower" ' - to_entries[] | select(.key | ascii_downcase == $username) | .key' "$USERS_FILE") - - if [ -n "$user_exists" ]; then - echo "User already exists." - exit 1 - fi - - jq --arg username "$username_lower" --arg password "$password" --argjson traffic "$traffic" --argjson expiration_days "$expiration_days" --arg creation_date "$creation_date" \ - '.[$username] = {password: $password, max_download_bytes: $traffic, expiration_days: $expiration_days, account_creation_date: $creation_date, blocked: false}' \ - "$USERS_FILE" > "${USERS_FILE}.temp" && mv "${USERS_FILE}.temp" "$USERS_FILE" - - echo -e "User $username added successfully." -} - -add_user "$1" "$2" "$3" "$4" "$5"