Refactor: Implement kick users in Python
This commit is contained in:
Whispering Wind
@ -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 -
|
||||
|
||||
}
|
||||
|
||||
|
||||
178
core/scripts/hysteria2/kick.py
Normal file
178
core/scripts/hysteria2/kick.py
Normal file
@ -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()
|
||||
@ -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
|
||||
43
upgrade.sh
43
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
|
||||
|
||||
Reference in New Issue
Block a user