Add edit user (need fucking test)

core/cli.py

@@ -14,6 +14,7 @@ import validator
 SCRIPT_DIR = '/etc/hysteria/core/scripts'
 DEBUG = True
 
+
 class Command(Enum):
     '''Constais path to command's script'''
     INSTALL_HYSTERIA2 = os.path.join(SCRIPT_DIR, 'hysteria2', 'install.sh')
@@ -21,6 +22,7 @@ class Command(Enum):
     UPDATE_HYSTERIA2 = os.path.join(SCRIPT_DIR, 'hysteria2', 'update.sh')
     RESTART_HYSTERIA2 = os.path.join(SCRIPT_DIR, 'hysteria2', 'restart.sh')
     CHANGE_PORT_HYSTERIA2 = os.path.join(SCRIPT_DIR, 'hysteria2', 'change_port.sh')
+    GET_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'get_user.sh')
     ADD_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'add_user.sh')
     EDIT_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'edit_user.sh')
     REMOVE_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'remove_user.sh')
@@ -39,10 +41,13 @@ def run_cmd(command:list[str]):
     Runs a command and returns the output.
     Could raise subprocess.CalledProcessError
     '''
+    if DEBUG:
+        print(' '.join(command))
     result = subprocess.check_output(command, shell=False)
     if DEBUG:
         print(result.decode().strip())
 
+
 def generate_password() -> str:
     '''
     Generates a random password using pwgen for user.
@@ -58,6 +63,8 @@ def cli():
     pass
 
 # region hysteria2 menu options
+
+
 @cli.command('install-hysteria2')
 @click.option('--port', '-p', required=True, help='New port for Hysteria2', type=int, callback=validator.validate_port)
 def install_hysteria2(port: int):
@@ -68,10 +75,12 @@ def install_hysteria2(port:int):
 def uninstall_hysteria2():
     run_cmd(['bash', Command.UNINSTALL_HYSTERIA2.value])
 
+
 @cli.command('update-hysteria2')
 def update_hysteria2():
     run_cmd(['bash', Command.UPDATE_HYSTERIA2.value])
 
+
 @cli.command('restart-hysteria2')
 def restart_hysteria2():
     run_cmd(['bash', Command.RESTART_HYSTERIA2.value])
@@ -82,6 +91,13 @@ def restart_hysteria2():
 def change_hysteria2_port(port: int):
     run_cmd(['bash', Command.CHANGE_PORT_HYSTERIA2.value, str(port)])
 
+
+@cli.command('get-user')
+@click.option('--username', '-u', required=True, help='Username for the user to get', type=str)
+def get_user(username: str):
+    run_cmd(['bash', Command.GET_USER.value, username])
+
+
 @cli.command('add-user')
 @click.option('--username', '-u', required=True, help='Username for the new user', type=str)
 @click.option('--traffic-limit', '-t', required=True, help='Traffic limit for the new user in GB', type=float)
