From a9d962c1f142e2f88df2a0ace082d95ba4c4a360 Mon Sep 17 00:00:00 2001
From: CaptainB <bin@fit2cloud.com>
Date: Wed, 23 Apr 2025 10:33:26 +0800
Subject: [PATCH] feat: implement RestrictedUnpickler for secure
 deserialization and enhance import functionality

---
 apps/tools/serializers/tool.py | 56 ++++++++++++++++++++++------------
 1 file changed, 36 insertions(+), 20 deletions(-)

diff --git a/apps/tools/serializers/tool.py b/apps/tools/serializers/tool.py
index 3781b73ea..daad9dc64 100644
--- a/apps/tools/serializers/tool.py
+++ b/apps/tools/serializers/tool.py
@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+import io
 import json
 import pickle
 import re
@@ -26,6 +27,22 @@ class ToolInstance:
         self.version = version
 
 
+ALLOWED_CLASSES = {
+    ("builtins", "dict"),
+    ('uuid', 'UUID'),
+    ("tools.serializers.tool", "ToolInstance")
+}
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+
+    def find_class(self, module, name):
+        if (module, name) in ALLOWED_CLASSES:
+            return super().find_class(module, name)
+        raise pickle.UnpicklingError("global '%s.%s' is forbidden" %
+                                     (module, name))
+
+
 def encryption(message: str):
     """
         加密敏感字段数据  加密方式是 如果密码是 1234567890  那么给前端则是 123******890
@@ -258,26 +275,25 @@ class ToolSerializer(serializers.Serializer):
         def import_(self):
             self.is_valid()
 
-            # user_id = self.data.get('user_id')
-            # flib_instance_bytes = self.data.get('file').read()
-            # try:
-            #     RestrictedUnpickler(io.BytesIO(s)).load()
-            #     flib_instance = restricted_loads(flib_instance_bytes)
-            # except Exception as e:
-            #     raise AppApiException(1001, _("Unsupported file format"))
-            # tool = flib_instance.tool
-            # tool_model = Tool(
-            #     id=uuid.uuid7(),
-            #     name=tool.get('name'),
-            #     desc=tool.get('desc'),
-            #     code=tool.get('code'),
-            #     user_id=user_id,
-            #     input_field_list=tool.get('input_field_list'),
-            #     init_field_list=tool.get('init_field_list', []),
-            #     scope=ToolScope.WORKSPACE,
-            #     is_active=False
-            # )
-            # tool_model.save()
+            user_id = self.data.get('user_id')
+            tool_instance_bytes = self.data.get('file').read()
+            try:
+                tool_instance = RestrictedUnpickler(io.BytesIO(tool_instance_bytes)).load()
+            except Exception as e:
+                raise AppApiException(1001, _("Unsupported file format"))
+            tool = tool_instance.tool
+            tool_model = Tool(
+                id=uuid.uuid7(),
+                name=tool.get('name'),
+                desc=tool.get('desc'),
+                code=tool.get('code'),
+                user_id=user_id,
+                input_field_list=tool.get('input_field_list'),
+                init_field_list=tool.get('init_field_list', []),
+                scope=ToolScope.WORKSPACE,
+                is_active=False
+            )
+            tool_model.save()
             return True