modify: pyarmor 加密支持
This commit is contained in:
@@ -0,0 +1,295 @@
|
||||
"""
|
||||
集成测试：测试 PyArmor 加密打包 -> 存储 -> 加载 -> 执行 全流程
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import tempfile
|
||||
import shutil
|
||||
|
||||
# 添加项目模块路径
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
|
||||
|
||||
TEST_DB = os.path.join(os.path.dirname(__file__), "test_pyarmor_integration.db")
|
||||
|
||||
|
||||
def create_test_module():
|
||||
"""创建测试用的 Python 模块"""
|
||||
return """
|
||||
def greet(name):
|
||||
return f"Hello, {name}! This is a PyArmor protected module."
|
||||
|
||||
class Calculator:
|
||||
def add(self, a, b):
|
||||
return a + b
|
||||
|
||||
def multiply(self, a, b):
|
||||
return a * b
|
||||
|
||||
VERSION = "1.0.0"
|
||||
"""
|
||||
|
||||
|
||||
def test_pyarmor_obfuscate():
|
||||
"""测试 PyArmor 混淆功能"""
|
||||
print("=" * 60)
|
||||
print("测试 1: PyArmor 混淆功能")
|
||||
print("=" * 60)
|
||||
|
||||
from module_bank.pyarmor_encryption import PyArmorEncryption
|
||||
|
||||
enc = PyArmorEncryption(obf_code=1, obf_mod=1)
|
||||
|
||||
# 检查 PyArmor 是否可用
|
||||
if not enc.check_pyarmor_available():
|
||||
print("ERROR: PyArmor 未安装或不可用")
|
||||
return False
|
||||
|
||||
source = create_test_module()
|
||||
print(f"原始代码长度: {len(source)} 字节")
|
||||
|
||||
try:
|
||||
obfuscated_code, runtime_data = enc.obfuscate_source(
|
||||
source, "test_protected_module"
|
||||
)
|
||||
print(f"混淆后代码长度: {len(obfuscated_code)} 字节")
|
||||
print(f"运行时文件: {'有' if runtime_data else '无'}")
|
||||
print(f"运行时数据大小: {len(runtime_data) if runtime_data else 0} 字节")
|
||||
|
||||
# 验证是否被保护
|
||||
assert enc.is_pyarmor_protected(obfuscated_code), "代码应被标记为 PyArmor 保护"
|
||||
print(f"PyArmor 保护标记: 正确")
|
||||
|
||||
# 提取运行时模块名
|
||||
runtime_name = enc.extract_runtime_info(obfuscated_code)
|
||||
print(f"运行时模块名: {runtime_name}")
|
||||
|
||||
print("测试 1 通过!\n")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"ERROR: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
|
||||
def test_pack_with_pyarmor():
|
||||
"""测试使用 PyArmor 打包到 SQLite"""
|
||||
print("=" * 60)
|
||||
print("测试 2: 使用 PyArmor 打包到 SQLite")
|
||||
print("=" * 60)
|
||||
|
||||
# 清理旧的测试数据库
|
||||
if os.path.exists(TEST_DB):
|
||||
os.remove(TEST_DB)
|
||||
|
||||
try:
|
||||
from module_bank.python_to_sqlite import PythonToSQLite
|
||||
|
||||
# 使用 PyArmor 创建打包器
|
||||
packer = PythonToSQLite(
|
||||
TEST_DB,
|
||||
pyarmor_obfuscate=True,
|
||||
pyarmor_obf_code=1,
|
||||
pyarmor_obf_mod=1,
|
||||
)
|
||||
|
||||
source = create_test_module()
|
||||
packer.importer.add_module("test_protected_module", source, is_package=False)
|
||||
print("模块已打包到数据库")
|
||||
|
||||
# 验证模块已存储
|
||||
modules = packer.list_modules()
|
||||
module_names = [m["module_name"] for m in modules]
|
||||
assert "test_protected_module" in module_names, "模块应在数据库中存在"
|
||||
print(f"数据库中的模块: {module_names}")
|
||||
|
||||
# 验证运行时已存储
|
||||
runtimes = packer.importer.list_runtimes()
|
||||
print(f"数据库中的运行时: {[r['runtime_name'] for r in runtimes]}")
|
||||
|
||||
# 验证 pyarmor_protected 标记
|
||||
cursor = packer.importer.connection.execute(
|
||||
"SELECT pyarmor_protected FROM python_modules WHERE module_name = ?",
|
||||
("test_protected_module",),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
assert row["pyarmor_protected"] == 1, "模块应标记为 pyarmor_protected"
|
||||
print("PyArmor 保护标记: 已正确设置")
|
||||
|
||||
print("测试 2 通过!\n")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"ERROR: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
|
||||
def test_import_pyarmor_module():
|
||||
"""测试从 SQLite 导入并执行 PyArmor 保护的模块"""
|
||||
print("=" * 60)
|
||||
print("测试 3: 从 SQLite 导入并执行 PyArmor 保护模块")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
from module_bank.python_to_sqlite import PythonToSQLite
|
||||
from module_bank.sqlite_meta_path_finder import SQLiteMetaPathFinder
|
||||
|
||||
packer = PythonToSQLite(TEST_DB)
|
||||
|
||||
# 安装导入器
|
||||
finder = packer.install_importer()
|
||||
|
||||
# 尝试导入模块
|
||||
import importlib
|
||||
|
||||
mod = importlib.import_module("test_protected_module")
|
||||
print(f"模块导入成功: {mod}")
|
||||
print(f"模块名: {mod.__name__}")
|
||||
print(f"模块文件: {mod.__file__}")
|
||||
|
||||
# 调用模块函数
|
||||
result = mod.greet("World")
|
||||
print(f"调用 greet('World'): {result}")
|
||||
assert result == "Hello, World! This is a PyArmor protected module."
|
||||
|
||||
# 使用模块类
|
||||
calc = mod.Calculator()
|
||||
assert calc.add(3, 4) == 7
|
||||
assert calc.multiply(3, 4) == 12
|
||||
print(f"Calculator.add(3, 4) = {calc.add(3, 4)}")
|
||||
print(f"Calculator.multiply(3, 4) = {calc.multiply(3, 4)}")
|
||||
|
||||
# 访问模块变量
|
||||
print(f"VERSION = {mod.VERSION}")
|
||||
assert mod.VERSION == "1.0.0"
|
||||
|
||||
print("测试 3 通过!\n")
|
||||
|
||||
# 清理
|
||||
sys.meta_path.remove(finder)
|
||||
if "test_protected_module" in sys.modules:
|
||||
del sys.modules["test_protected_module"]
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"ERROR: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
|
||||
def test_cli_with_pyarmor():
|
||||
"""测试 CLI 的 PyArmor 参数"""
|
||||
print("=" * 60)
|
||||
print("测试 4: CLI --pyarmor 参数")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
import tempfile
|
||||
|
||||
# 创建临时源文件
|
||||
tmpdir = tempfile.mkdtemp()
|
||||
src_file = os.path.join(tmpdir, "hello_module.py")
|
||||
with open(src_file, "w", encoding="utf-8") as f:
|
||||
f.write("def hello():\n return 'Hello from PyArmor CLI!'\n")
|
||||
|
||||
db_path = os.path.join(tmpdir, "test_cli.db")
|
||||
|
||||
# 测试解析 --pyarmor 参数
|
||||
from module_bank.cli import main as cli_main
|
||||
|
||||
# 保存原始argv
|
||||
old_argv = sys.argv
|
||||
try:
|
||||
sys.argv = [
|
||||
"mbank",
|
||||
"pack",
|
||||
src_file,
|
||||
"--db",
|
||||
db_path,
|
||||
"--name",
|
||||
"hello_cli_module",
|
||||
"--pyarmor",
|
||||
]
|
||||
cli_main()
|
||||
print("CLI pack with --pyarmor 执行成功")
|
||||
|
||||
# 验证模块已用 PyArmor 保护
|
||||
from module_bank.sqlite_module_importer import SQLiteModuleImporter
|
||||
|
||||
importer = SQLiteModuleImporter(db_path)
|
||||
cursor = importer.connection.execute(
|
||||
"SELECT pyarmor_protected FROM python_modules WHERE module_name = ?",
|
||||
("hello_cli_module",),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
assert row and row["pyarmor_protected"] == 1, "CLI 打包的模块应启用 PyArmor"
|
||||
print("CLI 打包的模块已正确启用 PyArmor 保护")
|
||||
|
||||
# 验证运行时存储
|
||||
runtimes = importer.list_runtimes()
|
||||
assert len(runtimes) > 0, "CLI 打包后应有运行时文件存储在数据库中"
|
||||
print(f"运行时已存储: {[r['runtime_name'] for r in runtimes]}")
|
||||
|
||||
finally:
|
||||
sys.argv = old_argv
|
||||
shutil.rmtree(tmpdir, ignore_errors=True)
|
||||
|
||||
print("测试 4 通过!\n")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"ERROR: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
|
||||
def cleanup():
|
||||
"""清理测试文件"""
|
||||
if os.path.exists(TEST_DB):
|
||||
os.remove(TEST_DB)
|
||||
print(f"已清理测试数据库: {TEST_DB}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("\n" + "=" * 60)
|
||||
print("PyArmor 集成测试套件")
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
tests = [
|
||||
("PyArmor 混淆功能", test_pyarmor_obfuscate),
|
||||
("PyArmor 打包到 SQLite", test_pack_with_pyarmor),
|
||||
("从 SQLite 导入 PyArmor 模块", test_import_pyarmor_module),
|
||||
("CLI --pyarmor 参数", test_cli_with_pyarmor),
|
||||
]
|
||||
|
||||
passed = 0
|
||||
failed = 0
|
||||
|
||||
for name, func in tests:
|
||||
print(f"\n>>> 运行测试: {name}")
|
||||
try:
|
||||
if func():
|
||||
passed += 1
|
||||
else:
|
||||
failed += 1
|
||||
except Exception as e:
|
||||
print(f"测试异常: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
failed += 1
|
||||
|
||||
cleanup()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print(f"测试完成: {passed} 通过, {failed} 失败")
|
||||
print("=" * 60)
|
||||
|
||||
sys.exit(1 if failed > 0 else 0)
|
||||
Reference in New Issue
Block a user