Refactor: Implement kick users in Python

core/scripts/hysteria2/install.sh

@@ -80,12 +80,12 @@ install_hysteria() {
     fi
     
     chmod +x /etc/hysteria/core/scripts/hysteria2/user.sh
-    chmod +x /etc/hysteria/core/scripts/hysteria2/kick.sh
+    chmod +x /etc/hysteria/core/scripts/hysteria2/kick.py
 
     (crontab -l ; echo "*/1 * * * * /bin/bash -c 'source /etc/hysteria/hysteria2_venv/bin/activate && python3 /etc/hysteria/core/cli.py traffic-status' >/dev/null 2>&1") | crontab -
     (crontab -l ; echo "0 3 */3 * * /bin/bash -c 'source /etc/hysteria/hysteria2_venv/bin/activate && python3 /etc/hysteria/core/cli.py restart-hysteria2' >/dev/null 2>&1") | crontab -
     (crontab -l ; echo "0 */6 * * * /bin/bash -c 'source /etc/hysteria/hysteria2_venv/bin/activate && python3 /etc/hysteria/core/cli.py backup-hysteria' >/dev/null 2>&1") | crontab -
-    (crontab -l ; echo "*/1 * * * * /etc/hysteria/core/scripts/hysteria2/kick.sh >/dev/null 2>&1") | crontab -
+    (crontab -l ; echo "*/1 * * * * /bin/bash -c 'source /etc/hysteria/hysteria2_venv/bin/activate && python3 /etc/hysteria/core/scripts/hysteria2/kick.py' >/dev/null 2>&1") | crontab -
 
 }
 

core/scripts/hysteria2/kick.py

@@ -0,0 +1,178 @@
+#!/usr/bin/env python3
+
+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
+logging.basicConfig(
+    stream=sys.stdout,
+    level=logging.INFO,
+    format='%(asctime)s: [%(levelname)s] %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+logger = logging.getLogger()
+
+LOCKFILE = "/tmp/kick.lock"
+BACKUP_FILE = f"{USERS_FILE}.bak"
+MAX_WORKERS = 8
+
+def acquire_lock():
+    try:
+        lock_file = open(LOCKFILE, 'w')
+        fcntl.flock(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        return lock_file
+    except IOError:
+        logger.warning("Another instance is already running. Exiting.")
+        sys.exit(1)
+
+def kick_users(usernames, 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
+
+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
+    
+    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
+            
+    except Exception as e:
+        logger.error(f"Error processing user {username}: {str(e)}")
+        return None
+    
+    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)
+            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)
+            
+        users_to_kick = []
+        logger.info(f"Processing {len(users_data)} users in parallel with {MAX_WORKERS} workers")
+        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()
+            }
+            
+            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")
+        
+        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 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")
+                        
+    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)
+        sys.exit(1)
+    finally:
+        fcntl.flock(lock_file, fcntl.LOCK_UN)
+        lock_file.close()
+        logger.info("Script completed")
+
+
+if __name__ == "__main__":
+    main()

core/scripts/hysteria2/kick.sh

@@ -1,76 +0,0 @@
-#!/bin/bash
-
-source /etc/hysteria/core/scripts/path.sh
-
-LOCKFILE="/tmp/kick.lock"
-exec 200>$LOCKFILE
-flock -n 200 || exit 1
-
-LOGFILE="/var/log/kick.log"
-BACKUP_FILE="${USERS_FILE}.bak"
-
-cp "$USERS_FILE" "$BACKUP_FILE"
-
-kick_user() {
-  local username=$1
-  local secret=$2
-  local kick_endpoint="http://127.0.0.1:25413/kick"
-  curl -s -H "Authorization: $secret" -X POST -d "[\"$username\"]" "$kick_endpoint"
-}
-
-SECRET=$(jq -r '.trafficStats.secret' "$CONFIG_FILE")
-
-if ! jq empty "$USERS_FILE"; then
-  echo "$(date): [ERROR] Invalid users.json. Restoring backup." >> $LOGFILE
-  cp "$BACKUP_FILE" "$USERS_FILE"
-  exit 1
-fi
-
-handle_error() {
-  echo "$(date): [ERROR] An error occurred. Restoring backup." >> $LOGFILE
-  cp "$BACKUP_FILE" "$USERS_FILE"
-  exit 1
-}
-
-trap handle_error ERR
-
-for USERNAME in $(jq -r 'keys[]' "$USERS_FILE"); do
-  BLOCKED=$(jq -r --arg user "$USERNAME" '.[$user].blocked // false' "$USERS_FILE")
-
-  if [ "$BLOCKED" == "true" ]; then
-    echo "$(date): [INFO] Skipping $USERNAME as they are already blocked." >> $LOGFILE
-    continue
-  fi
-
-  MAX_DOWNLOAD_BYTES=$(jq -r --arg user "$USERNAME" '.[$user].max_download_bytes // 0' "$USERS_FILE")
-  EXPIRATION_DAYS=$(jq -r --arg user "$USERNAME" '.[$user].expiration_days // 0' "$USERS_FILE")
-  ACCOUNT_CREATION_DATE=$(jq -r --arg user "$USERNAME" '.[$user].account_creation_date' "$USERS_FILE")
-  CURRENT_DOWNLOAD_BYTES=$(jq -r --arg user "$USERNAME" '.[$user].download_bytes // 0' "$USERS_FILE")
-  CURRENT_UPLOAD_BYTES=$(jq -r --arg user "$USERNAME" '.[$user].upload_bytes // 0' "$USERS_FILE")
-
-  TOTAL_BYTES=$((CURRENT_DOWNLOAD_BYTES + CURRENT_UPLOAD_BYTES))
-
-  if [ -z "$ACCOUNT_CREATION_DATE" ]; then
-    echo "$(date): [INFO] Skipping $USERNAME due to missing account creation date." >> $LOGFILE
-    continue
-  fi
-
-  CURRENT_DATE=$(date +%s)
-  EXPIRATION_DATE=$(date -d "$ACCOUNT_CREATION_DATE + $EXPIRATION_DAYS days" +%s)
-
-  if [ "$MAX_DOWNLOAD_BYTES" -gt 0 ] && [ "$TOTAL_BYTES" -ge 0 ] && [ "$EXPIRATION_DAYS" -gt 0 ]; then
-    if [ "$TOTAL_BYTES" -ge "$MAX_DOWNLOAD_BYTES" ] || [ "$CURRENT_DATE" -ge "$EXPIRATION_DATE" ]; then
-      for i in {1..3}; do
-        jq --arg user "$USERNAME" '.[$user].blocked = true' "$USERS_FILE" > temp.json && mv temp.json "$USERS_FILE" && break
-        sleep 1
-      done
-      kick_user "$USERNAME" "$SECRET"
-      echo "$(date): [INFO] Blocked and kicked user $USERNAME." >> $LOGFILE
-    fi
-  else
-    echo "$(date): [INFO] Skipping $USERNAME due to invalid or missing data." >> $LOGFILE
-  fi
-done
-
-# echo "$(date): [INFO] Kick script completed successfully." >> $LOGFILE
-# exit 0

