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
|
||||
Reference in New Issue
Block a user