refactor(core): centralize on-hold user logic in traffic.py
This commit is contained in:
@ -35,7 +35,7 @@ def add_bulk_users(traffic_gb, expiration_days, count, prefix, start_number, unl
|
|||||||
|
|
||||||
existing_users_lower = {u.lower() for u in users_data}
|
existing_users_lower = {u.lower() for u in users_data}
|
||||||
new_users_to_add = {}
|
new_users_to_add = {}
|
||||||
creation_date = datetime.now().strftime("%Y-%m-%d")
|
creation_date = None
|
||||||
|
|
||||||
try:
|
try:
|
||||||
password_process = subprocess.run(['pwgen', '-s', '32', str(count)], capture_output=True, text=True, check=True)
|
password_process = subprocess.run(['pwgen', '-s', '32', str(count)], capture_output=True, text=True, check=True)
|
||||||
|
|||||||
@ -1,13 +1,9 @@
|
|||||||
#!/usr/bin/env python3
|
#!/usr/bin/env python3
|
||||||
import os
|
|
||||||
import sys
|
|
||||||
import time
|
import time
|
||||||
import schedule
|
import schedule
|
||||||
import logging
|
import logging
|
||||||
import subprocess
|
import subprocess
|
||||||
import fcntl
|
import fcntl
|
||||||
import datetime
|
|
||||||
import json
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from paths import *
|
from paths import *
|
||||||
|
|
||||||
@ -72,46 +68,7 @@ def check_traffic_status():
|
|||||||
try:
|
try:
|
||||||
success = run_command(f"python3 {CLI_PATH} traffic-status --no-gui", log_success=False)
|
success = run_command(f"python3 {CLI_PATH} traffic-status --no-gui", log_success=False)
|
||||||
if not success:
|
if not success:
|
||||||
logger.error("Failed to run traffic-status command. Aborting check.")
|
pass
|
||||||
return
|
|
||||||
|
|
||||||
if not os.path.exists(USERS_FILE):
|
|
||||||
logger.warning(f"{USERS_FILE} not found. Skipping on-hold user check.")
|
|
||||||
return
|
|
||||||
|
|
||||||
try:
|
|
||||||
with open(USERS_FILE, 'r') as f:
|
|
||||||
users_data = json.load(f)
|
|
||||||
except (json.JSONDecodeError, IOError) as e:
|
|
||||||
logger.error(f"Error reading or parsing {USERS_FILE}: {e}")
|
|
||||||
return
|
|
||||||
|
|
||||||
users_updated = False
|
|
||||||
today_date = datetime.datetime.now().strftime("%Y-%m-%d")
|
|
||||||
|
|
||||||
for username, user_data in users_data.items():
|
|
||||||
is_on_hold = not user_data.get("account_creation_date")
|
|
||||||
|
|
||||||
if is_on_hold:
|
|
||||||
is_online = user_data.get("status") == "Online"
|
|
||||||
|
|
||||||
if is_online:
|
|
||||||
logger.info(f"On-hold user '{username}' connected. Activating account with creation date {today_date}.")
|
|
||||||
user_data["account_creation_date"] = today_date
|
|
||||||
users_updated = True
|
|
||||||
else:
|
|
||||||
if user_data.get("status") != "On-hold":
|
|
||||||
user_data["status"] = "On-hold"
|
|
||||||
users_updated = True
|
|
||||||
|
|
||||||
if users_updated:
|
|
||||||
try:
|
|
||||||
with open(USERS_FILE, 'w') as f:
|
|
||||||
json.dump(users_data, f, indent=4)
|
|
||||||
logger.info("Successfully updated users.json for on-hold users.")
|
|
||||||
except IOError as e:
|
|
||||||
logger.error(f"Error writing updates to {USERS_FILE}: {e}")
|
|
||||||
|
|
||||||
finally:
|
finally:
|
||||||
release_lock(lock_fd)
|
release_lock(lock_fd)
|
||||||
|
|
||||||
@ -132,10 +89,8 @@ def main():
|
|||||||
schedule.every(1).minutes.do(check_traffic_status)
|
schedule.every(1).minutes.do(check_traffic_status)
|
||||||
schedule.every(6).hours.do(backup_hysteria)
|
schedule.every(6).hours.do(backup_hysteria)
|
||||||
|
|
||||||
# logger.info("Performing initial runs on startup...")
|
|
||||||
check_traffic_status()
|
check_traffic_status()
|
||||||
backup_hysteria()
|
backup_hysteria()
|
||||||
# logger.info("Initial runs complete. Entering main loop.")
|
|
||||||
|
|
||||||
while True:
|
while True:
|
||||||
try:
|
try:
|
||||||
|
|||||||
@ -149,6 +149,7 @@ async def remove_user_api(username: str):
|
|||||||
cli_api.kick_user_by_name(username)
|
cli_api.kick_user_by_name(username)
|
||||||
cli_api.traffic_status(display_output=False)
|
cli_api.traffic_status(display_output=False)
|
||||||
cli_api.remove_user(username)
|
cli_api.remove_user(username)
|
||||||
|
cli_api.traffic_status(display_output=False)
|
||||||
return DetailResponse(detail=f'User {username} has been removed.')
|
return DetailResponse(detail=f'User {username} has been removed.')
|
||||||
except HTTPException:
|
except HTTPException:
|
||||||
|
|
||||||
|
|||||||
151
core/traffic.py
151
core/traffic.py
@ -16,16 +16,6 @@ LOCKFILE = "/tmp/kick.lock"
|
|||||||
BACKUP_FILE = f"{USERS_FILE}.bak"
|
BACKUP_FILE = f"{USERS_FILE}.bak"
|
||||||
MAX_WORKERS = 8
|
MAX_WORKERS = 8
|
||||||
|
|
||||||
# import logging
|
|
||||||
# logging.basicConfig(
|
|
||||||
# level=logging.INFO,
|
|
||||||
# format='%(asctime)s: [%(levelname)s] %(message)s',
|
|
||||||
# datefmt='%Y-%m-%d %H:%M:%S'
|
|
||||||
# )
|
|
||||||
# logger = logging.getLogger()
|
|
||||||
# null_handler = logging.NullHandler()
|
|
||||||
# logger.handlers = [null_handler]
|
|
||||||
|
|
||||||
def acquire_lock():
|
def acquire_lock():
|
||||||
"""Acquires a lock file to prevent concurrent execution"""
|
"""Acquires a lock file to prevent concurrent execution"""
|
||||||
try:
|
try:
|
||||||
@ -36,14 +26,7 @@ def acquire_lock():
|
|||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
def traffic_status(no_gui=False):
|
def traffic_status(no_gui=False):
|
||||||
"""Updates and retrieves traffic statistics for all users.
|
"""Updates and retrieves traffic statistics for all users."""
|
||||||
|
|
||||||
Args:
|
|
||||||
no_gui (bool): If True, suppresses output to console
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
dict: User data including upload/download bytes and status
|
|
||||||
"""
|
|
||||||
green = '\033[0;32m'
|
green = '\033[0;32m'
|
||||||
cyan = '\033[0;36m'
|
cyan = '\033[0;36m'
|
||||||
NC = '\033[0m'
|
NC = '\033[0m'
|
||||||
@ -90,8 +73,7 @@ def traffic_status(no_gui=False):
|
|||||||
users_data[user_id]["status"] = "Online" if status.is_online else "Offline"
|
users_data[user_id]["status"] = "Online" if status.is_online else "Offline"
|
||||||
else:
|
else:
|
||||||
users_data[user_id] = {
|
users_data[user_id] = {
|
||||||
"upload_bytes": 0,
|
"upload_bytes": 0, "download_bytes": 0,
|
||||||
"download_bytes": 0,
|
|
||||||
"status": "Online" if status.is_online else "Offline"
|
"status": "Online" if status.is_online else "Offline"
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -102,11 +84,23 @@ def traffic_status(no_gui=False):
|
|||||||
else:
|
else:
|
||||||
online = user_id in online_status and online_status[user_id].is_online
|
online = user_id in online_status and online_status[user_id].is_online
|
||||||
users_data[user_id] = {
|
users_data[user_id] = {
|
||||||
"upload_bytes": stats.upload_bytes,
|
"upload_bytes": stats.upload_bytes, "download_bytes": stats.download_bytes,
|
||||||
"download_bytes": stats.download_bytes,
|
|
||||||
"status": "Online" if online else "Offline"
|
"status": "Online" if online else "Offline"
|
||||||
}
|
}
|
||||||
|
|
||||||
|
today_date = datetime.datetime.now().strftime("%Y-%m-%d")
|
||||||
|
for username, user_data in users_data.items():
|
||||||
|
is_on_hold = not user_data.get("account_creation_date")
|
||||||
|
|
||||||
|
if is_on_hold:
|
||||||
|
is_online = user_data.get("status") == "Online"
|
||||||
|
has_traffic = user_data.get("download_bytes", 0) > 0 or user_data.get("upload_bytes", 0) > 0
|
||||||
|
|
||||||
|
if is_online or has_traffic:
|
||||||
|
user_data["account_creation_date"] = today_date
|
||||||
|
else:
|
||||||
|
user_data["status"] = "On-hold"
|
||||||
|
|
||||||
with open(USERS_FILE, 'w') as users_file:
|
with open(USERS_FILE, 'w') as users_file:
|
||||||
json.dump(users_data, users_file, indent=4)
|
json.dump(users_data, users_file, indent=4)
|
||||||
|
|
||||||
@ -137,66 +131,49 @@ def display_traffic_data(data, green, cyan, NC):
|
|||||||
print(f"{user:<15} {green}{formatted_tx:<15}{NC} {cyan}{formatted_rx:<15}{NC} {status:<10}")
|
print(f"{user:<15} {green}{formatted_tx:<15}{NC} {cyan}{formatted_rx:<15}{NC} {status:<10}")
|
||||||
print("-------------------------------------------------")
|
print("-------------------------------------------------")
|
||||||
|
|
||||||
def format_bytes(bytes):
|
def format_bytes(bytes_val):
|
||||||
"""Format bytes as human-readable string"""
|
"""Format bytes as human-readable string"""
|
||||||
if bytes < 1024:
|
if bytes_val < 1024: return f"{bytes_val}B"
|
||||||
return f"{bytes}B"
|
elif bytes_val < 1048576: return f"{bytes_val / 1024:.2f}KB"
|
||||||
elif bytes < 1048576:
|
elif bytes_val < 1073741824: return f"{bytes_val / 1048576:.2f}MB"
|
||||||
return f"{bytes / 1024:.2f}KB"
|
elif bytes_val < 1099511627776: return f"{bytes_val / 1073741824:.2f}GB"
|
||||||
elif bytes < 1073741824:
|
else: return f"{bytes_val / 1099511627776:.2f}TB"
|
||||||
return f"{bytes / 1048576:.2f}MB"
|
|
||||||
elif bytes < 1099511627776:
|
|
||||||
return f"{bytes / 1073741824:.2f}GB"
|
|
||||||
else:
|
|
||||||
return f"{bytes / 1099511627776:.2f}TB"
|
|
||||||
|
|
||||||
def kick_users(usernames, secret):
|
def kick_users(usernames, secret):
|
||||||
"""Kicks specified users from the server"""
|
"""Kicks specified users from the server"""
|
||||||
try:
|
try:
|
||||||
client = Hysteria2Client(
|
client = Hysteria2Client(base_url=API_BASE_URL, secret=secret)
|
||||||
base_url=API_BASE_URL,
|
|
||||||
secret=secret
|
|
||||||
)
|
|
||||||
|
|
||||||
client.kick_clients(usernames)
|
client.kick_clients(usernames)
|
||||||
return True
|
return True
|
||||||
except Exception:
|
except Exception:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
def process_user(username, user_data, config_secret, users_data):
|
def process_user(username, user_data, users_data):
|
||||||
"""Process a single user to check if they should be kicked"""
|
"""Process a single user to check if they should be kicked"""
|
||||||
blocked = user_data.get('blocked', False)
|
if user_data.get('blocked', False): return None
|
||||||
|
|
||||||
if blocked:
|
account_creation_date = user_data.get('account_creation_date')
|
||||||
return None
|
if not account_creation_date: return None
|
||||||
|
|
||||||
max_download_bytes = user_data.get('max_download_bytes', 0)
|
max_download_bytes = user_data.get('max_download_bytes', 0)
|
||||||
expiration_days = user_data.get('expiration_days', 0)
|
expiration_days = user_data.get('expiration_days', 0)
|
||||||
account_creation_date = user_data.get('account_creation_date')
|
total_bytes = user_data.get('download_bytes', 0) + user_data.get('upload_bytes', 0)
|
||||||
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:
|
|
||||||
return None
|
|
||||||
|
|
||||||
|
should_block = False
|
||||||
try:
|
try:
|
||||||
current_date = datetime.datetime.now().timestamp()
|
if expiration_days > 0:
|
||||||
creation_date = datetime.datetime.fromisoformat(account_creation_date.replace('Z', '+00:00'))
|
creation_date = datetime.datetime.strptime(account_creation_date, "%Y-%m-%d")
|
||||||
expiration_date = (creation_date + datetime.timedelta(days=expiration_days)).timestamp()
|
expiration_date = creation_date + datetime.timedelta(days=expiration_days)
|
||||||
|
if datetime.datetime.now() >= expiration_date:
|
||||||
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
|
should_block = True
|
||||||
|
|
||||||
if should_block:
|
if not should_block and max_download_bytes > 0 and total_bytes >= max_download_bytes:
|
||||||
users_data[username]['blocked'] = True
|
should_block = True
|
||||||
return username
|
|
||||||
|
|
||||||
except Exception:
|
if should_block:
|
||||||
|
users_data[username]['blocked'] = True
|
||||||
|
return username
|
||||||
|
except (ValueError, TypeError):
|
||||||
return None
|
return None
|
||||||
|
|
||||||
return None
|
return None
|
||||||
@ -206,59 +183,39 @@ def kick_expired_users():
|
|||||||
lock_file = acquire_lock()
|
lock_file = acquire_lock()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
if not os.path.exists(USERS_FILE): return
|
||||||
shutil.copy2(USERS_FILE, BACKUP_FILE)
|
shutil.copy2(USERS_FILE, BACKUP_FILE)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
with open(CONFIG_FILE, 'r') as f:
|
with open(CONFIG_FILE, 'r') as f:
|
||||||
config = json.load(f)
|
config = json.load(f)
|
||||||
secret = config.get('trafficStats', {}).get('secret', '')
|
secret = config.get('trafficStats', {}).get('secret', '')
|
||||||
if not secret:
|
if not secret: sys.exit(1)
|
||||||
sys.exit(1)
|
|
||||||
except Exception:
|
except Exception:
|
||||||
shutil.copy2(BACKUP_FILE, USERS_FILE)
|
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
try:
|
with open(USERS_FILE, 'r') as f:
|
||||||
with open(USERS_FILE, 'r') as f:
|
users_data = json.load(f)
|
||||||
users_data = json.load(f)
|
|
||||||
except json.JSONDecodeError:
|
|
||||||
shutil.copy2(BACKUP_FILE, USERS_FILE)
|
|
||||||
sys.exit(1)
|
|
||||||
except Exception:
|
|
||||||
shutil.copy2(BACKUP_FILE, USERS_FILE)
|
|
||||||
sys.exit(1)
|
|
||||||
|
|
||||||
users_to_kick = []
|
users_to_kick = []
|
||||||
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
|
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
|
||||||
future_to_user = {
|
futures = [executor.submit(process_user, u, d, users_data) for u, d in users_data.items()]
|
||||||
executor.submit(process_user, username, user_data, secret, users_data): username
|
for future in futures:
|
||||||
for username, user_data in users_data.items()
|
result = future.result()
|
||||||
}
|
if result:
|
||||||
|
users_to_kick.append(result)
|
||||||
for future in future_to_user:
|
|
||||||
username = future.result()
|
|
||||||
if username:
|
|
||||||
users_to_kick.append(username)
|
|
||||||
|
|
||||||
if users_to_kick:
|
if users_to_kick:
|
||||||
for retry in range(3):
|
with open(USERS_FILE, 'w') as f:
|
||||||
try:
|
json.dump(users_data, f, indent=4)
|
||||||
with open(USERS_FILE, 'w') as f:
|
|
||||||
json.dump(users_data, f, indent=2)
|
|
||||||
break
|
|
||||||
except Exception:
|
|
||||||
time.sleep(1)
|
|
||||||
if retry == 2:
|
|
||||||
raise
|
|
||||||
|
|
||||||
if users_to_kick:
|
for i in range(0, len(users_to_kick), 50):
|
||||||
batch_size = 50
|
batch = users_to_kick[i:i+50]
|
||||||
for i in range(0, len(users_to_kick), batch_size):
|
|
||||||
batch = users_to_kick[i:i+batch_size]
|
|
||||||
kick_users(batch, secret)
|
kick_users(batch, secret)
|
||||||
|
|
||||||
except Exception:
|
except Exception:
|
||||||
shutil.copy2(BACKUP_FILE, USERS_FILE)
|
if os.path.exists(BACKUP_FILE):
|
||||||
|
shutil.copy2(BACKUP_FILE, USERS_FILE)
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
finally:
|
finally:
|
||||||
fcntl.flock(lock_file, fcntl.LOCK_UN)
|
fcntl.flock(lock_file, fcntl.LOCK_UN)
|
||||||
|
|||||||
Reference in New Issue
Block a user