zhuchaowe/maxkb
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: implement RestrictedUnpickler for secure deserialization and enhance import functionality

Browse Source
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 -*- # -*- coding: utf-8 -*-
import io
import json import json
import pickle import pickle
import re import re
@ -26,6 +27,22 @@ class ToolInstance:
self.version = version 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): def encryption(message: str):
""" """
加密敏感字段数据 加密方式是 如果密码是 1234567890 那么给前端则是 123******890 加密敏感字段数据 加密方式是 如果密码是 1234567890 那么给前端则是 123******890
@ -258,26 +275,25 @@ class ToolSerializer(serializers.Serializer):
def import_(self): def import_(self):
self.is_valid() self.is_valid()
# user_id = self.data.get('user_id') user_id = self.data.get('user_id')
# flib_instance_bytes = self.data.get('file').read() tool_instance_bytes = self.data.get('file').read()
# try: try:
# RestrictedUnpickler(io.BytesIO(s)).load() tool_instance = RestrictedUnpickler(io.BytesIO(tool_instance_bytes)).load()
# flib_instance = restricted_loads(flib_instance_bytes) except Exception as e:
# except Exception as e: raise AppApiException(1001, _("Unsupported file format"))
# raise AppApiException(1001, _("Unsupported file format")) tool = tool_instance.tool
# tool = flib_instance.tool tool_model = Tool(
# tool_model = Tool( id=uuid.uuid7(),
# id=uuid.uuid7(), name=tool.get('name'),
# name=tool.get('name'), desc=tool.get('desc'),
# desc=tool.get('desc'), code=tool.get('code'),
# code=tool.get('code'), user_id=user_id,
# user_id=user_id, input_field_list=tool.get('input_field_list'),
# input_field_list=tool.get('input_field_list'), init_field_list=tool.get('init_field_list', []),
# init_field_list=tool.get('init_field_list', []), scope=ToolScope.WORKSPACE,
# scope=ToolScope.WORKSPACE, is_active=False
# is_active=False )
# ) tool_model.save()
# tool_model.save()
return True return True