# mandatory_model_crypto.py - 强制加密模型验证和解密模块
import os
import tempfile
import hashlib
import pickle
import traceback
import time
import secrets
import string
from pathlib import Path
import base64

import requests
import torch
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

from log import logger


class MandatoryModelValidator:
    """强制加密模型验证器 - 只负责验证和解密"""

    @staticmethod
    def verify_model_encryption(encrypted_path):
        """验证模型是否被正确加密"""
        try:
            if not os.path.exists(encrypted_path):
                return {'valid': False, 'error': '加密模型文件不存在'}

            with open(encrypted_path, 'rb') as f:
                data = pickle.load(f)

            # 检查必要的加密字段
            required_fields = ['salt', 'data', 'encrypted', 'model_hash', 'version']
            for field in required_fields:
                if field not in data:
                    return {'valid': False, 'error': f'缺少加密字段: {field}'}

            if data.get('encrypted', False) is not True:
                return {'valid': False, 'error': '模型未加密'}

            # 验证版本兼容性
            if data.get('version') not in ['1.0', '2.0']:
                return {'valid': False, 'error': f'不支持的加密版本: {data.get("version")}'}

            return {'valid': True, 'encrypted_data': data}

        except Exception as e:
            return {'valid': False, 'error': f'验证失败: {str(e)}'}

    @staticmethod
    def is_properly_encrypted(encrypted_path):
        """检查模型是否被正确加密（简化的验证）"""
        try:
            with open(encrypted_path, 'rb') as f:
                data = pickle.load(f)

            # 基本字段检查
            if not isinstance(data, dict):
                return False

            required_fields = ['salt', 'data', 'encrypted', 'model_hash']
            for field in required_fields:
                if field not in data:
                    return False

            return data.get('encrypted', False) is True

        except Exception as e:
            logger.warning(f"检查加密格式失败: {str(e)}")
            return False

    @staticmethod
    def decrypt_and_verify(encrypted_path, password, verify_key=True):
        """解密模型文件并验证完整性"""
        try:
            # 先验证加密格式
            verify_result = MandatoryModelValidator.verify_model_encryption(encrypted_path)
            if not verify_result['valid']:
                return {'success': False, 'error': verify_result['error']}

            encrypted_payload = verify_result['encrypted_data']
            salt = encrypted_payload['salt']
            encrypted_data = encrypted_payload['data']
            expected_hash = encrypted_payload['model_hash']

            # 生成解密密钥
            kdf = PBKDF2HMAC(
                algorithm=hashes.SHA256(),
                length=32,
                salt=salt,
                iterations=100000,
            )
            key = base64.urlsafe_b64encode(kdf.derive(password.encode()))

            # 解密数据
            fernet = Fernet(key)

            try:
                decrypted_data = fernet.decrypt(encrypted_data)
            except Exception as e:
                if "InvalidToken" in str(e) or "Invalid signature" in str(e):
                    return {'success': False, 'error': '解密密钥错误'}
                return {'success': False, 'error': f'解密失败: {str(e)}'}

            # 验证模型哈希
            actual_hash = hashlib.sha256(decrypted_data).hexdigest()
            if actual_hash != expected_hash:
                return {
                    'success': False,
                    'error': f'模型哈希不匹配: 期望{expected_hash[:16]}..., 实际{actual_hash[:16]}...'
                }

            # 验证文件大小
            original_size = encrypted_payload.get('original_size', 0)
            if original_size > 0 and len(decrypted_data) != original_size:
                return {
                    'success': False,
                    'error': f'文件大小不匹配: 期望{original_size}, 实际{len(decrypted_data)}'
                }

            logger.info(f"模型解密验证成功: {encrypted_path}")
            logger.debug(f"模型哈希: {actual_hash[:16]}...")

            return {
                'success': True,
                'model_hash': actual_hash,
                'original_size': len(decrypted_data),
                'decrypted_data': decrypted_data,  # 返回解密数据，供后续处理
                'version': encrypted_payload.get('version', '1.0')
            }

        except Exception as e:
            error_msg = str(e)
            if "InvalidToken" in error_msg or "Invalid signature" in error_msg:
                return {'success': False, 'error': '解密密钥错误'}
            return {'success': False, 'error': f'解密失败: {error_msg}'}

    @staticmethod
    def decrypt_model(encrypted_path, password, verify_key=False):
        """解密模型文件 - 兼容旧接口"""
        return MandatoryModelValidator.decrypt_and_verify(
            encrypted_path, password, verify_key=verify_key
        )


