"""
PyArmor + AES 双层保护完整示例

本示例演示如何同时使用 PyArmor 混淆和 AES 字节码加密
来提供双重保护。
"""

import sys
import os
import tempfile
import shutil

# 添加项目源码路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

# ========== 1. 创建一个需要保护的模块 ==========

SECRET_MODULE_CODE = """
# 这是需要保护的敏感代码（支付逻辑、授权验证等）
import hashlib
import hmac

SECRET_KEY = "this-is-a-top-secret-key"

def verify_license(license_key: str) -> bool:
    \"\"\"验证许可证密钥（敏感业务逻辑）\"\"\"
    expected = hmac.new(
        SECRET_KEY.encode(),
        license_key.encode(),
        hashlib.sha256
    ).hexdigest()
    return expected.startswith("a1b2")

def calculate_commission(amount: float, level: int) -> float:
    \"\"\"计算佣金（商业机密算法）\"\"\"
    rates = {1: 0.1, 2: 0.15, 3: 0.25, 4: 0.35}
    rate = rates.get(level, 0.05)
    # 敏感计算公式
    commission = amount * rate * (1.0 - 0.05)
    return round(commission, 2)

class PaymentProcessor:
    \"\"\"支付处理器（核心商业逻辑）\"\"\"
    def __init__(self, merchant_id: str):
        self.merchant_id = merchant_id
    
    def process(self, order_id: str, amount: float) -> dict:
        # 模拟支付处理
        return {
            "status": "success",
            "order_id": order_id,
            "amount": amount,
            "merchant": self.merchant_id,
            "transaction_id": hashlib.sha256(
                f"{order_id}{self.merchant_id}".encode()
            ).hexdigest()[:16]
        }

VERSION = "2.1.0"
BUILD = "2026-05-04"
"""


def main():
    print("=" * 70)
    print("PyArmor + AES 双层保护完整示例")
    print("=" * 70)

    # 创建临时工作目录
    # work_dir = tempfile.mkdtemp(prefix="dual_protection_")
    # db_path = os.path.join(work_dir, "protected_modules.db")
    db_path = os.path.join("dual_protection_", "protected_modules.db")
    key_file = os.path.join("dual_protection_", "encryption_key.txt")

    try:
        # ========== 步骤 1: 引入模块 ==========
        from module_bank.encryption import Encryption
        from module_bank.python_to_sqlite import PythonToSQLite

        # ========== 步骤 2: 生成 AES 加密密钥 ==========
        print("\n[步骤 1/5] 生成 AES 加密密钥...")
        aes_key = Encryption.generate_key()
        with open(key_file, "w") as f:
            f.write(aes_key)
        print(f"  AES 密钥已生成: {aes_key[:20]}...")

        # ========== 步骤 3: 使用双层保护打包模块 ==========
        print("\n[步骤 2/5] 使用 PyArmor + AES 双层保护打包模块...")
        packer = PythonToSQLite(
            db_path,
            key=aes_key,  # AES 字节码加密
            pyarmor_obfuscate=True,  # PyArmor 源代码混淆
            pyarmor_obf_code=1,  # 启用代码混淆
            pyarmor_obf_mod=1,  # 启用模块混淆
        )

        packer.importer.add_module(
            "payment_system",
            SECRET_MODULE_CODE,
            is_package=False,
            metadata={
                "version": "2.1.0",
                "protection": "pyarmor_obfuscation + aes_encryption",
            },
        )
        print("  ✓ 模块已使用双层保护存储到数据库")

        # ========== 步骤 4: 验证数据库中的保护状态 ==========
        print("\n[步骤 3/5] 验证数据库中的保护状态...")
        cursor = packer.importer.connection.execute(
            """SELECT module_name, encrypted, pyarmor_protected, source_code,
                      LENGTH(bytecode) as bytecode_size,
                      LENGTH(source_code) as source_size
               FROM python_modules"""
        )
        row = cursor.fetchone()
        print(f"  模块名: {row['module_name']}")
        print(f"  AES 加密: {'✓ 是' if row['encrypted'] else '✗ 否'}")
        print(f"  PyArmor 混淆: {'✓ 是' if row['pyarmor_protected'] else '✗ 否'}")
        print(f"  字节码大小: {row['bytecode_size']} 字节")
        print(f"  源代码大小: {row['source_size']} 字节")

        # 验证 PyArmor 运行时也已存储
        runtimes = packer.importer.list_runtimes()
        for rt in runtimes:
            print(f"  PyArmor 运行时: {rt['runtime_name']}")

        # 显示数据库中存储的混淆代码片段
        print("\n[验证] 数据库中的混淆代码（前200字符）:")
        obfuscated = row["source_code"]
        print(f"  {obfuscated[:200]}...")

        # ========== 步骤 5: 从数据库导入并执行 ==========
        print("\n[步骤 4/5] 从数据库导入模块（需要提供 AES 密钥）...")
        # 清除之前的导入
        if "payment_system" in sys.modules:
            del sys.modules["payment_system"]

        # 使用密钥安装导入器
        finder = packer.install_importer(key=aes_key)

        # 导入模块
        import importlib

        mod = importlib.import_module("payment_system")

        print(f"  ✓ 模块导入成功: {mod}")
        print(f"  模块文件: {mod.__file__}")

        # ========== 步骤 6: 验证模块功能正常 ==========
        print("\n[步骤 5/5] 验证模块功能...")

        # 调用验证许可证函数
        result = mod.verify_license("TEST-1234-5678")
        print(f"  verify_license('TEST-1234-5678') → {result}")

        # 调用佣金计算
        commission = mod.calculate_commission(10000.0, 3)
        print(f"  calculate_commission(10000, 3) → ¥{commission}")

        # 使用支付处理器
        processor = mod.PaymentProcessor("MERCHANT_001")
        payment_result = processor.process("ORDER_20260504", 299.00)
        print(f"  PaymentProcessor.process() → {payment_result['status']}")
        print(f"  交易ID: {payment_result['transaction_id']}")

        # 访问模块常量
        print(f"\n  模块版本: {mod.VERSION}")
        print(f"  构建日期: {mod.BUILD}")

        # ========== 安全验证 ==========
        print("\n" + "=" * 70)
        print("安全验证")
        print("=" * 70)
        print("""
  1. PyArmor 混淆: 源代码已转换为不可读的混淆代码
  2. AES 加密:     字节码已使用 AES-128-CBC 加密
  3. 运行时保护:    PyArmor 运行时二进制文件无法直接加载
  4. 双重解密:      需要先 AES 解密 → 再通过 PyArmor 运行时执行
  
  即使有人获取了数据库文件:
    - 没有 AES 密钥 → 无法解密字节码
    - 没有 PyArmor 运行时 → 无法执行混淆代码
    - 两者同时具备 → 才能解密和执行
        """)

    except Exception as e:
        print(f"\n❌ 错误: {e}")
        import traceback

        traceback.print_exc()
    finally:
        # 清理
        print("\n[清理] 清除临时文件...")
        # shutil.rmtree(work_dir, ignore_errors=True)
        if "payment_system" in sys.modules:
            del sys.modules["payment_system"]
        print("  完成！")


if __name__ == "__main__":
    main()