Refactor: Implement user addition functionality in Python

This commit is contained in:
Whispering Wind
2025-05-02 13:07:53 +03:30
committed by GitHub
parent 5a253258e8
commit 8ba46b9878
3 changed files with 124 additions and 66 deletions

View File

@ -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):

View File

@ -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]} <username> <traffic_limit_GB> <expiration_days> [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]} <username> <traffic_limit_GB> <expiration_days> [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)

View File

@ -1,62 +0,0 @@
#!/bin/bash
source /etc/hysteria/core/scripts/path.sh
add_user() {
if [ $# -ne 3 ] && [ $# -ne 5 ]; then
echo "Usage: $0 <username> <traffic_limit_GB> <expiration_days> [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"