class ModelEncryptionService:
    """模型加密服务 - 处理密钥生成、验证和加密流程"""

    @staticmethod
    def generate_secure_key(length=32):
        """生成安全的加密密钥"""
        # 使用加密安全的随机数生成器
        alphabet = string.ascii_letters + string.digits + "!@#$%^&*"
        key = ''.join(secrets.choice(alphabet) for _ in range(length))

        # 计算密钥哈希
        key_hash = hashlib.sha256(key.encode()).hexdigest()

        return {
            'key': key,
            'key_hash': key_hash,
            'short_hash': key_hash[:16]
        }

    @staticmethod
    def validate_key_strength(key):
        """验证密钥强度"""
        if len(key) < 16:
            return False, "密钥长度至少16位"

        # 检查字符种类
        has_upper = any(c.isupper() for c in key)
        has_lower = any(c.islower() for c in key)
        has_digit = any(c.isdigit() for c in key)
        has_special = any(c in "!@#$%^&*" for c in key)

        if not (has_upper and has_lower and has_digit and has_special):
            return False, "密钥应包含大小写字母、数字和特殊字符"

        return True, "密钥强度足够"

    @staticmethod
    def create_key_pair():
        """创建密钥对（用于客户端-服务器通信）"""
        # 生成主密钥
        master_key = secrets.token_urlsafe(32)

        # 生成验证令牌
        verification_token = hashlib.sha256(
            (master_key + str(time.time())).encode()
        ).hexdigest()[:20]

        return {
            'master_key': master_key,
            'verification_token': verification_token,
            'created_at': time.time()
        }


class PreTaskModelValidator:
    """任务前模型验证器 - 在创建任务前验证所有模型"""

    def __init__(self, config=None):
        self.config = config or {}
        self.encrypted_models_dir = self.config.get('encrypted_models_dir', 'encrypted_models')
        os.makedirs(self.encrypted_models_dir, exist_ok=True)

        # 验证状态缓存
        self.validation_cache = {}
        self.cache_lock = threading.Lock()
        self.cache_expiry = 300  # 5分钟缓存有效期

    def validate_models_before_task(self, task_config):
        """在创建任务前验证所有模型和密钥"""
        try:
            models_config = task_config.get('models', [])
            if not models_config:
                return {
                    'success': False,
                    'error': '任务配置中未找到模型列表'
                }

            validation_results = []
            all_valid = True

            for i, model_config in enumerate(models_config):
                # 检查必要的配置项
                required_fields = ['path', 'encryption_key']
                for field in required_fields:
                    if field not in model_config:
                        return {
                            'success': False,
                            'error': f'模型 {i} 缺少必要字段: {field}'
                        }

                model_path = model_config['path']
                encryption_key = model_config['encryption_key']

                # 构建完整的模型文件路径
                full_model_path = self._get_full_model_path(model_path)

                # 验证模型文件是否存在
                if not os.path.exists(full_model_path):
                    return {
                        'success': False,
                        'error': f'模型文件不存在: {full_model_path}'
                    }

                # 验证密钥强度
                key_valid, key_msg = ModelEncryptionService.validate_key_strength(encryption_key)
                if not key_valid:
                    return {
                        'success': False,
                        'error': f'模型 {i} 密钥强度不足: {key_msg}'
                    }

                # 验证密钥是否正确（尝试解密）
                validator = MandatoryModelValidator()
                decrypt_result = validator.decrypt_and_verify(full_model_path, encryption_key)

                validation_result = {
                    'model_index': i,
                    'model_path': model_path,
                    'full_path': full_model_path,
                    'file_exists': True,
                    'key_valid': decrypt_result['success'],
                    'model_hash': decrypt_result.get('model_hash', '')[:16] if decrypt_result['success'] else '',
                    'model_size': decrypt_result.get('original_size', 0) if decrypt_result['success'] else 0,
                    'validation_time': time.time()
                }

                if not decrypt_result['success']:
                    validation_result['error'] = decrypt_result.get('error', '验证失败')
                    all_valid = False

                validation_results.append(validation_result)

            return {
                'success': all_valid,
                'valid': all_valid,
                'total_models': len(models_config),
                'valid_models': sum(1 for r in validation_results if r['key_valid']),
                'validation_results': validation_results,
                'timestamp': time.time()
            }

        except Exception as e:
            logger.error(f"预验证模型失败: {str(e)}")
            return {
                'success': False,
                'error': f'预验证失败: {str(e)}'
            }

    def verify_single_model(self, model_path, encryption_key):
        """验证单个模型"""
        try:
            # 构建完整路径
            full_model_path = self._get_full_model_path(model_path)

            # 检查文件是否存在
            if not os.path.exists(full_model_path):
                return {
                    'success': False,
                    'error': f'模型文件不存在: {full_model_path}'
                }

            # 验证密钥强度
            key_valid, key_msg = ModelEncryptionService.validate_key_strength(encryption_key)
            if not key_valid:
                return {
                    'success': False,
                    'error': f'密钥强度不足: {key_msg}'
                }

            # 解密验证
            validator = MandatoryModelValidator()
            decrypt_result = validator.decrypt_and_verify(full_model_path, encryption_key)

            if decrypt_result['success']:
                return {
                    'success': True,
                    'valid': True,
                    'model_hash': decrypt_result.get('model_hash', '')[:16],
                    'model_size': decrypt_result.get('original_size', 0),
                    'file_path': full_model_path,
                    'timestamp': time.time()
                }
            else:
                return {
                    'success': False,
                    'valid': False,
                    'error': decrypt_result.get('error', '验证失败'),
                    'timestamp': time.time()
                }

        except Exception as e:
            logger.error(f"验证单个模型失败: {str(e)}")
            return {
                'success': False,
                'error': f'验证失败: {str(e)}'
            }

    def _get_full_model_path(self, model_path):
        """获取完整的模型文件路径"""
        if os.path.isabs(model_path):
            return model_path

        # 如果路径以 encrypted_models/ 开头
        if model_path.startswith('encrypted_models/'):
            model_filename = os.path.basename(model_path)
            return os.path.join(self.encrypted_models_dir, model_filename)

        # 否则直接使用文件名
        model_filename = os.path.basename(model_path)
        return os.path.join(self.encrypted_models_dir, model_filename)

    def get_model_info(self, model_path):
        """获取模型信息（不验证密钥）"""
        try:
            full_model_path = self._get_full_model_path(model_path)

            if not os.path.exists(full_model_path):
                return None

            verify_result = MandatoryModelValidator().verify_model_encryption(full_model_path)

            if verify_result['valid']:
                data = verify_result['encrypted_data']
                return {
                    'encrypted': True,
                    'model_hash': data.get('model_hash', '')[:16],
                    'original_size': data.get('original_size', 0),
                    'version': data.get('version', 'unknown'),
                    'file_path': full_model_path,
                    'file_size': os.path.getsize(full_model_path)
                }
            return None

        except Exception as e:
            logger.error(f"获取模型信息失败: {str(e)}")
            return None


