Merge pull request #5 from ReturnFI/Dev

Fix Reset User

core/cli.py

@@ -25,6 +25,7 @@ class Command(Enum):
     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')
+    RESET_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'reset_user.sh')
     REMOVE_USER = os.path.join(SCRIPT_DIR, 'hysteria2', 'remove_user.sh')
     SHOW_USER_URI = os.path.join(SCRIPT_DIR, 'hysteria2', 'show_user_uri.sh')
     TRAFFIC_STATUS = 'traffic.py'  # won't be call directly (it's a python module)
@@ -177,13 +178,16 @@ def edit_user(username: str, new_username: str, new_traffic_limit: int, new_expi
 
     run_cmd(command_args)
 
+@ cli.command('reset-user')
+@ click.option('--username', '-u', required=True, help='Username for the user to Reset', type=str)
+def reset_user(username: str):
+    run_cmd(['bash', Command.RESET_USER.value, username])
 
 @ 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)
 @click.option('--qrcode', '-qr', is_flag=True, help='Generate QR code for the URI')

core/scripts/hysteria2/edit_user.sh

@@ -124,13 +124,38 @@ update_user_info() {
         .blocked = $blocked
     )' "$USERS_FILE" > tmp.$$.json && mv tmp.$$.json "$USERS_FILE"
 
-
     if [ $? -ne 0 ]; then
-        echo "Error: Failed to update user '$old_username'."
+        echo "Error: Failed to update user '$old_username' in '$USERS_FILE'."
         return 1
     fi
 
-    echo "User '$old_username' updated successfully."
+    # Only update traffic_data.json and restart service if the username has changed
+    if [ "$old_username" != "$new_username" ]; then
+        # Update username in traffic_data.json
+        if [ -f "$TRAFFIC_FILE" ]; then
+            jq --arg old_username "$old_username" \
+               --arg new_username "$new_username" \
+               '
+               .[$new_username] = .[$old_username] |
+               del(.[$old_username])
+               ' "$TRAFFIC_FILE" > tmp.$$.json && mv tmp.$$.json "$TRAFFIC_FILE"
+
+            if [ $? -ne 0 ]; then
+                echo "Error: Failed to update username in '$TRAFFIC_FILE'."
+                return 1
+            fi
+        else
+            echo "Warning: '$TRAFFIC_FILE' not found. Skipping traffic data update."
+        fi
+
+        # Restart Hysteria service after updating user information
+        python3 $CLI_PATH restart-hysteria2
+
+        if [ $? -ne 0 ]; then
+            echo "Error: Failed to restart Hysteria service."
+            return 1
+        fi
+    fi
 }
 
 # Main function to edit user
@@ -177,8 +202,18 @@ edit_user() {
 
     # 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"
-}
 
+    # Restart Hysteria service after updating user information
+    echo "Restarting Hysteria service..."
+    python3 $CLI_PATH restart-hysteria2
+
+    if [ $? -ne 0 ]; then
+        echo "Error: Failed to restart Hysteria service."
+        exit 1
+    fi
+
+    echo "Hysteria service restarted successfully."
+}
 
 # Run the script
 edit_user "$1" "$2" "$3" "$4" "$5" "$6" "$7"

core/scripts/hysteria2/reset_user.sh

@@ -0,0 +1,56 @@
+#!/bin/bash
+
+source /etc/hysteria/core/scripts/utils.sh
+source /etc/hysteria/core/scripts/path.sh
+
+reset_user() {
+    local username=$1
+
+    if [ ! -f "$USERS_FILE" ]; then
+        echo "Error: File '$USERS_FILE' not found."
+        return 1
+    fi
+
+    user_exists=$(jq -e --arg username "$username" '.[$username]' "$USERS_FILE")
+    if [ $? -ne 0 ]; then
+        echo "Error: User '$username' not found in '$USERS_FILE'."
+        return 1
+    fi
+
+    today=$(date +%Y-%m-%d)
+    jq --arg username "$username" \
+       --arg today "$today" \
+       '
+       .[$username].account_creation_date = $today |
+       .[$username].blocked = false
+       ' "$USERS_FILE" > tmp.$$.json && mv tmp.$$.json "$USERS_FILE"
+
+    if [ $? -ne 0 ]; then
+        echo "Error: Failed to reset user '$username' in '$USERS_FILE'."
+        return 1
+    fi
+
+    if [ ! -f "$TRAFFIC_FILE" ]; then
+        echo "Warning: File '$TRAFFIC_FILE' not found. Skipping traffic data reset."
+    else
+        jq --arg username "$username" \
+           '
+           .[$username].upload_bytes = 0 |
+           .[$username].download_bytes = 0
+           ' "$TRAFFIC_FILE" > tmp.$$.json && mv tmp.$$.json "$TRAFFIC_FILE"
+
+        if [ $? -ne 0 ]; then
+            echo "Error: Failed to reset traffic data for user '$username' in '$TRAFFIC_FILE'."
+            return 1
+        fi
+    fi
+
+    echo "User '$username' has been reset successfully."
+}
+
+if [ $# -eq 0 ]; then
+    echo "Usage: $0 <username>"
+    exit 1
+fi
+
+reset_user "$1"

core/scripts/telegrambot/tbot.py

@@ -122,6 +122,7 @@ def process_show_user(message):
                    types.InlineKeyboardButton("Renew Password", callback_data=f"renew_password:{username}"))
         markup.add(types.InlineKeyboardButton("Renew Creation Date", callback_data=f"renew_creation:{username}"),
                    types.InlineKeyboardButton("Block User", callback_data=f"block_user:{username}"))
+        markup.add(types.InlineKeyboardButton("Reset User", callback_data=f"reset_user:{username}"))
 
         bot.send_photo(message.chat.id, bio_v4, caption=f"User Details:\n{formatted_details}\n\nIPv4 URI: {uri_v4}", reply_markup=markup)
 
@@ -131,7 +132,7 @@ def server_info(message):
     result = run_cli_command(command)
     bot.reply_to(message, result)
 
-@bot.callback_query_handler(func=lambda call: call.data.startswith('edit_') or call.data.startswith('renew_') or call.data.startswith('block_'))
+@bot.callback_query_handler(func=lambda call: call.data.startswith('edit_') or call.data.startswith('renew_') or call.data.startswith('block_') or call.data.startswith('reset_'))
 def handle_edit_callback(call):
     action, username = call.data.split(':')
     if action == 'edit_username':
@@ -156,6 +157,10 @@ def handle_edit_callback(call):
         markup.add(types.InlineKeyboardButton("True", callback_data=f"confirm_block:{username}:true"),
                    types.InlineKeyboardButton("False", callback_data=f"confirm_block:{username}:false"))
         bot.send_message(call.message.chat.id, f"Set block status for {username}:", reply_markup=markup)
+    elif action == 'reset_user':
+        command = f"python3 {CLI_PATH} reset-user -u {username}"
+        result = run_cli_command(command)
+        bot.send_message(call.message.chat.id, result)
 
 @bot.callback_query_handler(func=lambda call: call.data.startswith('confirm_block:'))
 def handle_block_confirmation(call):