upgrade.sh

@@ -26,28 +26,6 @@ for FILE in "${FILES[@]}"; do
     cp "$FILE" "$TEMP_DIR/$FILE"
 done
 
-# echo "Checking and renaming old systemd service files"
-# declare -A SERVICE_MAP=(
-#     ["/etc/systemd/system/hysteria-bot.service"]="hysteria-telegram-bot.service"
-#     ["/etc/systemd/system/singbox.service"]="hysteria-singbox.service"
-#     ["/etc/systemd/system/normalsub.service"]="hysteria-normal-sub.service"
-# )
-
-# for OLD_SERVICE in "${!SERVICE_MAP[@]}"; do
-#     NEW_SERVICE="/etc/systemd/system/${SERVICE_MAP[$OLD_SERVICE]}"
-
-#     if [[ -f "$OLD_SERVICE" ]]; then
-#         echo "Stopping old service: $(basename "$OLD_SERVICE")"
-#         systemctl stop "$(basename "$OLD_SERVICE")" 2>/dev/null
-
-#         echo "Renaming $OLD_SERVICE to $NEW_SERVICE"
-#         mv "$OLD_SERVICE" "$NEW_SERVICE"
-
-#         echo "Reloading systemd daemon"
-#         systemctl daemon-reload
-#     fi
-# done
-
 echo "Removing /etc/hysteria directory"
 rm -rf /etc/hysteria/
 
@@ -63,24 +41,6 @@ for FILE in "${FILES[@]}"; do
     cp "$TEMP_DIR/$FILE" "$FILE"
 done
 
-# CADDYFILE="/etc/hysteria/core/scripts/webpanel/Caddyfile"
-
-# if [ -f "$CADDYFILE" ]; then
-#     echo "Updating Caddyfile port from 8080 to 28260"
-
-#     sed -i 's/\(:[[:space:]]*\)8080/\128260/g' "$CADDYFILE"
-#     sed -i 's/0\.0\.0\.0:8080/0.0.0.0:28260/g' "$CADDYFILE"
-#     sed -i 's/127\.0\.0\.1:8080/127.0.0.1:28260/g' "$CADDYFILE"
-
-
-#     if ! grep -q ':28260' "$CADDYFILE"; then
-#         echo "Warning: Caddyfile does not contain port 8080 in expected formats.  Port replacement may have already been done."
-#     fi
-# else
-#     echo "Error: Caddyfile not found at $CADDYFILE.  Cannot update port."
-# fi
-
-
 CONFIG_ENV="/etc/hysteria/.configs.env"
 if [ ! -f "$CONFIG_ENV" ]; then
     echo ".configs.env not found, creating it with default values."
@@ -155,7 +115,6 @@ systemctl restart hysteria-caddy.service
 echo "Restarting other hysteria services"
 systemctl restart hysteria-server.service
 systemctl restart hysteria-telegram-bot.service
-# systemctl restart hysteria-singbox.service
 systemctl restart hysteria-normal-sub.service
 systemctl restart hysteria-webpanel.service
 
@@ -169,6 +128,8 @@ fi
 
 echo "Restoring cron jobs"
 crontab /tmp/crontab_backup
+echo "Updating kick.sh cron job to kick.py"
+( crontab -l | sed "s|/etc/hysteria/core/scripts/hysteria2/kick.sh|/bin/bash -c 'source /etc/hysteria/hysteria2_venv/bin/activate && python3 /etc/hysteria/core/scripts/hysteria2/kick.py'|g" ) | crontab -
 rm /tmp/crontab_backup
 
 chmod +x menu.sh