class SecureModelManager:
    """安全模型管理器 - 负责下载、验证、加载加密模型"""

    def __init__(self, config):
        self.config = config
        self.models_dir = config.get('models_dir', 'models')
        self.encrypted_models_dir = config.get('encrypted_models_dir', 'encrypted_models')
        os.makedirs(self.models_dir, exist_ok=True)
        os.makedirs(self.encrypted_models_dir, exist_ok=True)

        # 验证器实例
        self.validator = MandatoryModelValidator()
        self.pre_validator = PreTaskModelValidator(config)

        # 模型缓存
        self.model_cache = {}
        self.cache_lock = threading.Lock()
        self.cache_expiry = 600  # 10分钟缓存

        # 验证状态记录
        self.verification_status = {}
        self.verification_expiry = 300  # 5分钟验证缓存

    def ensure_model_available(self, model_config):
        """确保模型可用：验证模型文件存在"""
        try:
            model_path = model_config.get('path')
            if not model_path:
                return {'available': False, 'error': '模型路径未配置'}

            # 获取模型本地路径
            local_path = self._get_local_model_path(model_path, model_config)

            # 检查模型文件是否存在
            if not os.path.exists(local_path):
                return {'available': False, 'error': f'模型文件不存在: {local_path}'}

            # 验证加密格式
            verify_result = self.validator.verify_model_encryption(local_path)
            if not verify_result['valid']:
                return {'available': False, 'error': f'模型加密格式无效: {verify_result.get("error", "未知错误")}'}

            return {
                'available': True,
                'local_path': local_path,
                'model_hash': verify_result['encrypted_data'].get('model_hash', '')[:16]
            }

        except Exception as e:
            return {'available': False, 'error': f'确保模型可用失败: {str(e)}'}

    def verify_model_key(self, model_config, encryption_key):
        """验证模型密钥 - 这是创建任务的前置条件"""
        try:
            # 确保模型文件存在
            availability = self.ensure_model_available(model_config)
            if not availability['available']:
                return {'valid': False, 'error': availability['error']}

            local_path = availability['local_path']
            model_hash_short = availability.get('model_hash', 'unknown')

            # 检查是否已验证过（避免重复验证）
            cache_key = f"{local_path}_{hashlib.md5(encryption_key.encode()).hexdigest()[:8]}"
            with self.cache_lock:
                if cache_key in self.verification_status:
                    cached_result = self.verification_status[cache_key]
                    if time.time() - cached_result.get('verified_at', 0) < self.verification_expiry:
                        if cached_result['valid']:
                            logger.info(f"使用缓存的验证结果: {model_hash_short}")
                            return cached_result
                    else:
                        # 缓存过期，删除
                        del self.verification_status[cache_key]

            # 解密验证
            decrypt_result = self.validator.decrypt_and_verify(local_path, encryption_key)

            if decrypt_result['success']:
                result = {
                    'valid': True,
                    'model_hash': decrypt_result.get('model_hash', '')[:16],
                    'original_size': decrypt_result.get('original_size', 0),
                    'local_path': local_path,
                    'verified_at': time.time()
                }

                # 缓存验证结果
                with self.cache_lock:
                    self.verification_status[cache_key] = result

                logger.info(f"模型密钥验证成功: {model_hash_short}")
                return result
            else:
                result = {'valid': False, 'error': decrypt_result.get('error', '验证失败')}
                logger.warning(f"模型密钥验证失败: {model_hash_short} - {result['error']}")
                return result

        except Exception as e:
            error_msg = f"验证模型密钥失败: {str(e)}"
            logger.error(error_msg)
            return {'valid': False, 'error': error_msg}

    def load_verified_model(self, model_config, encryption_key):
        """加载已验证的模型 - 仅在验证通过后调用"""
        try:
            # 先验证密钥
            verify_result = self.verify_model_key(model_config, encryption_key)
            if not verify_result['valid']:
                raise ValueError(f"模型验证未通过: {verify_result.get('error', '未知错误')}")

            local_path = verify_result['local_path']

            # 检查缓存
            cache_key = f"{local_path}_{hashlib.md5(encryption_key.encode()).hexdigest()[:8]}"
            with self.cache_lock:
                if cache_key in self.model_cache:
                    cached_info = self.model_cache[cache_key]
                    if time.time() - cached_info.get('loaded_at', 0) < self.cache_expiry:
                        logger.info(f"使用缓存的模型: {local_path}")
                        return cached_info['model']

            # 解密模型（使用已验证的密钥）
            decrypt_result = self.validator.decrypt_and_verify(local_path, encryption_key)
            if not decrypt_result['success']:
                raise ValueError(f"模型解密失败: {decrypt_result.get('error', '未知错误')}")

            # 保存解密数据到临时文件
            with tempfile.NamedTemporaryFile(delete=False, suffix='.pt') as tmp:
                tmp.write(decrypt_result['decrypted_data'])
                temp_path = tmp.name

            # 加载YOLO模型
            from ultralytics import YOLO
            model = YOLO(temp_path)

            # 应用设备配置
            device = model_config.get('device', 'cuda:0' if torch.cuda.is_available() else 'cpu')
            model = model.to(device)

            # 应用半精度配置
            if model_config.get('half', False) and 'cuda' in device:
                model = model.half()
                logger.info(f"启用半精度推理: {model_config.get('path', '未知模型')}")

            # 清理临时文件
            try:
                os.unlink(temp_path)
            except:
                pass

            # 缓存模型
            with self.cache_lock:
                self.model_cache[cache_key] = {
                    'model': model,
                    'device': device,
                    'loaded_at': time.time(),
                    'original_size': decrypt_result.get('original_size', 0)
                }

            logger.info(f"加密模型加载成功: {model_config.get('path', '未知模型')} -> {device}")

            return model

        except Exception as e:
            logger.error(f"加载已验证模型失败: {str(e)}")
            logger.error(traceback.format_exc())
            return None

    def _get_local_model_path(self, model_path, model_config):
        """获取模型本地路径"""
        # 如果已经是本地路径，直接返回
        if os.path.exists(model_path):
            return model_path

        # 检查是否是加密模型
        encrypted = model_config.get('encrypted', False)

        if encrypted:
            # 从加密模型目录查找
            model_filename = os.path.basename(model_path)
            if not model_filename.endswith('.enc'):
                model_filename += '.enc'
            return os.path.join(self.encrypted_models_dir, model_filename)
        else:
            # 从普通模型目录查找
            model_filename = os.path.basename(model_path)
            return os.path.join(self.models_dir, model_filename)

    def get_model_info(self, model_config):
        """获取模型信息（不验证密钥）"""
        try:
            availability = self.ensure_model_available(model_config)
            if not availability['available']:
                return None

            local_path = availability['local_path']
            verify_result = self.validator.verify_model_encryption(local_path)

            if verify_result['valid']:
                data = verify_result['encrypted_data']
                return {
                    'encrypted': True,
                    'model_hash': data.get('model_hash', '')[:16],
                    'original_size': data.get('original_size', 0),
                    'version': data.get('version', 'unknown'),
                    'local_path': local_path,
                    'file_size': os.path.getsize(local_path)
                }
            return None

        except Exception as e:
            logger.error(f"获取模型信息失败: {str(e)}")
            return None

    def clear_cache(self):
        """清空缓存"""
        with self.cache_lock:
            self.model_cache.clear()
            self.verification_status.clear()
            logger.info("模型缓存已清空")

    def get_cache_info(self):
        """获取缓存信息"""
        with self.cache_lock:
            return {
                'model_cache_size': len(self.model_cache),
                'verification_cache_size': len(self.verification_status),
                'model_cache_keys': list(self.model_cache.keys()),
                'verification_cache_keys': list(self.verification_status.keys())
            }