@@ -100,27 +116,80 @@ def add_user(username:str, traffic_limit:float, expiration_days:int,password:str
 
     run_cmd(['bash', Command.ADD_USER.value, username, str(traffic_limit), str(expiration_days), password, creation_date])
 
-@cli.command('edit-user')
+
+@click.command('edit-user')
 @click.option('--username', '-u', required=True, help='Username for the user to edit', type=str)
-@click.option('--traffic-limit','-t', required=True, help='Traffic limit for the new user in GB',type=float)
+@click.option('--new-username', '-nu', required=False, help='New username for the user', type=str)
-@click.option('--expiration-days','-e', required=True, help='Expiration days for the new user',type=int)
+@click.option('--new-traffic-limit', '-nt', required=False, help='Traffic limit for the new user in GB', type=float)
-def edit_user(username:str, traffic_limit:float, expiration_days:int):
+@click.option('--new-expiration-days', '-ne', required=False, help='Expiration days for the new user', type=int)
-    run_cmd(['bash', Command.EDIT_USER.value, username, str(traffic_limit), str(expiration_days)])
+@click.option('--renew-password', '-rp', is_flag=True, help='Renew password for the user')
+@click.option('--renew-creation-date', '-rc', is_flag=True, help='Renew creation date for the user')
+@click.option('--blocked', '-b', is_flag=True, help='Block the user')
+def edit_user(username: str, new_username: str, new_traffic_limit: float, new_expiration_days: int, renew_password: bool, renew_creation_date: bool, blocked: bool):
+    if not username:
+        print('Error: username is required')
+        exit(1)
+
+    if not any([new_username, new_traffic_limit, new_expiration_days, renew_password, renew_creation_date, blocked is not None]):
+        print('Error: at least one option is required')
+        exit(1)
+
+    if new_traffic_limit is not None and new_traffic_limit <= 0:
+        print('Error: traffic limit must be greater than 0')
+        exit(1)
+
+    if new_expiration_days is not None and new_expiration_days <= 0:
+        print('Error: expiration days must be greater than 0')
+        exit(1)
+
+    # Handle renewing password and creation date
+    if renew_password:
+        try:
+            password = generate_password()
+        except subprocess.CalledProcessError as e:
+            print(f'Error: failed to generate password\n{e}')
+            exit(1)
+    else:
+        password = ""
+
+    if renew_creation_date:
+        creation_date = datetime.now().strftime('%Y-%m-%d')
+    else:
+        creation_date = ""
+
+    # Prepare arguments for the command
+    command_args = [
+        'bash',
+        'edit_user.sh',  # Replace with the actual path to your script
+        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,
+        str(blocked).lower() if blocked is not None else 'false'
+    ]
+
+    run_cmd(command_args)
+
 
 @ cli.command('remove-user')
 @ click.option('--username', '-u', required=True, help='Username for the user to remove', type=str)
 def remove_user(username: str):
     run_cmd(['bash', Command.REMOVE_USER.value, username])
 
+
 @ cli.command('show-user-uri')
 @ click.option('--username', '-u', required=True, help='Username for the user to show the URI', type=str)
 def show_user_uri(username: str):
     run_cmd(['bash', Command.SHOW_USER_URI.value, username])
 
+
 @ cli.command('traffic-status')
 def traffic_status():
     traffic.traffic_status()
 
+
 @ cli.command('list-users')
 def list_users():
     run_cmd(['bash', Command.LIST_USERS.value])
@@ -129,18 +198,22 @@ def list_users():
 
 # region advanced menu
 
+
 @ cli.command('install-tcp-brutal')
 def install_tcp_brutal():
     run_cmd(['bash', Command.INSTALL_TCP_BRUTAL.value])
 
+
 @ cli.command('install-warp')
 def install_warp():
     run_cmd(['bash', Command.INSTALL_WARP.value])
 
+
 @ cli.command('uninstall-warp')
 def uninstall_warp():
     run_cmd(['bash', Command.UNINSTALL_WARP.value])
 
+
 @ cli.command('configure-warp')
 @ click.option('--warp-mode', '-m', required=True, help='Warp mode', type=click.Choice(['proxy', 'direct', 'reject']))
 @ click.option('--block-porn', '-p', required=False, help='Block porn', type=bool)
@@ -149,5 +222,6 @@ def configure_warp(warp_mode:str, block_porn:bool):
 
 # endregion
 
+
 if __name__ == '__main__':
     cli()

core/scripts/hysteria2/edit_user.sh

@@ -0,0 +1,154 @@
+#!/bin/bash
+
+source /etc/hysteria/core/scripts/utils.sh
+source /etc/hysteria/core/scripts/path.sh
+
+# Function to validate all user input fields
+validate_inputs() {
+    local new_username=$1
+    local new_password=$2
+    local new_traffic_limit=$3
+    local new_expiration_days=$4
+    local new_creation_date=$5
+    local new_blocked=$6
+
+    # Validate username
+    if [ -n "$new_username" ]; then
+        if ! [[ "$new_username" =~ ^[a-z0-9]+$ ]]; then
+            echo -e "${red}Error:${NC} Username can only contain lowercase letters and numbers."
+            exit 1
+        fi
+    fi
+
+    # Validate traffic limit
+    if [ -n "$new_traffic_limit" ]; then
+        if ! [[ "$new_traffic_limit" =~ ^[0-9]+$ ]]; then
+            echo -e "${red}Error:${NC} Traffic limit must be a valid integer."
+            exit 1
+        fi
+    fi
+
+    # Validate expiration days
+    if [ -n "$new_expiration_days" ]; then
+        if ! [[ "$new_expiration_days" =~ ^[0-9]+$ ]]; then
+            echo -e "${red}Error:${NC} Expiration days must be a valid integer."
+            exit 1
+        fi
+    fi
+
+    # Validate date format
+    if [ -n "$new_creation_date" ]; then
+        if ! [[ "$new_creation_date" =~ ^[0-9]{4}-[0-9]{2}-[0-9]{2}$ ]]; then
+            echo "Invalid date format. Expected YYYY-MM-DD."
+            exit 1
+        elif ! date -d "$new_creation_date" >/dev/null 2>&1; then
+            echo "Invalid date. Please provide a valid date in YYYY-MM-DD format."
+            exit 1
+        fi
+    fi
+
+    # Validate blocked status
+    if [ -n "$new_blocked" ]; then
+        if [ "$new_blocked" != "true" ] && [ "$new_blocked" != "false" ]; then
+            echo -e "${red}Error:${NC} Blocked status must be 'true' or 'false'."
+            exit 1
+        fi
+    fi
+}
+
+# Function to get user info
+get_user_info() {
+    local username=$1
+    python3 /etc/hysteria/core/scripts/get_user.py "$username"
+}
+
+# Function to update user info in JSON
+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
+
+    if [ ! -f "$USERS_FILE" ]; then
+        echo "Error: File '$USERS_FILE' not found."
+        return 1
+    fi
+
+    # Check if the old username exists
+    user_exists=$(jq -e --arg username "$old_username" '.[$username]' "$USERS_FILE")
+    if [ $? -ne 0 ]; then
+        echo "Error: User '$old_username' not found."
+        return 1
+    fi
+
+    # Prepare jq filter to update the fields
+    jq_filter='.[$old_username] = 
+    if $new_password != "" then .password = $new_password else . end |
+    if $new_max_download_bytes != null then .max_download_bytes = $new_max_download_bytes else . end |
+    if $new_expiration_days != null then .expiration_days = $new_expiration_days else . end |
+    if $new_account_creation_date != "" then .account_creation_date = $new_account_creation_date else . end |
+    if $new_blocked != null then .blocked = $new_blocked else . end'
+
+    # Rename the user if new_username is provided
+    if [ -n "$new_username" ]; then
+        jq_filter=$(echo "$jq_filter" | sed "s|.$old_username|.$new_username|")
+    fi
+
+    jq --arg old_username "$old_username" \
+       --arg new_username "$new_username" \
+       --arg new_password "$new_password" \
+       --argjson new_max_download_bytes "$new_max_download_bytes" \
+       --argjson new_expiration_days "$new_expiration_days" \
+       --arg new_account_creation_date "$new_account_creation_date" \
+       --argjson new_blocked "$new_blocked" \
+       "$jq_filter" \
+       "$USERS_FILE" > tmp.$$.json && mv tmp.$$.json "$USERS_FILE"
+
+    echo "User '$old_username' updated successfully."
+}
+
+# Main function to edit user
+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
+
+    # Get user info
+    user_info=$(get_user_info "$username")
+    if [ -z "$user_info" ]; then
+        echo -e "${red}Error:${NC} User '$username' not found."
+        exit 1
+    fi
+
+    # Extract user info
+    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')
+
+    # Validate all inputs
+    validate_inputs "$new_username" "$new_password" "$new_traffic_limit" "$new_expiration_days" "$new_creation_date" "$new_blocked"
+
+    # Set new values with validation
+    new_username=${new_username:-$username}
+    new_password=${new_password:-$password}
+    new_traffic_limit=${new_traffic_limit:-$traffic_limit}
+    new_traffic_limit=$(echo "$new_traffic_limit * 1073741824" | bc)
+    new_expiration_days=${new_expiration_days:-$expiration_days}
+    new_creation_date=${new_creation_date:-$creation_date}
+    new_blocked=${new_blocked:-$blocked}
+
+    # Update user info in JSON file
+    update_user_info "$username" "$new_username" "$new_password" "$new_traffic_limit" "$new_expiration_days" "$new_creation_date" "$new_blocked"
+}
+
+# Run the script
+edit_user "$1" "$2" "$3" "$4" "$5" "$6" "$7"