15308555518/ MaxKB
1
0
Fork 0
mirror of https://github.com/1Panel-dev/MaxKB.git synced 2025-12-29 07:52:50 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

feat: implement RestrictedUnpickler for secure deserialization and enhance import functionality

This commit is contained in:
CaptainB 2025-04-23 10:33:26 +08:00
parent cc0e431605
commit a9d962c1f1
1 changed files with 36 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
apps/tools/serializers/tool.py
View File

@ -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