"""
安全管理系统模块
负责处理认证、授权和防止攻击
"""

import time
import hashlib
import secrets
import hmac
from typing import Dict, Any, List, Optional

class SecurityManager:
    """
    安全管理器
    负责处理认证、授权和防止攻击
    """
    
    def __init__(self, plugin):
        """
        初始化安全管理器
        
        Args:
            plugin: 服务器架构插件实例
        """
        self.plugin = plugin
        self.enabled = False
        self.initialized = False
        
        # 安全配置
        self.security_config = {
            "enable_authentication": True,
            "enable_encryption": False,
            "enable_message_authentication": True,
            "enable_rate_limiting": True,
            "enable_ddos_protection": True,
            "enable_session_management": True,
            "session_timeout": 3600,  # 1小时
            "max_login_attempts": 5,
            "login_lockout_duration": 900,  # 15分钟
            "enable_ip_filtering": True,
            "enable_captcha": False,
            "captcha_threshold": 10,
            "password_min_length": 8,
            "password_require_complexity": True,
            "enable_two_factor_auth": False,
            "token_lifetime": 86400,  # 24小时
            "enable_audit_logging": True
        }
        
        # 安全状态
        self.security_state = {
            "active_sessions": {},
            "blocked_ips": {},
            "login_attempts": {},
            "security_events": [],
            "last_security_check": 0.0
        }
        
        # 安全统计
        self.security_stats = {
            "successful_authentications": 0,
            "failed_authentications": 0,
            "blocked_connections": 0,
            "security_events": 0,
            "active_sessions": 0,
            "errors": 0
        }
        
        # 密钥管理
        self.secret_key = self._generate_secret_key()
        self.session_keys = {}
        
        # 回调函数
        self.security_callbacks = {
            "authentication_success": [],
            "authentication_failed": [],
            "session_created": [],
            "session_expired": [],
            "security_violation": [],
            "ip_blocked": []
        }
        
        # 时间戳记录
        self.last_security_event = 0.0
        self.last_session_cleanup = 0.0
        
        print("✓ 安全管理器已创建")
        
    def initialize(self) -> bool:
        """
        初始化安全管理器
        
        Returns:
            是否初始化成功
        """
        try:
            print("正在初始化安全管理器...")
            
            # 生成密钥
            self.secret_key = self._generate_secret_key()
            
            self.initialized = True
            print("✓ 安全管理器初始化完成")
            return True
            
        except Exception as e:
            print(f"✗ 安全管理器初始化失败: {e}")
            self.security_stats["errors"] += 1
            import traceback
            traceback.print_exc()
            return False
            
    def _generate_secret_key(self) -> str:
        """
        生成密钥
        
        Returns:
            生成的密钥
        """
        try:
            return secrets.token_hex(32)
        except Exception as e:
            print(f"✗ 密钥生成失败: {e}")
            # 使用默认密钥（仅用于开发环境）
            return "default_secret_key_for_development_purposes_only"
            
    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.security_stats["errors"] += 1
            import traceback
            traceback.print_exc()
            return False
            
    def disable(self):
        """禁用安全管理器"""
        try:
            self.enabled = False
            
            # 清理会话
            self.security_state["active_sessions"].clear()
            self.session_keys.clear()
            
            print("✓ 安全管理器已禁用")
            
        except Exception as e:
            print(f"✗ 安全管理器禁用失败: {e}")
            self.security_stats["errors"] += 1
            import traceback
            traceback.print_exc()
            
    def finalize(self):
        """清理安全管理器资源"""
        try:
            # 禁用安全管理器
            if self.enabled:
                self.disable()
                
            # 清理回调
            self.security_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_session_cleanup > 300:  # 每5分钟清理一次
                self._cleanup_expired_sessions(current_time)
                self.last_session_cleanup = current_time
                
            # 定期清理登录尝试记录
            self._cleanup_login_attempts(current_time)
            
            self.security_state["last_security_check"] = current_time
            
        except Exception as e:
            print(f"✗ 安全管理器更新失败: {e}")
            self.security_stats["errors"] += 1
            import traceback
            traceback.print_exc()
            
    def _cleanup_expired_sessions(self, current_time: float):
        """清理过期会话"""
        try:
            expired_sessions = []
            for session_id, session_data in self.security_state["active_sessions"].items():
                if current_time - session_data["login_time"] > self.security_config["session_timeout"]:
                    expired_sessions.append(session_id)
                    
            for session_id in expired_sessions:
                del self.security_state["active_sessions"][session_id]
                if session_id in self.session_keys:
                    del self.session_keys[session_id]
                    
                self.security_stats["active_sessions"] -= 1
                
                # 触发会话过期回调
                self._trigger_security_callback("session_expired", {
                    "session_id": session_id,
                    "timestamp": current_time
                })
                
        except Exception as e:
            print(f"✗ 过期会话清理失败: {e}")
            self.security_stats["errors"] += 1
            
    def _cleanup_login_attempts(self, current_time: float):
        """清理登录尝试记录"""
        try:
            expired_attempts = []
            for ip_address, attempt_data in self.security_state["login_attempts"].items():
                if current_time - attempt_data["last_attempt"] > self.security_config["login_lockout_duration"]:
                    expired_attempts.append(ip_address)
                    
            for ip_address in expired_attempts:
                del self.security_state["login_attempts"][ip_address]
                
        except Exception as e:
            print(f"✗ 登录尝试记录清理失败: {e}")
            self.security_stats["errors"] += 1
            
    def authenticate_client(self, client_id: str, auth_data: Dict[str, Any], client_address: tuple) -> bool:
        """
        认证客户端
        
        Args:
            client_id: 客户端ID
            auth_data: 认证数据
            client_address: 客户端地址
            
        Returns:
            是否认证成功
        """
        try:
            if not self.enabled or not self.security_config["enable_authentication"]:
                # 如果未启用认证，直接通过
                return True
                
            client_ip = client_address[0] if isinstance(client_address, tuple) else str(client_address)
            
            # 检查IP是否被阻止
            if self._is_ip_blocked(client_ip):
                self.security_stats["blocked_connections"] += 1
                return False
                
            # 检查登录尝试次数
            if self._is_login_blocked(client_ip):
                self.security_stats["blocked_connections"] += 1
                return False
                
            # 执行认证
            auth_result = self._perform_authentication(client_id, auth_data, client_address)
            
            if auth_result:
                self.security_stats["successful_authentications"] += 1
                
                # 触发认证成功回调
                self._trigger_security_callback("authentication_success", {
                    "client_id": client_id,
                    "client_address": client_address,
                    "timestamp": time.time()
                })
                
                # 记录安全事件
                self._log_security_event("authentication_success", client_id, client_address)
                
            else:
                self.security_stats["failed_authentications"] += 1
                
                # 记录失败的登录尝试
                self._record_failed_login(client_ip)
                
                # 触发认证失败回调
                self._trigger_security_callback("authentication_failed", {
                    "client_id": client_id,
                    "client_address": client_address,
                    "timestamp": time.time()
                })
                
                # 记录安全事件
                self._log_security_event("authentication_failed", client_id, client_address)
                
            return auth_result
            
        except Exception as e:
            print(f"✗ 客户端认证失败: {e}")
            self.security_stats["errors"] += 1
            return False
            
    def _is_ip_blocked(self, client_ip: str) -> bool:
        """
        检查IP是否被阻止
        
        Args:
            client_ip: 客户端IP
            
        Returns:
            IP是否被阻止
        """
        try:
            if not self.security_config["enable_ip_filtering"]:
                return False
                
            if client_ip in self.security_state["blocked_ips"]:
                block_data = self.security_state["blocked_ips"][client_ip]
                if time.time() - block_data["block_time"] < block_data["duration"]:
                    return True
                else:
                    # 阻止时间已过期，移除
                    del self.security_state["blocked_ips"][client_ip]
                    
            return False
            
        except Exception as e:
            print(f"✗ IP阻止检查失败: {e}")
            self.security_stats["errors"] += 1
            return False
            
    def _is_login_blocked(self, client_ip: str) -> bool:
        """
        检查登录是否被阻止
        
        Args:
            client_ip: 客户端IP
            
        Returns:
            登录是否被阻止
        """
        try:
            if client_ip in self.security_state["login_attempts"]:
                attempts = self.security_state["login_attempts"][client_ip]
                if (attempts["count"] >= self.security_config["max_login_attempts"] and
                    time.time() - attempts["last_attempt"] < self.security_config["login_lockout_duration"]):
                    return True
                    
            return False
            
        except Exception as e:
            print(f"✗ 登录阻止检查失败: {e}")
            self.security_stats["errors"] += 1
            return False
            
    def _record_failed_login(self, client_ip: str):
        """
        记录失败的登录尝试
        
        Args:
            client_ip: 客户端IP
        """
        try:
            current_time = time.time()
            
            if client_ip not in self.security_state["login_attempts"]:
                self.security_state["login_attempts"][client_ip] = {
                    "count": 1,
                    "last_attempt": current_time
                }
            else:
                self.security_state["login_attempts"][client_ip]["count"] += 1
                self.security_state["login_attempts"][client_ip]["last_attempt"] = current_time
                
            # 如果超过最大尝试次数，阻止IP
            if (self.security_state["login_attempts"][client_ip]["count"] >= 
                self.security_config["max_login_attempts"]):
                self._block_ip(client_ip, self.security_config["login_lockout_duration"])
                
        except Exception as e:
            print(f"✗ 失败登录记录失败: {e}")
            self.security_stats["errors"] += 1
            
    def _perform_authentication(self, client_id: str, auth_data: Dict[str, Any], client_address: tuple) -> bool:
        """
        执行认证
        
        Args:
            client_id: 客户端ID
            auth_data: 认证数据
            client_address: 客户端地址
            
        Returns:
            是否认证成功
        """
        try:
            # 检查认证类型
            auth_type = auth_data.get("auth_type", "token")
            
            if auth_type == "token":
                return self._authenticate_with_token(client_id, auth_data, client_address)
            elif auth_type == "password":
                return self._authenticate_with_password(client_id, auth_data, client_address)
            elif auth_type == "session":
                return self._authenticate_with_session(client_id, auth_data, client_address)
            else:
                # 默认认证方式
                return self._authenticate_with_token(client_id, auth_data, client_address)
                
        except Exception as e:
            print(f"✗ 认证执行失败: {e}")
            self.security_stats["errors"] += 1
            return False
            
    def _authenticate_with_token(self, client_id: str, auth_data: Dict[str, Any], client_address: tuple) -> bool:
        """
        使用令牌认证
        
        Args:
            client_id: 客户端ID
            auth_data: 认证数据
            client_address: 客户端地址
            
        Returns:
            是否认证成功
        """
        try:
            token = auth_data.get("token")
            if not token:
                return False
                
            # 验证令牌（简化实现）
            # 在实际应用中，这里应该连接到用户数据库验证令牌
            if token == "valid_test_token":
                # 创建会话
                self._create_session(client_id, auth_data, client_address)
                return True
            else:
                return False
                
        except Exception as e:
            print(f"✗ 令牌认证失败: {e}")
            self.security_stats["errors"] += 1
            return False
            
    def _authenticate_with_password(self, client_id: str, auth_data: Dict[str, Any], client_address: tuple) -> bool:
        """
        使用密码认证
        
        Args:
            client_id: 客户端ID
            auth_data: 认证数据
            client_address: 客户端地址
            
        Returns:
            是否认证成功
        """
        try:
            username = auth_data.get("username")
            password = auth_data.get("password")
            
            if not username or not password:
                return False
                
            # 验证密码（简化实现）
            # 在实际应用中，这里应该连接到用户数据库验证用户名和密码
            if username == "test_user" and password == "test_password":
                # 创建会话
                self._create_session(client_id, auth_data, client_address)
                return True
            else:
                return False
                
        except Exception as e:
            print(f"✗ 密码认证失败: {e}")
            self.security_stats["errors"] += 1
            return False
            
    def _authenticate_with_session(self, client_id: str, auth_data: Dict[str, Any], client_address: tuple) -> bool:
        """
        使用会话认证
        
        Args:
            client_id: 客户端ID
            auth_data: 认证数据
            client_address: 客户端地址
            
        Returns:
            是否认证成功
        """
        try:
            session_id = auth_data.get("session_id")
            if not session_id:
                return False
                
            # 检查会话是否存在且有效
            if session_id in self.security_state["active_sessions"]:
                session_data = self.security_state["active_sessions"][session_id]
                if time.time() - session_data["login_time"] < self.security_config["session_timeout"]:
                    # 更新会话最后活动时间
                    session_data["last_activity"] = time.time()
                    return True
                    
            return False
            
        except Exception as e:
            print(f"✗ 会话认证失败: {e}")
            self.security_stats["errors"] += 1
            return False
            
    def _create_session(self, client_id: str, auth_data: Dict[str, Any], client_address: tuple):
        """
        创建会话
        
        Args:
            client_id: 客户端ID
            auth_data: 认证数据
            client_address: 客户端地址
        """
        try:
            session_id = self._generate_session_id()
            
            session_data = {
                "client_id": client_id,
                "client_address": client_address,
                "login_time": time.time(),
                "last_activity": time.time(),
                "user_id": auth_data.get("user_id", ""),
                "username": auth_data.get("username", ""),
                "permissions": auth_data.get("permissions", [])
            }
            
            self.security_state["active_sessions"][session_id] = session_data
            self.security_stats["active_sessions"] += 1
            
            # 生成会话密钥
            if self.security_config["enable_encryption"]:
                self.session_keys[session_id] = self._generate_session_key()
                
            # 触发会话创建回调
            self._trigger_security_callback("session_created", {
                "session_id": session_id,
                "client_id": client_id,
                "client_address": client_address,
                "timestamp": time.time()
            })
            
        except Exception as e:
            print(f"✗ 会话创建失败: {e}")
            self.security_stats["errors"] += 1
            
    def _generate_session_id(self) -> str:
        """
        生成会话ID
        
        Returns:
            生成的会话ID
        """
        try:
            return secrets.token_urlsafe(32)
        except Exception as e:
            print(f"✗ 会话ID生成失败: {e}")
            self.security_stats["errors"] += 1
            return f"session_{int(time.time() * 1000000)}"
            
    def _generate_session_key(self) -> str:
        """
        生成会话密钥
        
        Returns:
            生成的会话密钥
        """
        try:
            return secrets.token_hex(16)
        except Exception as e:
            print(f"✗ 会话密钥生成失败: {e}")
            self.security_stats["errors"] += 1
            return f"key_{int(time.time() * 1000000)}"
            
    def validate_token(self, token: str) -> bool:
        """
        验证令牌
        
        Args:
            token: 令牌
            
        Returns:
            令牌是否有效
        """
        try:
            if not token:
                return False
                
            # 简化实现：检查令牌格式
            # 在实际应用中，这里应该连接到令牌数据库验证令牌
            return len(token) > 10 and token.isalnum()
            
        except Exception as e:
            print(f"✗ 令牌验证失败: {e}")
            self.security_stats["errors"] += 1
            return False
            
    def hash_password(self, password: str, salt: str = None) -> Dict[str, str]:
        """
        哈希密码
        
        Args:
            password: 密码
            salt: 盐值（可选）
            
        Returns:
            包含哈希值和盐值的字典
        """
        try:
            if salt is None:
                salt = secrets.token_hex(16)
                
            # 使用PBKDF2哈希密码
            password_hash = hashlib.pbkdf2_hmac('sha256', password.encode('utf-8'), salt.encode('utf-8'), 100000)
            return {
                "hash": password_hash.hex(),
                "salt": salt
            }
            
        except Exception as e:
            print(f"✗ 密码哈希失败: {e}")
            self.security_stats["errors"] += 1
            return {}
            
    def verify_password(self, password: str, password_hash: str, salt: str) -> bool:
        """
        验证密码
        
        Args:
            password: 密码
            password_hash: 密码哈希值
            salt: 盐值
            
        Returns:
            密码是否正确
        """
        try:
            # 重新哈希密码
            new_hash = hashlib.pbkdf2_hmac('sha256', password.encode('utf-8'), salt.encode('utf-8'), 100000)
            return hmac.compare_digest(new_hash.hex(), password_hash)
            
        except Exception as e:
            print(f"✗ 密码验证失败: {e}")
            self.security_stats["errors"] += 1
            return False
            
    def _block_ip(self, client_ip: str, duration: float):
        """
        阻止IP地址
        
        Args:
            client_ip: 客户端IP
            duration: 阻止持续时间（秒）
        """
        try:
            self.security_state["blocked_ips"][client_ip] = {
                "block_time": time.time(),
                "duration": duration
            }
            
            # 触发IP阻止回调
            self._trigger_security_callback("ip_blocked", {
                "ip_address": client_ip,
                "duration": duration,
                "timestamp": time.time()
            })
            
            # 记录安全事件
            self._log_security_event("ip_blocked", None, client_ip)
            
            print(f"✓ IP地址已被阻止: {client_ip} (持续 {duration} 秒)")
            
        except Exception as e:
            print(f"✗ IP阻止失败: {e}")
            self.security_stats["errors"] += 1
            
    def unblock_ip(self, client_ip: str) -> bool:
        """
        解除IP阻止
        
        Args:
            client_ip: 客户端IP
            
        Returns:
            是否解除成功
        """
        try:
            if client_ip in self.security_state["blocked_ips"]:
                del self.security_state["blocked_ips"][client_ip]
                print(f"✓ IP地址阻止已解除: {client_ip}")
                return True
            else:
                print(f"✗ IP地址未被阻止: {client_ip}")
                return False
                
        except Exception as e:
            print(f"✗ IP阻止解除失败: {e}")
            self.security_stats["errors"] += 1
            return False
            
    def _log_security_event(self, event_type: str, client_id: str = None, client_address=None):
        """
        记录安全事件
        
        Args:
            event_type: 事件类型
            client_id: 客户端ID（可选）
            client_address: 客户端地址（可选）
        """
        try:
            if not self.security_config["enable_audit_logging"]:
                return
                
            event = {
                "event_type": event_type,
                "client_id": client_id,
                "client_address": client_address,
                "timestamp": time.time()
            }
            
            self.security_state["security_events"].append(event)
            self.security_stats["security_events"] += 1
            
            # 保持事件列表大小
            if len(self.security_state["security_events"]) > 1000:
                self.security_state["security_events"] = self.security_state["security_events"][-1000:]
                
            self.last_security_event = time.time()
            
        except Exception as e:
            print(f"✗ 安全事件记录失败: {e}")
            self.security_stats["errors"] += 1
            
    def get_active_sessions(self) -> Dict[str, Any]:
        """
        获取活跃会话
        
        Returns:
            活跃会话字典
        """
        try:
            return self.security_state["active_sessions"].copy()
        except Exception as e:
            print(f"✗ 获取活跃会话失败: {e}")
            self.security_stats["errors"] += 1
            return {}
            
    def get_blocked_ips(self) -> Dict[str, Any]:
        """
        获取被阻止的IP地址
        
        Returns:
            被阻止的IP地址字典
        """
        try:
            return self.security_state["blocked_ips"].copy()
        except Exception as e:
            print(f"✗ 获取被阻止的IP地址失败: {e}")
            self.security_stats["errors"] += 1
            return {}
            
    def get_security_stats(self) -> Dict[str, Any]:
        """
        获取安全统计信息
        
        Returns:
            安全统计字典
        """
        return {
            "state": self.security_state.copy(),
            "stats": self.security_stats.copy(),
            "config": self.security_config.copy()
        }
        
    def reset_stats(self):
        """重置安全统计信息"""
        try:
            self.security_stats = {
                "successful_authentications": 0,
                "failed_authentications": 0,
                "blocked_connections": 0,
                "security_events": 0,
                "active_sessions": 0,
                "errors": 0
            }
            print("✓ 安全统计信息已重置")
        except Exception as e:
            print(f"✗ 安全统计信息重置失败: {e}")
            
    def set_security_config(self, config: Dict[str, Any]) -> bool:
        """
        设置安全配置
        
        Args:
            config: 安全配置字典
            
        Returns:
            是否设置成功
        """
        try:
            self.security_config.update(config)
            print(f"✓ 安全配置已更新: {self.security_config}")
            return True
        except Exception as e:
            print(f"✗ 安全配置设置失败: {e}")
            return False
            
    def get_security_config(self) -> Dict[str, Any]:
        """
        获取安全配置
        
        Returns:
            安全配置字典
        """
        return self.security_config.copy()
        
    def _trigger_security_callback(self, callback_type: str, data: Dict[str, Any]):
        """
        触发安全回调
        
        Args:
            callback_type: 回调类型
            data: 回调数据
        """
        try:
            if callback_type in self.security_callbacks:
                for callback in self.security_callbacks[callback_type]:
                    try:
                        callback(data)
                    except Exception as e:
                        print(f"✗ 安全回调执行失败: {callback_type} - {e}")
        except Exception as e:
            print(f"✗ 安全回调触发失败: {e}")
            
    def register_security_callback(self, callback_type: str, callback: callable):
        """
        注册安全回调
        
        Args:
            callback_type: 回调类型
            callback: 回调函数
        """
        try:
            if callback_type in self.security_callbacks:
                self.security_callbacks[callback_type].append(callback)
                print(f"✓ 安全回调已注册: {callback_type}")
            else:
                print(f"✗ 无效的回调类型: {callback_type}")
        except Exception as e:
            print(f"✗ 安全回调注册失败: {e}")
            
    def unregister_security_callback(self, callback_type: str, callback: callable):
        """
        注销安全回调
        
        Args:
            callback_type: 回调类型
            callback: 回调函数
        """
        try:
            if callback_type in self.security_callbacks:
                if callback in self.security_callbacks[callback_type]:
                    self.security_callbacks[callback_type].remove(callback)
                    print(f"✓ 安全回调已注销: {callback_type}")
            else:
                print(f"✗ 无效的回调类型: {callback_type}")
        except Exception as e:
            print(f"✗ 安全回调注销失败: {e}")