Blitz-Proxy/core/traffic.py

#!/usr/bin/env python3

import json
import os
import sys
import fcntl
import datetime
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'
API_BASE_URL = 'http://127.0.0.1:25413'
LOCKFILE = "/tmp/hysteria_traffic.lock"

def acquire_lock():
    try:
        lock_file = open(LOCKFILE, 'w')
        fcntl.flock(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
        return lock_file
    except IOError:
        sys.exit(1)

def get_secret():
    try:
        with open(CONFIG_FILE, 'r') as f:
            config = json.load(f)
        return config.get('trafficStats', {}).get('secret')
    except (json.JSONDecodeError, FileNotFoundError):
        return None

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):
    if not data:
        print("No traffic data to display.")
        return

    print("Traffic Data:")
    print("-------------------------------------------------")
    print(f"{'User':<15} {'Upload (TX)':<15} {'Download (RX)':<15} {'Status':<10}")
    print("-------------------------------------------------")

    for user, entry in data.items():
        upload_bytes = entry.get("upload_bytes", 0)
        download_bytes = entry.get("download_bytes", 0)
        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 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

    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)
    except Exception as e:
        print(f"Failed to kick users via API: {e}", file=sys.stderr)

def kick_expired_users():
    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
    
    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:
            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
            
            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__":
    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:
            traffic_status(no_gui=False)
    finally:
        fcntl.flock(lock_file, fcntl.LOCK_UN)
        lock_file.close()