Exeprotector/encryptor_secure.py

"""
安全的EXE加密器（V2改进版）
- 使用持久化缓存的验证器
- 激活一次，长期使用
- 到期提醒功能
- 智能验证策略

作者：太一
"""

import os
import hashlib
import json
import tempfile
import subprocess
import shutil
import uuid
from typing import Tuple
from datetime import datetime
from cryptography.fernet import Fernet
import base64

# 加密文件的特殊标识（嵌入在验证器中）
ENCRYPTION_MARKER = "SECURE_EXE_VALIDATOR_V2_TAIYI"

class SecureEXEEncryptor:
    """安全的EXE加密器 V2（改进版）- 支持版本管理"""
    
    def __init__(self):
        self.encryption_key = None
        self.metadata_file = "files/file_metadata.json"
    
    def generate_encryption_key(self) -> bytes:
        """生成加密密钥"""
        return Fernet.generate_key()
    
    def _is_encrypted_file(self, file_path: str) -> bool:
        """检测文件是否已经被加密过"""
        try:
            # 方法1: 检查文件中是否包含加密标识
            with open(file_path, 'rb') as f:
                content = f.read()
                # 检查是否包含验证器的特殊标记
                if ENCRYPTION_MARKER.encode() in content:
                    return True
                # 检查是否包含验证器代码特征
                if b"ENCRYPTED_RESOURCE" in content and b"smart_validate" in content:
                    return True
            return False
        except Exception as e:
            print(f"检测文件时出错: {e}")
            return False
    
    def _load_metadata(self) -> dict:
        """加载元数据文件"""
        try:
            if os.path.exists(self.metadata_file):
                with open(self.metadata_file, 'r', encoding='utf-8') as f:
                    return json.load(f)
        except Exception as e:
            print(f"加载元数据失败: {e}")
        return {}
    
    def _save_metadata(self, metadata: dict):
        """保存元数据文件"""
        try:
            os.makedirs(os.path.dirname(self.metadata_file), exist_ok=True)
            with open(self.metadata_file, 'w', encoding='utf-8') as f:
                json.dump(metadata, f, indent=2, ensure_ascii=False)
            return True
        except Exception as e:
            print(f"保存元数据失败: {e}")
            return False
    
    def _get_or_create_software_id(self, software_name: str, software_id: str = None) -> str:
        """获取或创建软件ID"""
        metadata = self._load_metadata()
        
        # 如果提供了软件ID，直接使用
        if software_id:
            return software_id
        
        # 查找是否已经有该软件名的记录
        for key, info in metadata.items():
            if info.get('software_name') == software_name and info.get('software_id'):
                print(f"✅ 找到已存在的软件ID: {info['software_id']}")
                return info['software_id']
        
        # 生成新的软件ID
        new_id = str(uuid.uuid4())
        print(f"🆕 生成新的软件ID: {new_id}")
        return new_id
    
    def _save_encryption_metadata(self, output_path: str, source_path: str, 
                                  software_name: str, software_id: str, 
                                  version: str, file_hash: str):
        """保存加密元数据"""
        try:
            metadata = self._load_metadata()
            
            # 生成唯一键
            file_key = str(uuid.uuid4())
            
            # 保存元数据
            metadata[file_key] = {
                'original_name': os.path.basename(source_path),
                'encrypted_path': output_path,
                'software_name': software_name,
                'software_id': software_id,
                'version': version,
                'hash': file_hash,
                'size': os.path.getsize(output_path),
                'encryption_time': str(datetime.now()),
                'source_path': source_path
            }
            
            # 保存到文件
            if self._save_metadata(metadata):
                print(f"✅ 元数据已保存")
            else:
                print(f"⚠️  元数据保存失败")
                
        except Exception as e:
            print(f"保存元数据时出错: {e}")
    
    def encrypt_exe(self, source_path: str, output_path: str, 
                   api_config: dict, software_name: str,
                   use_v2_validator: bool = True, 
                   software_id: str = None,
                   version: str = "1.0.0") -> Tuple[bool, str]:
        """
        加密EXE文件（使用V2改进版验证器）
        
        参数:
            source_path: 原始EXE路径
            output_path: 输出路径
            api_config: API配置（不包含数据库密码）
            software_name: 软件名称
            use_v2_validator: 保留参数（兼容性），现在统一使用V2验证器
            software_id: 软件唯一ID（用于版本管理），为None则自动生成
            version: 软件版本号，默认1.0.0
        """
        try:
            print(f"开始加密: {source_path}")
            
            # 1. 检测是否已经是加密文件
            if self._is_encrypted_file(source_path):
                return False, "❌ 检测到这是一个已加密的文件！\n\n为避免多层验证，请使用原始未加密的EXE文件。\n\n如需更新软件版本：\n1. 使用原始新版本EXE文件\n2. 在加密时使用相同的软件ID\n3. 这样用户的激活码可以继续使用"
            
            # 2. 读取原始EXE
            with open(source_path, 'rb') as f:
                original_content = f.read()
            
            print(f"原始文件大小: {len(original_content)} 字节")
            
            # 3. 获取或创建软件ID
            actual_software_id = self._get_or_create_software_id(software_name, software_id)
            print(f"软件ID: {actual_software_id}")
            print(f"软件版本: {version}")
            
            # 4. 生成加密密钥
            self.encryption_key = self.generate_encryption_key()
            
            # 5. 加密EXE内容
            fernet = Fernet(self.encryption_key)
            encrypted_content = fernet.encrypt(original_content)
            
            print(f"加密后大小: {len(encrypted_content)} 字节")
            
            # 6. 计算哈希
            file_hash = hashlib.sha256(original_content).hexdigest()
            
            # 7. 准备资源数据
            key_part1 = base64.b64encode(self.encryption_key[:16]).decode()
            key_part2 = base64.b64encode(self.encryption_key[16:]).decode()
            
            resource_data = {
                'data': base64.b64encode(encrypted_content).decode(),
                'hash': file_hash,
                'key_hint': key_part1,
                'key_part2': key_part2,
                'software_id': actual_software_id,
                'version': version,
            }
            
            # 6. 选择验证器模板（现在统一使用V2版本）
            validator_template_path = "validator_secure.py"
            
            if not os.path.exists(validator_template_path):
                return False, f"找不到验证器模板: {validator_template_path}"
            
            print(f"使用验证器模板: {validator_template_path}")
            
            with open(validator_template_path, 'r', encoding='utf-8') as f:
                template_code = f.read()
            
            # 7. 替换占位符
            final_code = template_code.replace(
                '# ENCRYPTED_RESOURCE_PLACEHOLDER\nENCRYPTED_RESOURCE = None',
                f'ENCRYPTED_RESOURCE = {json.dumps(resource_data)}'
            )
            
            # 替换API配置
            if 'api_url' in api_config:
                final_code = final_code.replace(
                    'API_BASE_URL = "http://localhost:5100/"',
                    f'API_BASE_URL = "{api_config["api_url"]}"'
                )
            
            if 'api_key' in api_config:
                final_code = final_code.replace(
                    'API_KEY = "taiyi1224taiyi1224taiyi1224taiyi1224taiyi1224taiyi1224taiyi1224"',
                    f'API_KEY = "{api_config["api_key"]}"'
                )
            
            # 替换软件名称（保留用于显示）
            final_code = final_code.replace(
                'SOFTWARE_NAME = "SOFTWARE_NAME_PLACEHOLDER"',
                f'SOFTWARE_NAME = "{software_name}"'
            )
            
            # 替换软件ID（用于验证）
            final_code = final_code.replace(
                'SOFTWARE_ID = "SOFTWARE_ID_PLACEHOLDER"',
                f'SOFTWARE_ID = "{actual_software_id}"'
            )
            
            # 8. 写入临时Python文件
            temp_py = tempfile.mktemp(suffix='.py')
            with open(temp_py, 'w', encoding='utf-8') as f:
                f.write(final_code)
            
            print(f"生成的Python文件: {temp_py}")
            
            # 9. 编译为EXE
            success, msg = self._compile_to_exe(temp_py, output_path)
            
            if success:
                # 保存密钥（用于测试和恢复）
                key_file = output_path + '.key'
                with open(key_file, 'wb') as f:
                    f.write(self.encryption_key)
                print(f"⚠️  密钥已保存到: {key_file}")
                print(f"⚠️  请妥善保管此文件，建议加密存储！")
                
                # 保存元数据
                self._save_encryption_metadata(
                    output_path=output_path,
                    source_path=source_path,
                    software_name=software_name,
                    software_id=actual_software_id,
                    version=version,
                    file_hash=file_hash
                )
                
                # 显示V2验证器功能说明
                print("\n" + "="*60)
                print("✅ 使用V2验证器（支持版本管理），包含以下功能：")
                print("  1. 持久化缓存存储（激活一次，长期使用）")
                print("  2. 智能验证策略（减少重复验证）")
                print("  3. 到期提醒功能（提前7天开始提醒）")
                print("  4. 增强的离线验证（直到许可证过期）")
                print("  5. 改善的用户界面（显示剩余天数）")
                print("  6. 软件版本管理（更新软件不影响激活状态）")
                print("="*60)
                print(f"\n📋 软件信息：")
                print(f"  软件名称: {software_name}")
                print(f"  软件ID: {actual_software_id}")
                print(f"  版本号: {version}")
                print(f"  文件哈希: {file_hash[:16]}...")
                print("\n💡 版本管理说明：")
                print(f"  • 软件ID是唯一标识，用于版本管理")
                print(f"  • 更新软件时请使用相同的软件ID")
                print(f"  • 这样用户的激活码可以继续在新版本上使用")
                print("="*60 + "\n")
            
            # 清理临时文件
            try:
                os.remove(temp_py)
            except:
                pass
            
            return success, msg
            
        except Exception as e:
            import traceback
            error_detail = traceback.format_exc()
            print(f"加密过程出错: {error_detail}")
            return False, f"加密过程出错: {str(e)}"
    
    def _compile_to_exe(self, python_file: str, output_path: str) -> Tuple[bool, str]:
        """编译Python文件为EXE"""
        build_dir = None
        try:
            # 1. 检查PyInstaller
            print("检查 PyInstaller...")
            result = subprocess.run(['pyinstaller', '--version'],
                                  capture_output=True, text=True, timeout=10)
            if result.returncode != 0:
                return False, "PyInstaller未安装或不可用。请运行: pip install pyinstaller"
            
            pyinstaller_version = result.stdout.strip()
            print(f"PyInstaller 版本: {pyinstaller_version}")
            
            # 2. 创建工作目录
            build_dir = os.path.join(os.path.dirname(os.path.abspath(output_path)), '.build_temp')
            os.makedirs(build_dir, exist_ok=True)
            print(f"工作目录: {build_dir}")
            
            # 3. 复制Python文件到工作目录
            py_name = os.path.basename(python_file)
            work_py = os.path.join(build_dir, py_name)
            shutil.copy2(python_file, work_py)
            print(f"工作文件: {work_py}")
            
            # 4. 构建PyInstaller命令
            exe_name = os.path.splitext(os.path.basename(output_path))[0]
            
            cmd = [
                'pyinstaller',
                '--clean',
                '--noconfirm',
                '--onefile',
                '--windowed',
                '--name', exe_name,
                '--distpath', build_dir,
                '--workpath', os.path.join(build_dir, 'build'),
                '--specpath', build_dir,
                # 添加隐藏导入
                '--hidden-import', 'io',
                '--hidden-import', 'cryptography.fernet',
                '--hidden-import', 'cryptography.hazmat.primitives',
                '--hidden-import', 'cryptography.hazmat.backends',
                '--hidden-import', 'requests',
                '--hidden-import', 'urllib3',
                '--hidden-import', 'tkinter',
                '--hidden-import', 'tkinter.messagebox',
                '--hidden-import', '_tkinter',
                '--hidden-import', 'sqlite3',
                # 收集所有cryptography包
                '--collect-all', 'cryptography',
                # 禁用UPX
                '--noupx',
                work_py
            ]
            
            print("开始编译...")
            print(f"命令: {' '.join(cmd)}")
            
            # 5. 执行编译
            result = subprocess.run(
                cmd, 
                capture_output=True, 
                text=True, 
                timeout=600,
                cwd=build_dir,
                encoding='utf-8',
                errors='replace'
            )
            
            # 6. 输出详细日志
            if result.stdout:
                print("=== 编译输出 ===")
                print(result.stdout[-2000:])
                
            if result.stderr:
                print("=== 编译错误 ===")
                print(result.stderr[-2000:])
            
            # 7. 检查编译结果
            if result.returncode == 0:
                compiled_exe = os.path.join(build_dir, f"{exe_name}.exe")
                print(f"查找编译文件: {compiled_exe}")
                
                if os.path.exists(compiled_exe):
                    # 移动到目标位置
                    os.makedirs(os.path.dirname(output_path), exist_ok=True)
                    if os.path.exists(output_path):
                        os.remove(output_path)
                    shutil.move(compiled_exe, output_path)
                    
                    print(f"✅ 编译成功: {output_path}")
                    print(f"文件大小: {os.path.getsize(output_path) / 1024 / 1024:.2f} MB")
                    
                    # 清理工作目录
                    try:
                        shutil.rmtree(build_dir)
                        print("已清理临时文件")
                    except:
                        pass
                    
                    return True, "编译成功"
                else:
                    return False, f"找不到编译输出文件: {compiled_exe}"
            else:
                # 提取关键错误信息
                error_lines = []
                if result.stderr:
                    error_lines = [line for line in result.stderr.split('\n') 
                                 if 'error' in line.lower() or 'failed' in line.lower()]
                
                error_msg = '\n'.join(error_lines[-10:]) if error_lines else result.stderr[-1000:]
                
                return False, f"编译失败 (返回码: {result.returncode})\n{error_msg}"
                
        except subprocess.TimeoutExpired:
            return False, "编译超时（超过10分钟）。请检查系统资源或Python依赖。"
        except Exception as e:
            import traceback
            error_detail = traceback.format_exc()
            print(f"编译异常: {error_detail}")
            return False, f"编译过程出错: {str(e)}\n\n详细信息:\n{error_detail[-500:]}"
        finally:
            # 清理临时文件（如果编译失败）
            if build_dir and os.path.exists(build_dir):
                try:
                    if not os.path.exists(output_path):
                        print(f"⚠️  编译失败，工作目录已保留用于调试: {build_dir}")
                except:
                    pass


# ========== 测试代码 ==========
if __name__ == '__main__':
    print("EXE加密器 (V2改进版) - 测试")
    print("=" * 60)
    
    # 测试加密
    encryptor = SecureEXEEncryptor()
    
    # 配置
    api_config = {
        'api_url': 'http://localhost:5100/',
        'api_key': 'taiyi1224taiyi1224taiyi1224taiyi1224taiyi1224taiyi1224taiyi1224'
    }
    
    # 示例：加密一个测试EXE
    # success, msg = encryptor.encrypt_exe(
    #     source_path='test.exe',
    #     output_path='test_encrypted.exe',
    #     api_config=api_config,
    #     software_name='TestApp'
    # )
    
    # print(f"\n结果: {success}")
    # print(f"消息: {msg}")
    
    print("\n使用说明：")
    print("1. 在main.py中使用此加密器")
    print("2. 现在默认使用V2验证器（改进版）")
    print("3. V2功能：持久化缓存、到期提醒、智能验证策略")