class SimpleModelEncryptor:
    """简化模型加密器 - 仅用于测试或本地转换"""

    @staticmethod
    def create_encrypted_payload(model_data, password):
        """创建加密数据包（模拟.NET服务器的加密过程）"""
        # 生成盐和密钥
        salt = os.urandom(16)
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        fernet = Fernet(key)

        # 加密数据
        encrypted_data = fernet.encrypt(model_data)
        model_hash = hashlib.sha256(model_data).hexdigest()

        # 创建加密数据包
        encrypted_payload = {
            'salt': salt,
            'data': encrypted_data,
            'model_hash': model_hash,
            'original_size': len(model_data),
            'encrypted': True,
            'version': '2.0',
            'created_at': time.time()
        }

        return encrypted_payload, key

    @staticmethod
    def encrypt_model_file(input_path, output_path, password):
        """加密模型文件"""
        try:
            # 读取原始模型文件
            with open(input_path, 'rb') as f:
                model_data = f.read()

            # 创建加密数据包
            encrypted_payload, key = SimpleModelEncryptor.create_encrypted_payload(model_data, password)

            # 保存加密文件
            with open(output_path, 'wb') as f:
                pickle.dump(encrypted_payload, f)

            logger.info(f"模型加密成功: {input_path} -> {output_path}")

            return {
                'success': True,
                'output_path': output_path,
                'model_hash': encrypted_payload['model_hash'][:16],
                'original_size': len(model_data),
                'encrypted_size': os.path.getsize(output_path)
            }

        except Exception as e:
            logger.error(f"加密模型文件失败: {str(e)}")
            return {'success': False, 'error': str(e)}

    @staticmethod
    def verify_encryption(encrypted_path, password):
        """验证加密文件"""
        return MandatoryModelValidator().decrypt_and_verify(encrypted_path, password)


