bivashy/Blitz-Proxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: migrate user management from json to mongodb

Browse Source
This commit is contained in:
Whispering Wind
2025-09-06 22:38:32 +03:30
committed by GitHub
parent 15eac57988
commit 80c21ccd07
17 changed files with 824 additions and 994 deletions
Download Patch File Download Diff File Expand all files Collapse all files

331
core/traffic.py
View File

@ -1,23 +1,22 @@
#!/usr/bin/env python3
import json
import os
import sys
import time
import fcntl
import shutil
import datetime
from concurrent.futures import ThreadPoolExecutor
from hysteria2_api import Hysteria2Client
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(SCRIPT_DIR, 'scripts'))
from db.database import db
CONFIG_FILE = '/etc/hysteria/config.json'
USERS_FILE = '/etc/hysteria/users.json'
API_BASE_URL = 'http://127.0.0.1:25413'
LOCKFILE = "/tmp/kick.lock"
BACKUP_FILE = f"{USERS_FILE}.bak"
MAX_WORKERS = 8
LOCKFILE = "/tmp/hysteria_traffic.lock"
def acquire_lock():
"""Acquires a lock file to prevent concurrent execution"""
try:
lock_file = open(LOCKFILE, 'w')
fcntl.flock(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
@ -25,92 +24,22 @@ def acquire_lock():
except IOError:
sys.exit(1)
def traffic_status(no_gui=False):
"""Updates and retrieves traffic statistics for all users."""
green = '\033[0;32m'
cyan = '\033[0;36m'
NC = '\033[0m'
def get_secret():
try:
with open(CONFIG_FILE, 'r') as config_file:
config = json.load(config_file)
secret = config.get('trafficStats', {}).get('secret')
except (json.JSONDecodeError, FileNotFoundError) as e:
if not no_gui:
print(f"Error: Failed to read secret from {CONFIG_FILE}. Details: {e}")
with open(CONFIG_FILE, 'r') as f:
config = json.load(f)
return config.get('trafficStats', {}).get('secret')
except (json.JSONDecodeError, FileNotFoundError):
return None
if not secret:
if not no_gui:
print("Error: Secret not found in config.json")
return None
client = Hysteria2Client(base_url=API_BASE_URL, secret=secret)
try:
traffic_stats = client.get_traffic_stats(clear=True)
online_status = client.get_online_clients()
except Exception as e:
if not no_gui:
print(f"Error communicating with Hysteria2 API: {e}")
return None
users_data = {}
if os.path.exists(USERS_FILE):
try:
with open(USERS_FILE, 'r') as users_file:
users_data = json.load(users_file)
except json.JSONDecodeError:
if not no_gui:
print("Error: Failed to parse existing users data JSON file.")
return None
for user in users_data:
users_data[user]["status"] = "Offline"
for user_id, status in online_status.items():
if user_id in users_data:
users_data[user_id]["status"] = "Online" if status.is_online else "Offline"
else:
users_data[user_id] = {
"upload_bytes": 0, "download_bytes": 0,
"status": "Online" if status.is_online else "Offline"
}
for user_id, stats in traffic_stats.items():
if user_id in users_data:
users_data[user_id]["upload_bytes"] = users_data[user_id].get("upload_bytes", 0) + stats.upload_bytes
users_data[user_id]["download_bytes"] = users_data[user_id].get("download_bytes", 0) + stats.download_bytes
else:
online = user_id in online_status and online_status[user_id].is_online
users_data[user_id] = {
"upload_bytes": stats.upload_bytes, "download_bytes": stats.download_bytes,
"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:
json.dump(users_data, users_file, indent=4)
if not no_gui:
display_traffic_data(users_data, green, cyan, NC)
return users_data
def format_bytes(bytes_val):
if bytes_val < 1024: return f"{bytes_val}B"
elif bytes_val < 1048576: return f"{bytes_val / 1024:.2f}KB"
elif bytes_val < 1073741824: return f"{bytes_val / 1048576:.2f}MB"
elif bytes_val < 1099511627776: return f"{bytes_val / 1073741824:.2f}GB"
else: return f"{bytes_val / 1099511627776:.2f}TB"
def display_traffic_data(data, green, cyan, NC):
"""Displays traffic data in a formatted table"""
if not data:
print("No traffic data to display.")
return
@ -123,113 +52,151 @@ def display_traffic_data(data, green, cyan, NC):
for user, entry in data.items():
upload_bytes = entry.get("upload_bytes", 0)
download_bytes = entry.get("download_bytes", 0)
status = entry.get("status", "Offline")
status = entry.get("status", "On-hold")
formatted_tx = format_bytes(upload_bytes)
formatted_rx = format_bytes(download_bytes)
print(f"{user:<15} {green}{formatted_tx:<15}{NC} {cyan}{formatted_rx:<15}{NC} {status:<10}")
print("-------------------------------------------------")
def format_bytes(bytes_val):
"""Format bytes as human-readable string"""
if bytes_val < 1024: return f"{bytes_val}B"
elif bytes_val < 1048576: return f"{bytes_val / 1024:.2f}KB"
elif bytes_val < 1073741824: return f"{bytes_val / 1048576:.2f}MB"
elif bytes_val < 1099511627776: return f"{bytes_val / 1073741824:.2f}GB"
else: return f"{bytes_val / 1099511627776:.2f}TB"
def traffic_status(no_gui=False):
green, cyan, NC = '\033[0;32m', '\033[0;36m', '\033[0m'
if db is None:
if not no_gui: print("Error: Database connection failed.")
return None
def kick_users(usernames, secret):
"""Kicks specified users from the server"""
secret = get_secret()
if not secret:
if not no_gui: print(f"Error: Secret not found or failed to read {CONFIG_FILE}.")
return None
client = Hysteria2Client(base_url=API_BASE_URL, secret=secret)
try:
traffic_stats = client.get_traffic_stats(clear=True)
online_status = client.get_online_clients()
except Exception as e:
if not no_gui: print(f"Error communicating with Hysteria2 API: {e}")
return None
try:
all_users = db.get_all_users()
initial_users_data = {user['_id']: user for user in all_users}
except Exception as e:
if not no_gui: print(f"Error fetching users from database: {e}")
return None
today_date = datetime.datetime.now().strftime("%Y-%m-%d")
users_to_update = {}
for username, user_data in initial_users_data.items():
updates = {}
is_online = username in online_status and online_status[username].is_online
if username in traffic_stats:
new_upload = user_data.get('upload_bytes', 0) + traffic_stats[username].upload_bytes
new_download = user_data.get('download_bytes', 0) + traffic_stats[username].download_bytes
if new_upload != user_data.get('upload_bytes'): updates['upload_bytes'] = new_upload
if new_download != user_data.get('download_bytes'): updates['download_bytes'] = new_download
is_activated = "account_creation_date" in user_data
if not is_activated:
current_traffic = traffic_stats.get(username)
has_activity = is_online or (current_traffic and (current_traffic.upload_bytes > 0 or current_traffic.download_bytes > 0))
if has_activity:
updates["account_creation_date"] = today_date
updates["status"] = "Online" if is_online else "Offline"
else:
if user_data.get("status") != "On-hold":
updates["status"] = "On-hold"
else:
new_status = "Online" if is_online else "Offline"
if user_data.get("status") != new_status:
updates["status"] = new_status
if updates:
users_to_update[username] = updates
if users_to_update:
try:
for username, update_data in users_to_update.items():
db.update_user(username, update_data)
except Exception as e:
if not no_gui: print(f"Error updating database: {e}")
return None
if not no_gui:
# For display, merge updates into the initial data
for username, updates in users_to_update.items():
initial_users_data[username].update(updates)
display_traffic_data(initial_users_data, green, cyan, NC)
return initial_users_data
def kick_api_call(usernames, secret):
try:
client = Hysteria2Client(base_url=API_BASE_URL, secret=secret)
client.kick_clients(usernames)
return True
except Exception:
return False
def process_user(username, user_data, users_data):
"""Process a single user to check if they should be kicked"""
if user_data.get('blocked', False): return None
account_creation_date = user_data.get('account_creation_date')
if not account_creation_date: return None
max_download_bytes = user_data.get('max_download_bytes', 0)
expiration_days = user_data.get('expiration_days', 0)
total_bytes = user_data.get('download_bytes', 0) + user_data.get('upload_bytes', 0)
should_block = False
try:
if expiration_days > 0:
creation_date = datetime.datetime.strptime(account_creation_date, "%Y-%m-%d")
expiration_date = creation_date + datetime.timedelta(days=expiration_days)
if datetime.datetime.now() >= expiration_date:
should_block = True
if not should_block and max_download_bytes > 0 and total_bytes >= max_download_bytes:
should_block = True
if should_block:
users_data[username]['blocked'] = True
return username
except (ValueError, TypeError):
return None
return None
except Exception as e:
print(f"Failed to kick users via API: {e}", file=sys.stderr)
def kick_expired_users():
"""Kicks users who have exceeded their data limits or whose accounts have expired"""
lock_file = acquire_lock()
if db is None:
print("Error: Database connection failed.", file=sys.stderr)
return
secret = get_secret()
if not secret:
print(f"Error: Secret not found or failed to read {CONFIG_FILE}.", file=sys.stderr)
return
try:
if not os.path.exists(USERS_FILE): return
shutil.copy2(USERS_FILE, BACKUP_FILE)
all_users = db.get_all_users()
users_to_kick, users_to_block = [], []
for user in all_users:
username = user.get('_id')
if not username or user.get('blocked', False) or not user.get('account_creation_date'):
continue
total_bytes = user.get('download_bytes', 0) + user.get('upload_bytes', 0)
should_block = False
try:
with open(CONFIG_FILE, 'r') as f:
config = json.load(f)
secret = config.get('trafficStats', {}).get('secret', '')
if not secret: sys.exit(1)
except Exception:
sys.exit(1)
if user.get('expiration_days', 0) > 0:
creation_date = datetime.datetime.strptime(user['account_creation_date'], "%Y-%m-%d")
if datetime.datetime.now() >= creation_date + datetime.timedelta(days=user['expiration_days']):
should_block = True
with open(USERS_FILE, 'r') as f:
users_data = json.load(f)
users_to_kick = []
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
futures = [executor.submit(process_user, u, d, users_data) for u, d in users_data.items()]
for future in futures:
result = future.result()
if result:
users_to_kick.append(result)
if users_to_kick:
with open(USERS_FILE, 'w') as f:
json.dump(users_data, f, indent=4)
for i in range(0, len(users_to_kick), 50):
batch = users_to_kick[i:i+50]
kick_users(batch, secret)
except Exception:
if os.path.exists(BACKUP_FILE):
shutil.copy2(BACKUP_FILE, USERS_FILE)
sys.exit(1)
finally:
fcntl.flock(lock_file, fcntl.LOCK_UN)
lock_file.close()
if not should_block and user.get('max_download_bytes', 0) > 0 and total_bytes >= user['max_download_bytes']:
should_block = True
if should_block:
users_to_kick.append(username)
users_to_block.append(username)
except (ValueError, TypeError):
continue
if users_to_block:
for username in users_to_block:
db.update_user(username, {'blocked': True})
if users_to_kick:
for i in range(0, len(users_to_kick), 50):
kick_api_call(users_to_kick[i:i+50], secret)
if __name__ == "__main__":
if len(sys.argv) > 1:
if sys.argv[1] == "kick":
kick_expired_users()
elif sys.argv[1] == "--no-gui":
traffic_status(no_gui=True)
kick_expired_users()
lock_file = acquire_lock()
try:
if len(sys.argv) > 1:
if sys.argv[1] == "kick":
kick_expired_users()
elif sys.argv[1] == "--no-gui":
traffic_status(no_gui=True)
kick_expired_users()
else:
print(f"Usage: python {sys.argv[0]} [kick|--no-gui]")
else:
print(f"Unknown argument: {sys.argv[1]}")
print("Usage: python traffic.py [kick|--no-gui]")
else:
traffic_status(no_gui=False)
traffic_status(no_gui=False)
finally:
fcntl.flock(lock_file, fcntl.LOCK_UN)
lock_file.close()