modify: 修复 pyarmor 模式无法使用字节码加载的问题

2026-05-04 20:04:31 +08:00
parent 1fdb5bd1bf
commit fa836f9453
4 changed files with 226 additions and 63 deletions
@@ -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
@@ -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())
@@ -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 混淆")

@@ -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,64 +94,44 @@ 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

-        # 提取运行时模块名
-        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"<sqlite_module:{self.module_name}>", "exec"
-            )
-            exec(bytecode, module.__dict__)
-            module.__file__ = f"sqlite://{self.db_path}/{self.module_name}"
-            return module
-
-        runtime_name = match.group(1)
-
-        # 从数据库获取运行时包数据
-        runtime_dict = None
-        if self.db_path is not None:
        try:
            from .sqlite_module_importer import SQLiteModuleImporter

            importer = SQLiteModuleImporter(self.db_path)
            runtime_dict = importer.get_runtime(runtime_name)
        except Exception:
-                pass
+            return None
+
+        if not runtime_dict:
+            return None

-        if runtime_dict:
        # 创建临时目录来存放运行时包
-            self._runtime_temp_dir = tempfile.mkdtemp(prefix="pyarmor_runtime_")
-            runtime_pkg_dir = os.path.join(self._runtime_temp_dir, runtime_name)
+        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 = f"pyarmor_runtime.pyd"
+                so_filename = "pyarmor_runtime.pyd"
            elif sys.platform == "darwin":
-                    so_filename = f"pyarmor_runtime.dylib"
+                so_filename = "pyarmor_runtime.dylib"
            else:
-                    so_filename = f"pyarmor_runtime.so"
+                so_filename = "pyarmor_runtime.so"

            so_path = os.path.join(runtime_pkg_dir, so_filename)
            with open(so_path, "wb") as f:
@@ -164,8 +145,61 @@ class SQLiteModuleLoader(importlib.abc.Loader):
                f.write(init_data)

        # 将临时目录添加到 sys.path，使运行时包可被导入
-            if self._runtime_temp_dir not in sys.path:
-                sys.path.insert(0, self._runtime_temp_dir)
+        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,
+        )
+        return match.group(1) if match else None
+
+    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
+
+    def _exec_pyarmor_module(self, module, row):
+        """
+        执行 PyArmor 保护的模块
+
+        PyArmor 混淆后的代码需要在运行时导入 pyarmor_runtime_xxx 包。
+        该包包含 __init__.py 和原生 .pyd/.so 文件。
+
+        支持两种模式：
+        1. 有源代码：从源代码编译执行（标准模式）
+        2. 无源代码：从已编译的字节码执行（delete_source 后的模式）
+        """
+        source_code = row["source_code"]
+        bytecode_blob = row["bytecode"]
+        encrypted = row["encrypted"]
+        metadata_str = row["metadata"]
+
+        # 确定运行时名：优先从源代码提取，其次从 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)
+
+        # 部署运行时
+        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,8 +210,29 @@ class SQLiteModuleLoader(importlib.abc.Loader):
            module.__path__ = []
            module.__package__ = self.module_name

-        # 执行混淆后的代码
-        bytecode = compile(source_code, f"<sqlite_module:{self.module_name}>", "exec")
+        # 执行代码
+        if source_code:
+            # 模式1：有源代码 - 编译后执行
+            bytecode = compile(
+                source_code, f"<sqlite_module:{self.module_name}>", "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