# ==================== 全局实例和接口函数 ====================

# 导入线程模块
import threading

# 全局实例
_model_manager = None
_pre_task_validator = None
_upload_manager = None


def get_model_manager(config=None):
    """获取模型管理器单例"""
    global _model_manager
    if _model_manager is None:
        if config is None:
            from config import get_default_config
            config = get_default_config()
        _model_manager = SecureModelManager(config)
    return _model_manager


def get_pre_task_validator(config=None):
    """获取预验证器单例"""
    global _pre_task_validator
    if _pre_task_validator is None:
        if config is None:
            from config import get_default_config
            config = get_default_config()
        _pre_task_validator = PreTaskModelValidator(config.get('upload', {}))
    return _pre_task_validator


def validate_models_before_task(task_config):
    """验证任务的所有模型（外部调用接口）"""
    validator = get_pre_task_validator()
    return validator.validate_models_before_task(task_config)


def verify_single_model_api(model_path, encryption_key):
    """验证单个模型（API接口）"""
    validator = get_pre_task_validator()
    return validator.verify_single_model(model_path, encryption_key)


def verify_model_key_api(model_config, encryption_key):
    """验证模型密钥（API接口）"""
    manager = get_model_manager()
    return manager.verify_model_key(model_config, encryption_key)


def load_verified_model_api(model_config, encryption_key):
    """加载已验证的模型（API接口）"""
    manager = get_model_manager()
    return manager.load_verified_model(model_config, encryption_key)


