diff --git a/pyproject.toml b/pyproject.toml index 6e1368f..d991a23 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ description = "Add your description here" name = "module_bank" readme = "README.md" requires-python = ">=3.7" -version = "0.3.0" +version = "0.3.1" [[tool.uv.index]] default = true diff --git a/scripts/test_delete_source.py b/scripts/test_delete_source.py new file mode 100644 index 0000000..261db11 --- /dev/null +++ b/scripts/test_delete_source.py @@ -0,0 +1,100 @@ +""" +测试:删除 PyArmor 模块的源代码后仍能正常加载 +""" + +import sys +import os +import tempfile +import shutil + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) + + +def main(): + work_dir = tempfile.mkdtemp(prefix="test_delete_source_") + db_path = os.path.join(work_dir, "test.db") + + try: + from module_bank.python_to_sqlite import PythonToSQLite + + # 1. 创建并打包 PyArmor 模块 + print("=" * 60) + print("测试:删除源代码后加载 PyArmor 模块") + print("=" * 60) + + print("\n[1/5] 使用 PyArmor 打包模块...") + packer = PythonToSQLite(db_path, pyarmor_obfuscate=True) + packer.importer.add_module( + "mymod", + 'def hello(): return "ok"', + is_package=False, + ) + print(" 模块打包完成") + + # 2. 验证 metadata 中已记录运行时名 + print("\n[2/5] 验证 metadata 包含运行时名...") + cursor = packer.importer.connection.execute( + "SELECT metadata, source_code FROM python_modules WHERE module_name = ?", + ("mymod",), + ) + row = cursor.fetchone() + print(f" metadata: {row['metadata']}") + assert "pyarmor_runtime" in row["metadata"], "metadata 缺少运行时名" + print(" ✓ metadata 正确包含运行时名") + assert row["source_code"] is not None, "源代码不应为空" + print(" ✓ 源代码存在") + + # 3. 导入一次(确保能工作) + print("\n[3/5] 导入模块(有源代码)...") + from module_bank.sqlite_meta_path_finder import SQLiteMetaPathFinder + + finder = SQLiteMetaPathFinder(db_path) + sys.meta_path.insert(0, finder) + import importlib + + mod = importlib.import_module("mymod") + assert mod.hello() == "ok" + print(f" ✓ 导入成功: hello() = {mod.hello()}") + + # 4. 删除源代码 + print("\n[4/5] 删除源代码...") + packer.delete_source_code("mymod") + cursor = packer.importer.connection.execute( + "SELECT source_code, metadata FROM python_modules WHERE module_name = ?", + ("mymod",), + ) + row = cursor.fetchone() + print(f" 删除后 source_code = {row['source_code']}") + assert row["source_code"] is None, "源代码未成功删除" + print(f" 删除后 metadata = {row['metadata']}") + assert row["metadata"] is not None, "metadata 不应为空" + print(" ✓ metadata 仍保留") + + # 5. 重新导入(无源代码) + print("\n[5/5] 重新导入模块(无源代码)...") + if "mymod" in sys.modules: + del sys.modules["mymod"] + + mod2 = importlib.import_module("mymod") + result = mod2.hello() + print(f" ✓ 无源代码导入成功: hello() = {result}") + assert result == "ok", "模块功能异常" + + print("\n" + "=" * 60) + print("全部通过!删除源代码后仍可正常加载 PyArmor 模块") + print("=" * 60) + + except Exception as e: + print(f"\n❌ 测试失败: {e}") + import traceback + + traceback.print_exc() + return 1 + finally: + shutil.rmtree(work_dir, ignore_errors=True) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/module_bank/sqlite_module_importer.py b/src/module_bank/sqlite_module_importer.py index 004bd92..26e3c09 100644 --- a/src/module_bank/sqlite_module_importer.py +++ b/src/module_bank/sqlite_module_importer.py @@ -101,8 +101,16 @@ class SQLiteModuleImporter: pyarmor_protected = 1 # 存储运行时文件 + runtime_name = None if runtime_dict is not None: + runtime_name = runtime_dict.get("runtime_name") self._store_runtime(runtime_dict) + + # 将运行时名存入 metadata,以便在删除源代码后仍能定位运行时 + if runtime_name and metadata is None: + metadata = {"pyarmor_runtime": runtime_name} + elif runtime_name and isinstance(metadata, dict): + metadata["pyarmor_runtime"] = runtime_name else: print(f"警告: PyArmor 未安装,跳过模块 '{module_name}' 的 PyArmor 混淆") diff --git a/src/module_bank/sqlite_module_loader.py b/src/module_bank/sqlite_module_loader.py index c3f90bf..45221ff 100644 --- a/src/module_bank/sqlite_module_loader.py +++ b/src/module_bank/sqlite_module_loader.py @@ -7,6 +7,7 @@ import os import sys import tempfile import shutil +import json from pathlib import Path from .encryption import Encryption @@ -50,7 +51,7 @@ class SQLiteModuleLoader(importlib.abc.Loader): def exec_module(self, module): # type: ignore """执行模块代码""" cursor = self.db_connection.execute( - "SELECT bytecode, is_package, encrypted, pyarmor_protected, source_code FROM python_modules WHERE module_name = ?", + "SELECT bytecode, is_package, encrypted, pyarmor_protected, source_code, metadata FROM python_modules WHERE module_name = ?", (self.module_name,), ) row = cursor.fetchone() @@ -60,7 +61,7 @@ class SQLiteModuleLoader(importlib.abc.Loader): pyarmor_protected = row["pyarmor_protected"] - # 如果是 PyArmor 保护的模块,通过源代码执行(需要先部署运行时) + # 如果是 PyArmor 保护的模块,需要先部署运行时再执行 if pyarmor_protected: return self._exec_pyarmor_module(module, row) @@ -93,79 +94,112 @@ class SQLiteModuleLoader(importlib.abc.Loader): return module - def _exec_pyarmor_module(self, module, row): + def _deploy_runtime(self, runtime_name: str) -> str | None: """ - 执行 PyArmor 保护的模块 + 部署 PyArmor 运行时到临时目录 - PyArmor 混淆后的代码需要在运行时导入 pyarmor_runtime_xxx 包。 - 该包包含 __init__.py 和原生 .pyd/.so 文件。 - 我们需要从数据库提取这些文件并重建包结构到临时目录。 + Args: + runtime_name: 运行时模块名,如 'pyarmor_runtime_000000' + + Returns: + 临时目录路径,如果部署失败则返回 None """ - source_code = row["source_code"] - if not source_code: - raise ImportError( - f"PyArmor protected module '{self.module_name}' has no source code available" - ) + if self.db_path is None: + return None - # 提取运行时模块名 + try: + from .sqlite_module_importer import SQLiteModuleImporter + + importer = SQLiteModuleImporter(self.db_path) + runtime_dict = importer.get_runtime(runtime_name) + except Exception: + return None + + if not runtime_dict: + return None + + # 创建临时目录来存放运行时包 + temp_dir = tempfile.mkdtemp(prefix="pyarmor_runtime_") + runtime_pkg_dir = os.path.join(temp_dir, runtime_name) + os.makedirs(runtime_pkg_dir, exist_ok=True) + + # 确定运行时二进制文件名(.pyd for Windows, .so for Linux/Mac) + pyd_data = runtime_dict.get("pyd_data") + if pyd_data: + if sys.platform == "win32": + so_filename = "pyarmor_runtime.pyd" + elif sys.platform == "darwin": + so_filename = "pyarmor_runtime.dylib" + else: + so_filename = "pyarmor_runtime.so" + + so_path = os.path.join(runtime_pkg_dir, so_filename) + with open(so_path, "wb") as f: + f.write(pyd_data) + + # 写入 __init__.py + init_data = runtime_dict.get("init_data") + if init_data: + init_path = os.path.join(runtime_pkg_dir, "__init__.py") + with open(init_path, "w", encoding="utf-8") as f: + f.write(init_data) + + # 将临时目录添加到 sys.path,使运行时包可被导入 + if temp_dir not in sys.path: + sys.path.insert(0, temp_dir) + + return temp_dir + + def _extract_runtime_name_from_source(self, source_code: str) -> str | None: + """从混淆代码中提取运行时模块名""" import re match = re.search( r"from\s+(pyarmor_runtime_\w+)\s+import\s+__pyarmor__", source_code, ) - if not match: - # 没有找到运行时引用,直接执行源代码 - bytecode = compile( - source_code, f"", "exec" - ) - exec(bytecode, module.__dict__) - module.__file__ = f"sqlite://{self.db_path}/{self.module_name}" - return module + return match.group(1) if match else None - runtime_name = match.group(1) + def _extract_runtime_name_from_metadata( + self, metadata_str: str | None + ) -> str | None: + """从 metadata JSON 中提取运行时模块名""" + if not metadata_str: + return None + try: + metadata = json.loads(metadata_str) + return metadata.get("pyarmor_runtime") + except (json.JSONDecodeError, TypeError): + return None - # 从数据库获取运行时包数据 - runtime_dict = None - if self.db_path is not None: - try: - from .sqlite_module_importer import SQLiteModuleImporter + def _exec_pyarmor_module(self, module, row): + """ + 执行 PyArmor 保护的模块 - importer = SQLiteModuleImporter(self.db_path) - runtime_dict = importer.get_runtime(runtime_name) - except Exception: - pass + PyArmor 混淆后的代码需要在运行时导入 pyarmor_runtime_xxx 包。 + 该包包含 __init__.py 和原生 .pyd/.so 文件。 - if runtime_dict: - # 创建临时目录来存放运行时包 - self._runtime_temp_dir = tempfile.mkdtemp(prefix="pyarmor_runtime_") - runtime_pkg_dir = os.path.join(self._runtime_temp_dir, runtime_name) - os.makedirs(runtime_pkg_dir, exist_ok=True) + 支持两种模式: + 1. 有源代码:从源代码编译执行(标准模式) + 2. 无源代码:从已编译的字节码执行(delete_source 后的模式) + """ + source_code = row["source_code"] + bytecode_blob = row["bytecode"] + encrypted = row["encrypted"] + metadata_str = row["metadata"] - # 确定运行时二进制文件名(.pyd for Windows, .so for Linux/Mac) - pyd_data = runtime_dict.get("pyd_data") - if pyd_data: - if sys.platform == "win32": - so_filename = f"pyarmor_runtime.pyd" - elif sys.platform == "darwin": - so_filename = f"pyarmor_runtime.dylib" - else: - so_filename = f"pyarmor_runtime.so" + # 确定运行时名:优先从源代码提取,其次从 metadata + runtime_name = None + if source_code: + runtime_name = self._extract_runtime_name_from_source(source_code) + if not runtime_name: + runtime_name = self._extract_runtime_name_from_metadata(metadata_str) - so_path = os.path.join(runtime_pkg_dir, so_filename) - with open(so_path, "wb") as f: - f.write(pyd_data) - - # 写入 __init__.py - init_data = runtime_dict.get("init_data") - if init_data: - init_path = os.path.join(runtime_pkg_dir, "__init__.py") - with open(init_path, "w", encoding="utf-8") as f: - f.write(init_data) - - # 将临时目录添加到 sys.path,使运行时包可被导入 - if self._runtime_temp_dir not in sys.path: - sys.path.insert(0, self._runtime_temp_dir) + # 部署运行时 + if runtime_name: + self._runtime_temp_dir = self._deploy_runtime(runtime_name) + if self._runtime_temp_dir: + print(f" 已部署 PyArmor 运行时: {runtime_name}") # 设置 __file__ 属性(PyArmor 的 __pyarmor__() 需要它) module_file = f"sqlite://{self.db_path}/{self.module_name}" @@ -176,9 +210,30 @@ class SQLiteModuleLoader(importlib.abc.Loader): module.__path__ = [] module.__package__ = self.module_name - # 执行混淆后的代码 - bytecode = compile(source_code, f"", "exec") - exec(bytecode, module.__dict__) + # 执行代码 + if source_code: + # 模式1:有源代码 - 编译后执行 + bytecode = compile( + source_code, f"", "exec" + ) + exec(bytecode, module.__dict__) + else: + # 模式2:无源代码 - 从已存储的字节码执行 + if not bytecode_blob: + raise ImportError( + f"PyArmor protected module '{self.module_name}' has no source code or bytecode available" + ) + + # 如果字节码是加密的,需要解密 + if encrypted: + if self.encryption.cipher is None: + raise ImportError( + f"Module '{self.module_name}' has encrypted bytecode but no decryption key provided." + ) + bytecode_blob = self.encryption.decrypt(bytecode_blob) + + bytecode = marshal.loads(bytecode_blob) + exec(bytecode, module.__dict__) return module