EG/plugins/user/realtime_communication/clients/client_manager.py

"""
客户端管理器模块
管理客户端连接和状态
"""

import time
import threading
import hashlib
from typing import Dict, Any, List, Optional

class ClientManager:
    """
    客户端管理器
    管理客户端连接和状态
    """
    
    def __init__(self, plugin):
        """
        初始化客户端管理器
        
        Args:
            plugin: 实时通信插件实例
        """
        self.plugin = plugin
        self.enabled = False
        self.initialized = False
        
        # 客户端配置
        self.client_config = {
            "enable_authentication": True,
            "auth_timeout": 30.0,
            "max_clients": 1000,
            "enable_heartbeat": True,
            "heartbeat_interval": 30.0,
            "enable_rate_limiting": True,
            "max_messages_per_second": 100,
            "enable_banning": True,
            "ban_duration": 3600,
            "client_timeout": 60.0
        }
        
        # 客户端状态
        self.client_state = {
            "current_clients": 0,
            "total_clients": 0,
            "authenticated_clients": 0,
            "banned_clients": 0
        }
        
        # 客户端存储
        self.clients = {}
        self.client_lock = threading.RLock()
        
        # 客户端统计
        self.client_stats = {
            "clients_joined": 0,
            "clients_left": 0,
            "authentication_success": 0,
            "authentication_failed": 0,
            "messages_sent": 0,
            "messages_received": 0,
            "banned_clients": 0,
            "errors": 0
        }
        
        # 黑名单管理
        self.banned_ips = {}
        self.banned_users = {}
        
        # 回调函数
        self.client_callbacks = {
            "client_joined": [],
            "client_left": [],
            "client_authenticated": [],
            "client_banned": [],
            "client_error": []
        }
        
        # 时间戳记录
        self.last_heartbeat_check = 0.0
        self.last_cleanup = 0.0
        
        print("✓ 客户端管理器已创建")
        
    def initialize(self) -> bool:
        """
        初始化客户端管理器
        
        Returns:
            是否初始化成功
        """
        try:
            print("正在初始化客户端管理器...")
            
            self.initialized = True
            print("✓ 客户端管理器初始化完成")
            return True
            
        except Exception as e:
            print(f"✗ 客户端管理器初始化失败: {e}")
            self.client_stats["errors"] += 1
            import traceback
            traceback.print_exc()
            return False
            
    def enable(self) -> bool:
        """
        启用客户端管理器
        
        Returns:
            是否启用成功
        """
        try:
            if not self.initialized:
                print("✗ 客户端管理器未初始化")
                return False
                
            self.enabled = True
            print("✓ 客户端管理器已启用")
            return True
            
        except Exception as e:
            print(f"✗ 客户端管理器启用失败: {e}")
            self.client_stats["errors"] += 1
            import traceback
            traceback.print_exc()
            return False
            
    def disable(self):
        """禁用客户端管理器"""
        try:
            self.enabled = False
            
            # 断开所有客户端
            self._disconnect_all_clients()
            
            print("✓ 客户端管理器已禁用")
            
        except Exception as e:
            print(f"✗ 客户端管理器禁用失败: {e}")
            self.client_stats["errors"] += 1
            import traceback
            traceback.print_exc()
            
    def finalize(self):
        """清理客户端管理器资源"""
        try:
            # 禁用客户端管理器
            if self.enabled:
                self.disable()
                
            # 清理回调
            self.client_callbacks.clear()
            
            self.initialized = False
            print("✓ 客户端管理器资源已清理")
            
        except Exception as e:
            print(f"✗ 客户端管理器资源清理失败: {e}")
            import traceback
            traceback.print_exc()
            
    def update(self, dt: float):
        """
        更新客户端管理器状态
        
        Args:
            dt: 时间增量（秒）
        """
        try:
            if not self.enabled:
                return
                
            current_time = time.time()
            
            # 定期检查心跳
            if current_time - self.last_heartbeat_check > self.client_config["heartbeat_interval"]:
                self._check_client_heartbeats(current_time)
                self.last_heartbeat_check = current_time
                
            # 定期清理
            if current_time - self.last_cleanup > 60.0:  # 每分钟清理一次
                self._cleanup_banned_clients(current_time)
                self.last_cleanup = current_time
                
        except Exception as e:
            print(f"✗ 客户端管理器更新失败: {e}")
            self.client_stats["errors"] += 1
            import traceback
            traceback.print_exc()
            
    def _check_client_heartbeats(self, current_time: float):
        """检查客户端心跳"""
        try:
            if not self.client_config["enable_heartbeat"]:
                return
                
            clients_to_disconnect = []
            with self.client_lock:
                for client_id, client_data in self.clients.items():
                    last_heartbeat = client_data.get("last_heartbeat", client_data["connect_time"])
                    if current_time - last_heartbeat > self.client_config["client_timeout"]:
                        clients_to_disconnect.append(client_id)
                        
            for client_id in clients_to_disconnect:
                self.disconnect_client(client_id, "timeout")
                
        except Exception as e:
            print(f"✗ 客户端心跳检查失败: {e}")
            self.client_stats["errors"] += 1
            
    def _cleanup_banned_clients(self, current_time: float):
        """清理过期的黑名单客户端"""
        try:
            expired_bans = []
            for ip, ban_data in self.banned_ips.items():
                if current_time - ban_data["ban_time"] > self.client_config["ban_duration"]:
                    expired_bans.append(ip)
                    
            for ip in expired_bans:
                del self.banned_ips[ip]
                self.client_state["banned_clients"] -= 1
                
        except Exception as e:
            print(f"✗ 黑名单清理失败: {e}")
            self.client_stats["errors"] += 1
            
    def _disconnect_all_clients(self):
        """断开所有客户端连接"""
        try:
            with self.client_lock:
                client_ids = list(self.clients.keys())
                
            for client_id in client_ids:
                self.disconnect_client(client_id, "server_shutdown")
                
        except Exception as e:
            print(f"✗ 断开所有客户端连接失败: {e}")
            self.client_stats["errors"] += 1
            
    def add_client(self, client_id: str, client_info: Dict[str, Any]) -> bool:
        """
        添加客户端
        
        Args:
            client_id: 客户端ID
            client_info: 客户端信息
            
        Returns:
            是否添加成功
        """
        try:
            if not self.enabled:
                return False
                
            # 检查客户端数量限制
            with self.client_lock:
                if len(self.clients) >= self.client_config["max_clients"]:
                    print("✗ 客户端数量已达上限")
                    return False
                    
            # 检查IP是否被禁止
            client_ip = client_info.get("ip_address", "")
            if self._is_ip_banned(client_ip):
                print(f"✗ IP地址被禁止: {client_ip}")
                return False
                
            # 添加客户端
            client_data = {
                "client_id": client_id,
                "client_info": client_info,
                "connect_time": time.time(),
                "last_heartbeat": time.time(),
                "is_authenticated": False,
                "user_id": "",
                "username": "",
                "permissions": [],
                "messages_sent": 0,
                "messages_received": 0,
                "bytes_sent": 0,
                "bytes_received": 0
            }
            
            with self.client_lock:
                self.clients[client_id] = client_data
                
            self.client_state["current_clients"] = len(self.clients)
            self.client_state["total_clients"] += 1
            self.client_stats["clients_joined"] += 1
            
            # 触发客户端加入回调
            self._trigger_client_callback("client_joined", {
                "client_id": client_id,
                "client_info": client_info,
                "timestamp": time.time()
            })
            
            print(f"✓ 客户端已添加: {client_id}")
            return True
            
        except Exception as e:
            print(f"✗ 添加客户端失败: {e}")
            self.client_stats["errors"] += 1
            return False
            
    def remove_client(self, client_id: str, reason: str = "unknown") -> bool:
        """
        移除客户端
        
        Args:
            client_id: 客户端ID
            reason: 移除原因
            
        Returns:
            是否移除成功
        """
        try:
            with self.client_lock:
                if client_id in self.clients:
                    client_data = self.clients[client_id]
                    if client_data["is_authenticated"]:
                        self.client_state["authenticated_clients"] -= 1
                    del self.clients[client_id]
                    
            self.client_state["current_clients"] = len(self.clients)
            self.client_stats["clients_left"] += 1
            
            # 触发客户端离开回调
            self._trigger_client_callback("client_left", {
                "client_id": client_id,
                "reason": reason,
                "timestamp": time.time()
            })
            
            print(f"✓ 客户端已移除: {client_id} (原因: {reason})")
            return True
            
        except Exception as e:
            print(f"✗ 移除客户端失败: {e}")
            self.client_stats["errors"] += 1
            return False
            
    def disconnect_client(self, client_id: str, reason: str = "kicked") -> bool:
        """
        断开客户端连接
        
        Args:
            client_id: 客户端ID
            reason: 断开原因
            
        Returns:
            是否断开成功
        """
        try:
            # 通知WebSocket服务器断开连接
            if self.plugin.websocket_server:
                self.plugin.websocket_server.disconnect_client(client_id, reason)
                
            # 移除客户端
            return self.remove_client(client_id, reason)
            
        except Exception as e:
            print(f"✗ 断开客户端连接失败: {e}")
            self.client_stats["errors"] += 1
            return False
            
    def authenticate_client(self, client_id: str, auth_data: Dict[str, Any]) -> bool:
        """
        认证客户端
        
        Args:
            client_id: 客户端ID
            auth_data: 认证数据
            
        Returns:
            是否认证成功
        """
        try:
            if not self.enabled or not self.client_config["enable_authentication"]:
                # 如果未启用认证，直接通过
                with self.client_lock:
                    if client_id in self.clients:
                        self.clients[client_id]["is_authenticated"] = True
                return True
                
            # 使用认证管理器进行认证
            if self.plugin.auth_manager:
                auth_result = self.plugin.auth_manager.authenticate_client(client_id, auth_data)
                if auth_result:
                    # 更新客户端认证状态
                    with self.client_lock:
                        if client_id in self.clients:
                            self.clients[client_id]["is_authenticated"] = True
                            self.clients[client_id]["user_id"] = auth_data.get("user_id", "")
                            self.clients[client_id]["username"] = auth_data.get("username", "")
                            self.clients[client_id]["permissions"] = auth_data.get("permissions", [])
                            
                    self.client_state["authenticated_clients"] += 1
                    self.client_stats["authentication_success"] += 1
                    
                    # 触发客户端认证回调
                    self._trigger_client_callback("client_authenticated", {
                        "client_id": client_id,
                        "user_id": auth_data.get("user_id", ""),
                        "username": auth_data.get("username", ""),
                        "timestamp": time.time()
                    })
                    
                    print(f"✓ 客户端认证成功: {client_id}")
                    return True
                else:
                    self.client_stats["authentication_failed"] += 1
                    print(f"✗ 客户端认证失败: {client_id}")
                    return False
            else:
                print("✗ 认证管理器不可用")
                return False
                
        except Exception as e:
            print(f"✗ 客户端认证失败: {e}")
            self.client_stats["errors"] += 1
            return False
            
    def update_client_heartbeat(self, client_id: str) -> bool:
        """
        更新客户端心跳
        
        Args:
            client_id: 客户端ID
            
        Returns:
            是否更新成功
        """
        try:
            with self.client_lock:
                if client_id in self.clients:
                    self.clients[client_id]["last_heartbeat"] = time.time()
                    return True
                else:
                    return False
                    
        except Exception as e:
            print(f"✗ 更新客户端心跳失败: {e}")
            self.client_stats["errors"] += 1
            return False
            
    def get_client(self, client_id: str) -> Optional[Dict[str, Any]]:
        """
        获取客户端信息
        
        Args:
            client_id: 客户端ID
            
        Returns:
            客户端信息或None
        """
        try:
            with self.client_lock:
                if client_id in self.clients:
                    return self.clients[client_id].copy()
                else:
                    return None
                    
        except Exception as e:
            print(f"✗ 获取客户端信息失败: {e}")
            self.client_stats["errors"] += 1
            return None
            
    def get_all_clients(self) -> Dict[str, Any]:
        """
        获取所有客户端信息
        
        Returns:
            所有客户端信息字典
        """
        try:
            with self.client_lock:
                return {k: v.copy() for k, v in self.clients.items()}
        except Exception as e:
            print(f"✗ 获取所有客户端信息失败: {e}")
            self.client_stats["errors"] += 1
            return {}
            
    def ban_client(self, client_id: str, reason: str = "violation", duration: float = 3600) -> bool:
        """
        禁止客户端
        
        Args:
            client_id: 客户端ID
            reason: 禁止原因
            duration: 禁止时长（秒）
            
        Returns:
            是否禁止成功
        """
        try:
            if not self.client_config["enable_banning"]:
                return False
                
            client_data = self.get_client(client_id)
            if not client_data:
                return False
                
            client_ip = client_data["client_info"].get("ip_address", "")
            if client_ip:
                # 添加到IP黑名单
                self.banned_ips[client_ip] = {
                    "client_id": client_id,
                    "ban_time": time.time(),
                    "reason": reason,
                    "duration": duration
                }
                
                self.client_state["banned_clients"] += 1
                self.client_stats["banned_clients"] += 1
                
                # 断开客户端连接
                self.disconnect_client(client_id, f"banned: {reason}")
                
                # 触发客户端禁止回调
                self._trigger_client_callback("client_banned", {
                    "client_id": client_id,
                    "ip_address": client_ip,
                    "reason": reason,
                    "duration": duration,
                    "timestamp": time.time()
                })
                
                print(f"✓ 客户端已被禁止: {client_id} (IP: {client_ip})")
                return True
            else:
                return False
                
        except Exception as e:
            print(f"✗ 禁止客户端失败: {e}")
            self.client_stats["errors"] += 1
            return False
            
    def _is_ip_banned(self, ip_address: str) -> bool:
        """
        检查IP是否被禁止
        
        Args:
            ip_address: IP地址
            
        Returns:
            IP是否被禁止
        """
        try:
            if not self.client_config["enable_banning"]:
                return False
                
            if ip_address in self.banned_ips:
                ban_data = self.banned_ips[ip_address]
                if time.time() - ban_data["ban_time"] < ban_data["duration"]:
                    return True
                else:
                    # 禁止时间已过期，移除
                    del self.banned_ips[ip_address]
                    self.client_state["banned_clients"] -= 1
                    
            return False
            
        except Exception as e:
            print(f"✗ IP禁止检查失败: {e}")
            self.client_stats["errors"] += 1
            return False
            
    def unban_ip(self, ip_address: str) -> bool:
        """
        解除IP禁止
        
        Args:
            ip_address: IP地址
            
        Returns:
            是否解除成功
        """
        try:
            if ip_address in self.banned_ips:
                del self.banned_ips[ip_address]
                self.client_state["banned_clients"] -= 1
                print(f"✓ IP禁止已解除: {ip_address}")
                return True
            else:
                print(f"✗ IP未被禁止: {ip_address}")
                return False
                
        except Exception as e:
            print(f"✗ 解除IP禁止失败: {e}")
            self.client_stats["errors"] += 1
            return False
            
    def is_client_authenticated(self, client_id: str) -> bool:
        """
        检查客户端是否已认证
        
        Args:
            client_id: 客户端ID
            
        Returns:
            客户端是否已认证
        """
        try:
            with self.client_lock:
                if client_id in self.clients:
                    return self.clients[client_id]["is_authenticated"]
                else:
                    return False
                    
        except Exception as e:
            print(f"✗ 客户端认证检查失败: {e}")
            self.client_stats["errors"] += 1
            return False
            
    def get_client_stats(self) -> Dict[str, Any]:
        """
        获取客户端统计信息
        
        Returns:
            客户端统计字典
        """
        return {
            "state": self.client_state.copy(),
            "stats": self.client_stats.copy(),
            "config": self.client_config.copy()
        }
        
    def reset_stats(self):
        """重置客户端统计信息"""
        try:
            self.client_stats = {
                "clients_joined": 0,
                "clients_left": 0,
                "authentication_success": 0,
                "authentication_failed": 0,
                "messages_sent": 0,
                "messages_received": 0,
                "banned_clients": 0,
                "errors": 0
            }
            print("✓ 客户端统计信息已重置")
        except Exception as e:
            print(f"✗ 客户端统计信息重置失败: {e}")
            
    def set_client_config(self, config: Dict[str, Any]) -> bool:
        """
        设置客户端配置
        
        Args:
            config: 客户端配置字典
            
        Returns:
            是否设置成功
        """
        try:
            self.client_config.update(config)
            print(f"✓ 客户端配置已更新: {self.client_config}")
            return True
        except Exception as e:
            print(f"✗ 客户端配置设置失败: {e}")
            return False
            
    def get_client_config(self) -> Dict[str, Any]:
        """
        获取客户端配置
        
        Returns:
            客户端配置字典
        """
        return self.client_config.copy()
        
    def _trigger_client_callback(self, callback_type: str, data: Dict[str, Any]):
        """
        触发客户端回调
        
        Args:
            callback_type: 回调类型
            data: 回调数据
        """
        try:
            if callback_type in self.client_callbacks:
                for callback in self.client_callbacks[callback_type]:
                    try:
                        callback(data)
                    except Exception as e:
                        print(f"✗ 客户端回调执行失败: {callback_type} - {e}")
        except Exception as e:
            print(f"✗ 客户端回调触发失败: {e}")
            
    def register_client_callback(self, callback_type: str, callback: callable):
        """
        注册客户端回调
        
        Args:
            callback_type: 回调类型
            callback: 回调函数
        """
        try:
            if callback_type in self.client_callbacks:
                self.client_callbacks[callback_type].append(callback)
                print(f"✓ 客户端回调已注册: {callback_type}")
            else:
                print(f"✗ 无效的回调类型: {callback_type}")
        except Exception as e:
            print(f"✗ 客户端回调注册失败: {e}")
            
    def unregister_client_callback(self, callback_type: str, callback: callable):
        """
        注销客户端回调
        
        Args:
            callback_type: 回调类型
            callback: 回调函数
        """
        try:
            if callback_type in self.client_callbacks:
                if callback in self.client_callbacks[callback_type]:
                    self.client_callbacks[callback_type].remove(callback)
                    print(f"✓ 客户端回调已注销: {callback_type}")
            else:
                print(f"✗ 无效的回调类型: {callback_type}")
        except Exception as e:
            print(f"✗ 客户端回调注销失败: {e}")