def generate_secure_key_api(length=32):
    """生成安全密钥（API接口）"""
    return ModelEncryptionService.generate_secure_key(length)


def validate_key_strength_api(key):
    """验证密钥强度（API接口）"""
    return ModelEncryptionService.validate_key_strength(key)


def test_decryption_api(encrypted_path, password):
    """测试解密（API接口）"""
    validator = MandatoryModelValidator()
    return validator.decrypt_and_verify(encrypted_path, password)


# ==================== 向后兼容的接口 ====================

class MandatoryModelEncryptor:
    """向后兼容的加密器类"""

    @staticmethod
    def encrypt_model(input_path, output_path, password, require_encryption=True):
        """加密模型 - 向后兼容接口"""
        return SimpleModelEncryptor.encrypt_model_file(input_path, output_path, password)

    @staticmethod
    def decrypt_model(encrypted_path, password, verify_key=False):
        """解密模型 - 向后兼容接口"""
        validator = MandatoryModelValidator()
        return validator.decrypt_model(encrypted_path, password, verify_key=verify_key)

    @staticmethod
    def is_properly_encrypted(encrypted_path):
        """检查是否被正确加密 - 向后兼容接口"""
        return MandatoryModelValidator.is_properly_encrypted(encrypted_path)

    @staticmethod
    def generate_secure_key():
        """生成安全密钥 - 向后兼容接口"""
        return ModelEncryptionService.generate_secure_key()

    @staticmethod
    def verify_model_key(encrypted_path, password):
        """验证模型密钥 - 向后兼容接口"""
        validator = MandatoryModelValidator()
        return validator.decrypt_and_verify(encrypted_path, password)


# ==================== 测试代码 ====================

if __name__ == "__main__":
    import time

    print("=== 强制加密模型验证模块测试 ===")

    # 配置示例
    config = {
        'models_dir': './models',
        'encrypted_models_dir': './encrypted_models',
        'cache_verification': True
    }

    # 模型配置示例
    model_config = {
        'path': 'yolov8n.enc',  # 加密模型文件路径
        'device': 'cpu',
        'half': False,
        'encrypted': True,
        'encryption_key': 'test-password-123!@#'  # 测试密钥
    }

    # 测试密钥生成
    print("\n1. 测试密钥生成:")
    key_info = ModelEncryptionService.generate_secure_key()
    print(f"   生成的密钥: {key_info['key']}")
    print(f"   密钥哈希: {key_info['key_hash']}")
    print(f"   短哈希: {key_info['short_hash']}")

    # 测试密钥强度验证
    print("\n2. 测试密钥强度验证:")
    test_key = "WeakKey123"
    valid, msg = ModelEncryptionService.validate_key_strength(test_key)
    print(f"   密钥 '{test_key}': {valid} - {msg}")

    test_key = "Strong@Password#2024!Complex"
    valid, msg = ModelEncryptionService.validate_key_strength(test_key)
    print(f"   密钥 '{test_key}': {valid} - {msg}")

    # 测试模型验证器
    print("\n3. 测试模型验证器:")
    validator = MandatoryModelValidator()

    # 测试不存在的文件
    result = validator.verify_model_encryption("nonexistent.enc")
    print(f"   不存在的文件验证: {result['valid']} - {result.get('error', '')}")

    # 测试预验证器
    print("\n4. 测试预验证器:")
    pre_validator = PreTaskModelValidator(config)

    task_config = {
        'models': [
            {
                'path': 'test_model.enc',
                'encryption_key': 'Test@Password#2024!Complex'
            }
        ]
    }

    result = pre_validator.validate_models_before_task(task_config)
    print(f"   任务验证结果: 成功={result['success']}, 错误={result.get('error', '无')}")

    # 测试安全模型管理器
    print("\n5. 测试安全模型管理器:")
    manager = SecureModelManager(config)

    # 测试模型信息获取
    model_info = manager.get_model_info(model_config)
    if model_info:
        print(f"   模型信息: 哈希={model_info['model_hash']}, 大小={model_info['original_size']}")
    else:
        print(f"   模型信息: 不可用")

    print("\n=== 测试完成 ===")