tool.py

# -*- coding: utf-8 -*-
import asyncio
import io
import json
import os
import pickle
import re
import tempfile
import zipfile
from typing import Dict

import requests
import uuid_utils.compat as uuid
from django.core import validators
from django.db import transaction
from django.db.models import QuerySet, Q
from django.http import HttpResponse
from django.utils import timezone
from django.utils.translation import gettext_lazy as _
from langchain_mcp_adapters.client import MultiServerMCPClient
from pylint.lint import Run
from pylint.reporters import JSON2Reporter
from rest_framework import serializers, status

from common.database_model_manage.database_model_manage import DatabaseModelManage
from common.db.search import page_search, native_page_search, native_search
from common.exception.app_exception import AppApiException
from common.field.common import UploadedImageField
from common.result import result
from common.utils.common import get_file_content
from common.utils.logger import maxkb_logger
from common.utils.rsa_util import rsa_long_decrypt, rsa_long_encrypt
from common.utils.tool_code import ToolExecutor
from knowledge.models import File, FileSourceType
from maxkb.const import CONFIG, PROJECT_DIR
from system_manage.models import AuthTargetType, WorkspaceUserResourcePermission
from system_manage.serializers.user_resource_permission import UserResourcePermissionSerializer
from tools.models import Tool, ToolScope, ToolFolder, ToolType
from users.serializers.user import is_workspace_manage

tool_executor = ToolExecutor(CONFIG.get('SANDBOX'))


class ToolInstance:
    def __init__(self, tool: dict, version: str):
        self.tool = tool
        self.version = version


ALLOWED_CLASSES = {
    ("builtins", "dict"),
    ('uuid', 'UUID'),
    ("tools.serializers.tool", "ToolInstance")
}


def to_dict(message, file_name):
    return {
        'line': message.line,
        'column': message.column,
        'endLine': message.end_line,
        'endColumn': message.end_column,
        'message': (message.msg or "").replace(file_name, 'code'),
        'type': message.category
    }


def get_file_name():
    file_name = f"{uuid.uuid7()}"
    pylint_dir = os.path.join(PROJECT_DIR, 'data', 'pylint')
    if not os.path.exists(pylint_dir):
        os.makedirs(pylint_dir, 0o700, exist_ok=True)
        os.chmod(os.path.dirname(pylint_dir), 0o700)
    return os.path.join(pylint_dir, file_name)


class RestrictedUnpickler(pickle.Unpickler):

    def find_class(self, folder, name):
        if (folder, name) in ALLOWED_CLASSES:
            return super().find_class(folder, name)
        raise pickle.UnpicklingError("global '%s.%s' is forbidden" %
                                     (folder, name))


def encryption(message: str):
    """
        加密敏感字段数据  加密方式是 如果密码是 1234567890  那么给前端则是 123******890
    :param message:
    :return:
    """
    if type(message) != str:
        return message
    if message == "":
        return ""
    max_pre_len = 8
    max_post_len = 4
    message_len = len(message)
    pre_len = int(message_len / 5 * 2)
    post_len = int(message_len / 5 * 1)
    pre_str = "".join([message[index] for index in
                       range(0,
                             max_pre_len if pre_len > max_pre_len else 1 if pre_len <= 0 else int(
                                 pre_len))])
    end_str = "".join(
        [message[index] for index in
         range(message_len - (int(post_len) if pre_len < max_post_len else max_post_len),
               message_len)])
    content = "***************"
    return pre_str + content + end_str


def validate_mcp_config(servers: Dict):
    async def validate():
        client = MultiServerMCPClient(servers)
        await client.get_tools()

    try:
        asyncio.run(validate())
    except Exception as e:
        maxkb_logger.error(f"validate mcp config error: {e}, servers: {servers}")
        raise serializers.ValidationError(_('MCP configuration is invalid'))


class ToolModelSerializer(serializers.ModelSerializer):
    class Meta:
        model = Tool
        fields = ['id', 'name', 'icon', 'desc', 'code', 'input_field_list', 'init_field_list', 'init_params',
                  'scope', 'is_active', 'user_id', 'template_id', 'workspace_id', 'folder_id', 'tool_type', 'label',
                  'version', 'create_time', 'update_time']


class ToolExportModelSerializer(serializers.ModelSerializer):
    class Meta:
        model = Tool
        fields = ['id', 'name', 'icon', 'desc', 'code', 'input_field_list', 'init_field_list',
                  'scope', 'is_active', 'user_id', 'template_id', 'workspace_id', 'folder_id', 'tool_type', 'label',
                  'create_time', 'update_time']


class UploadedFileField(serializers.FileField):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def to_representation(self, value):
        return value


class ToolInputField(serializers.Serializer):
    name = serializers.CharField(required=True, label=_('variable name'))
    is_required = serializers.BooleanField(required=True, label=_('required'))
    type = serializers.CharField(required=True, label=_('type'), validators=[
        validators.RegexValidator(regex=re.compile("^string|int|dict|array|float$"),
                                  message=_('fields only support string|int|dict|array|float'), code=500)
    ])
    source = serializers.CharField(required=True, label=_('source'), validators=[
        validators.RegexValidator(regex=re.compile("^custom|reference$"),
                                  message=_('The field only supports custom|reference'), code=500)
    ])


class InitField(serializers.Serializer):
    field = serializers.CharField(required=True, label=_('field name'))
    label = serializers.CharField(required=True, label=_('field label'))
    required = serializers.BooleanField(required=True, label=_('required'))
    input_type = serializers.CharField(required=True, label=_('input type'))
    default_value = serializers.CharField(required=False, allow_null=True, allow_blank=True)
    show_default_value = serializers.BooleanField(required=False, default=False)
    props_info = serializers.DictField(required=False, default=dict)
    attrs = serializers.DictField(required=False, default=dict)


class ToolCreateRequest(serializers.Serializer):
    name = serializers.CharField(required=True, label=_('tool name'))

    desc = serializers.CharField(required=False, allow_null=True, allow_blank=True, label=_('tool description'))

    code = serializers.CharField(required=True, label=_('tool content'))

    input_field_list = serializers.ListField(required=False, default=list, label=_('input field list'))

    init_field_list = serializers.ListField(required=False, default=list, label=_('init field list'))

    is_active = serializers.BooleanField(required=False, label=_('Is active'))

    folder_id = serializers.CharField(required=False, allow_null=True)


class ToolEditRequest(serializers.Serializer):
    name = serializers.CharField(required=False, label=_('tool name'), allow_null=True)
    desc = serializers.CharField(required=False, allow_null=True, allow_blank=True, label=_('tool description'))
    code = serializers.CharField(required=False, label=_('tool content'), allow_null=True, )
    input_field_list = serializers.ListField(required=False, default=list, allow_null=True, label=_('input field list'))
    init_field_list = serializers.ListField(required=False, default=list, allow_null=True, label=_('init field list'))
    init_params = serializers.DictField(required=False, default=dict, allow_null=True, label=_('init params'))
    is_active = serializers.BooleanField(required=False, label=_('Is active'), allow_null=True, )
    folder_id = serializers.CharField(required=False, allow_null=True)


class AddInternalToolRequest(serializers.Serializer):
    name = serializers.CharField(required=False, label=_("tool name"), allow_null=True, allow_blank=True)
    folder_id = serializers.CharField(required=False, allow_null=True, label=_("folder id"))


class DebugField(serializers.Serializer):
    name = serializers.CharField(required=True, label=_('variable name'))
    value = serializers.CharField(required=False, allow_blank=True, allow_null=True, label=_('variable value'))


class ToolDebugRequest(serializers.Serializer):
    code = serializers.CharField(required=True, label=_('tool content'))
    input_field_list = serializers.ListField(required=False, default=list, label=_('input field list'))
    init_field_list = serializers.ListField(required=False, default=list, label=_('init field list'))
    init_params = serializers.DictField(required=False, default=dict, label=_('init params'))
    debug_field_list = DebugField(required=True, many=True)


class PylintInstance(serializers.Serializer):
    code = serializers.CharField(required=True, allow_null=True, allow_blank=True, label=_('function content'))


class ToolSerializer(serializers.Serializer):
    class Query(serializers.Serializer):
        workspace_id = serializers.CharField(required=True, label=_('workspace id'))
        folder_id = serializers.CharField(required=False, allow_blank=True, allow_null=True, label=_('folder id'))
        name = serializers.CharField(required=False, allow_null=True, allow_blank=True, label=_('tool name'))
        user_id = serializers.UUIDField(required=False, allow_null=True, label=_('user id'))
        scope = serializers.CharField(required=True, label=_('scope'))
        tool_type = serializers.CharField(required=False, label=_('tool type'), allow_null=True, allow_blank=True)
        create_user = serializers.UUIDField(required=False, label=_('create user'), allow_null=True)

        def get_query_set(self, workspace_manage, is_x_pack_ee):
            tool_query_set = QuerySet(Tool).filter(workspace_id=self.data.get('workspace_id'))
            folder_query_set = QuerySet(ToolFolder)
            default_query_set = QuerySet(Tool)

            workspace_id = self.data.get('workspace_id')
            user_id = self.data.get('user_id')
            scope = self.data.get('scope')
            tool_type = self.data.get('tool_type')
            desc = self.data.get('desc')
            name = self.data.get('name')
            folder_id = self.data.get('folder_id')
            create_user = self.data.get('create_user')

            if workspace_id is not None:
                folder_query_set = folder_query_set.filter(workspace_id=workspace_id)
                default_query_set = default_query_set.filter(workspace_id=workspace_id)
            if folder_id is not None:
                folder_query_set = folder_query_set.filter(parent=folder_id)
                default_query_set = default_query_set.filter(folder_id=folder_id)
            if name is not None:
                folder_query_set = folder_query_set.filter(name__icontains=name)
                default_query_set = default_query_set.filter(name__icontains=name)
            if desc is not None:
                folder_query_set = folder_query_set.filter(desc__icontains=desc)
                default_query_set = default_query_set.filter(desc__icontains=desc)
            if create_user is not None:
                tool_query_set = tool_query_set.filter(user_id=create_user)
                folder_query_set = folder_query_set.filter(user_id=create_user)

            default_query_set = default_query_set.order_by("-create_time")

            if scope is not None:
                tool_query_set = tool_query_set.filter(scope=scope)
            if tool_type:
                tool_query_set = tool_query_set.filter(tool_type=tool_type)

            query_set_dict = {
                'folder_query_set': folder_query_set,
                'tool_query_set': tool_query_set,
                'default_query_set': default_query_set,
            }
            if not workspace_manage:
                query_set_dict['workspace_user_resource_permission_query_set'] = QuerySet(
                    WorkspaceUserResourcePermission).filter(
                    auth_target_type="TOOL",
                    workspace_id=workspace_id,
                    user_id=user_id
                )
            return query_set_dict

        def get_authorized_query_set(self):
            default_query_set = QuerySet(Tool)
            tool_type = self.data.get('tool_type')
            desc = self.data.get('desc')
            name = self.data.get('name')
            create_user = self.data.get('create_user')

            default_query_set = default_query_set.filter(workspace_id='None')
            default_query_set = default_query_set.filter(scope=ToolScope.SHARED)
            if name is not None:
                default_query_set = default_query_set.filter(name__icontains=name)
            if desc is not None:
                default_query_set = default_query_set.filter(desc__icontains=desc)
            if create_user is not None:
                default_query_set = default_query_set.filter(user_id=create_user)
            if tool_type:
                default_query_set = default_query_set.filter(tool_type=tool_type)

            default_query_set = default_query_set.order_by("-create_time")

            return default_query_set

        @staticmethod
        def is_x_pack_ee():
            workspace_user_role_mapping_model = DatabaseModelManage.get_model("workspace_user_role_mapping")
            role_permission_mapping_model = DatabaseModelManage.get_model("role_permission_mapping_model")
            return workspace_user_role_mapping_model is not None and role_permission_mapping_model is not None

        def get_tools(self):
            self.is_valid(raise_exception=True)

            workspace_manage = is_workspace_manage(self.data.get('user_id'), self.data.get('workspace_id'))
            is_x_pack_ee = self.is_x_pack_ee()
            results = native_search(
                self.get_query_set(workspace_manage, is_x_pack_ee),
                get_file_content(
                    os.path.join(
                        PROJECT_DIR,
                        "apps", "tools", 'sql',
                        'list_tool.sql' if workspace_manage else (
                            'list_tool_user_ee.sql' if is_x_pack_ee else 'list_tool_user.sql'
                        )
                    )
                ),
            )

            get_authorized_tool = DatabaseModelManage.get_model("get_authorized_tool")
            shared_queryset = QuerySet(Tool).none()
            if get_authorized_tool is not None:
                shared_queryset = self.get_authorized_query_set()
                shared_queryset = get_authorized_tool(shared_queryset, self.data.get('workspace_id'))

            return {
                'shared_tools': [
                    ToolModelSerializer(data).data for data in shared_queryset
                ],
                'tools': [
                    {
                        **tool,
                        'input_field_list': json.loads(tool.get('input_field_list', '[]')),
                        'init_field_list': json.loads(tool.get('init_field_list', '[]')),
                    } for tool in results if tool['resource_type'] == 'tool'
                ],
            }

    class Create(serializers.Serializer):
        user_id = serializers.UUIDField(required=True, label=_('user id'))
        workspace_id = serializers.CharField(required=True, label=_('workspace id'))

        @transaction.atomic
        def insert(self, instance, with_valid=True):
            if with_valid:
                self.is_valid(raise_exception=True)
                ToolCreateRequest(data=instance).is_valid(raise_exception=True)
                # 校验代码是否包括禁止的关键字
                ToolExecutor().validate_banned_keywords(instance.get('code', ''))
                if instance.get('tool_type') == ToolType.MCP:
                    ToolExecutor().validate_mcp_transport(instance.get('code', ''))

            tool_id = uuid.uuid7()
            Tool(
                id=tool_id,
                name=instance.get('name'),
                desc=instance.get('desc'),
                code=instance.get('code'),
                user_id=self.data.get('user_id'),
                workspace_id=self.data.get('workspace_id'),
                input_field_list=instance.get('input_field_list', []),
                init_field_list=instance.get('init_field_list', []),
                scope=instance.get('scope', ToolScope.WORKSPACE),
                tool_type=instance.get('tool_type', ToolType.CUSTOM),
                folder_id=instance.get('folder_id', self.data.get('workspace_id')),
                is_active=False
            ).save()

            # 自动授权给创建者
            UserResourcePermissionSerializer(data={
                'workspace_id': self.data.get('workspace_id'),
                'user_id': self.data.get('user_id'),
                'auth_target_type': AuthTargetType.TOOL.value
            }).auth_resource(str(tool_id))
            return ToolSerializer.Operate(data={
                'id': tool_id, 'workspace_id': self.data.get('workspace_id')
            }).one()

    class TestConnection(serializers.Serializer):
        workspace_id = serializers.CharField(required=True, label=_('workspace id'))
        code = serializers.CharField(required=True, label=_('tool content'))

        def test_connection(self):
            self.is_valid(raise_exception=True)
            # 校验代码是否包括禁止的关键字
            ToolExecutor().validate_banned_keywords(self.data.get('code', ''))
            ToolExecutor().validate_mcp_transport(self.data.get('code', ''))

            # 校验mcp json
            validate_mcp_config(json.loads(self.data.get('code')))
            return True

    class Debug(serializers.Serializer):
        user_id = serializers.UUIDField(required=True, label=_('user id'))
        workspace_id = serializers.CharField(required=True, label=_('workspace id'))

        def debug(self, debug_instance):
            self.is_valid(raise_exception=True)
            ToolDebugRequest(data=debug_instance).is_valid(raise_exception=True)
            input_field_list = debug_instance.get('input_field_list')
            code = debug_instance.get('code')
            debug_field_list = debug_instance.get('debug_field_list')
            init_params = debug_instance.get('init_params')
            params = {field.get('name'): self.convert_value(field.get('name'), field.get('value'), field.get('type'),
                                                            field.get('is_required'))
                      for field in
                      [{'value': self.get_field_value(debug_field_list, field.get('name'), field.get('is_required')),
                        **field} for field in
                       input_field_list]}
            # 合并初始化参数
            if init_params is not None:
                all_params = init_params | params
            else:
                all_params = params
            return tool_executor.exec_code(code, all_params)

        @staticmethod
        def get_field_value(debug_field_list, name, is_required):
            result = [field for field in debug_field_list if field.get('name') == name]
            if len(result) > 0:
                return result[-1].get('value')
            if is_required:
                raise AppApiException(500, f"{name}" + _('field has no value set'))
            return None

        @staticmethod
        def convert_value(name: str, value: str, _type: str, is_required: bool):
            if not is_required and value is None:
                return None
            try:
                if _type == 'int':
                    return int(value)
                if _type == 'float':
                    return float(value)
                if _type == 'dict':
                    v = json.loads(value)
                    if isinstance(v, dict):
                        return v
                    raise Exception(_('type error'))
                if _type == 'array':
                    v = json.loads(value)
                    if isinstance(v, list):
                        return v
                    raise Exception(_('type error'))
                return value
            except Exception as e:
                raise AppApiException(500, _('Field: {name} Type: {_type} Value: {value} Type conversion error').format(
                    name=name, type=_type, value=value
                ))

    class Operate(serializers.Serializer):
        id = serializers.UUIDField(required=True, label=_('tool id'))
        workspace_id = serializers.CharField(required=True, label=_('workspace id'))

        def is_one_valid(self, *, raise_exception=False):
            super().is_valid(raise_exception=True)
            workspace_id = self.data.get('workspace_id')
            query_set = QuerySet(Tool).filter(id=self.data.get('id'))
            if workspace_id:
                query_set = query_set.filter(workspace_id=workspace_id)
            if not query_set.exists():
                get_authorized_tool = DatabaseModelManage.get_model('get_authorized_tool')
                if get_authorized_tool:
                    if not get_authorized_tool(QuerySet(Tool).filter(id=self.data.get('id')), workspace_id).exists():
                        raise AppApiException(500, _('Tool id does not exist'))

        def is_valid(self, *, raise_exception=False):
            super().is_valid(raise_exception=True)
            workspace_id = self.data.get('workspace_id')
            query_set = QuerySet(Tool).filter(id=self.data.get('id'))
            if workspace_id:
                query_set = query_set.filter(workspace_id=workspace_id)
            if not query_set.exists():
                raise AppApiException(500, _('Tool id does not exist'))

        def edit(self, instance, with_valid=True):
            if with_valid:
                self.is_valid(raise_exception=True)
                ToolEditRequest(data=instance).is_valid(raise_exception=True)
                # 校验代码是否包括禁止的关键字
                ToolExecutor().validate_banned_keywords(instance.get('code', ''))
                if instance.get('tool_type') == ToolType.MCP:
                    ToolExecutor().validate_mcp_transport(instance.get('code', ''))

            if not QuerySet(Tool).filter(id=self.data.get('id')).exists():
                raise serializers.ValidationError(_('Tool not found'))

            edit_field_list = ['name', 'desc', 'code', 'icon', 'input_field_list', 'init_field_list', 'init_params',
                               'is_active', 'folder_id']
            edit_dict = {field: instance.get(field) for field in edit_field_list if (
                    field in instance and instance.get(field) is not None)}

            tool = QuerySet(Tool).filter(id=self.data.get('id')).first()
            if 'init_params' in edit_dict:
                if edit_dict['init_field_list'] is not None:
                    rm_key = []
                    for key in edit_dict['init_params']:
                        if key not in [field['field'] for field in edit_dict['init_field_list']]:
                            rm_key.append(key)
                    for key in rm_key:
                        edit_dict['init_params'].pop(key)
                if tool.init_params:
                    old_init_params = json.loads(rsa_long_decrypt(tool.init_params))
                    for key in edit_dict['init_params']:
                        if key in old_init_params and edit_dict['init_params'][key] == encryption(old_init_params[key]):
                            edit_dict['init_params'][key] = old_init_params[key]
                edit_dict['init_params'] = rsa_long_encrypt(json.dumps(edit_dict['init_params']))

            edit_dict['update_time'] = timezone.now()
            QuerySet(Tool).filter(id=self.data.get('id')).update(**edit_dict)

            return self.one()

        def delete(self):
            self.is_valid(raise_exception=True)
            tool = QuerySet(Tool).filter(id=self.data.get('id')).first()
            if tool.template_id is None and tool.icon != '':
                QuerySet(File).filter(id=tool.icon.split('/')[-1]).delete()
            QuerySet(WorkspaceUserResourcePermission).filter(target=tool.id).delete()
            QuerySet(Tool).filter(id=self.data.get('id')).delete()

        def one(self):
            self.is_one_valid(raise_exception=True)
            tool = QuerySet(Tool).filter(id=self.data.get('id')).select_related('user').first()
            nick_name = tool.user.nick_name if tool and tool.user else None
            if tool.init_params:
                tool.init_params = json.loads(rsa_long_decrypt(tool.init_params))
            if tool.init_field_list:
                password_fields = [i["field"] for i in tool.init_field_list if
                                   i.get("input_type") == "PasswordInput"]
                if tool.init_params:
                    for k in tool.init_params:
                        if k in password_fields and tool.init_params[k]:
                            tool.init_params[k] = encryption(tool.init_params[k])
            return {
                **ToolModelSerializer(tool).data,
                'init_params': tool.init_params if tool.init_params else {},
                'nick_name': nick_name
            }

        def export(self):
            try:
                self.is_valid()
                id = self.data.get('id')
                tool = QuerySet(Tool).filter(id=id).first()
                tool_dict = ToolExportModelSerializer(tool).data
                mk_instance = ToolInstance(tool_dict, 'v2')
                tool_pickle = pickle.dumps(mk_instance)
                response = HttpResponse(content_type='text/plain', content=tool_pickle)
                response['Content-Disposition'] = f'attachment; filename="{tool.name}.tool"'
                return response
            except Exception as e:
                return result.error(str(e), response_status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    class Pylint(serializers.Serializer):

        def run(self, instance, is_valid=True):
            if is_valid:
                self.is_valid(raise_exception=True)
                PylintInstance(data=instance).is_valid(raise_exception=True)
            code = instance.get('code')
            file_name = get_file_name()
            with open(file_name, 'w') as file:
                file.write(code)
            reporter = JSON2Reporter(output=io.StringIO())
            Run([file_name,
                 "--disable=line-too-long",
                 '--module-rgx=[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}'],
                reporter=reporter, exit=False)
            os.remove(file_name)
            return [to_dict(m, os.path.basename(file_name)) for m in reporter.messages]

    class Import(serializers.Serializer):
        file = UploadedFileField(required=True, label=_("file"))
        user_id = serializers.UUIDField(required=True, label=_("User ID"))
        workspace_id = serializers.CharField(required=True, label=_("workspace id"))
        folder_id = serializers.CharField(required=False, allow_null=True, label=_("folder id"))

        #
        @transaction.atomic
        def import_(self, scope=ToolScope.WORKSPACE):
            self.is_valid()

            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"))
            if self.data.get('folder_id') is None:
                folder_id = self.data.get('workspace_id')
            else:
                folder_id = self.data.get('folder_id')
            tool = tool_instance.tool
            tool_id = uuid.uuid7()
            tool_model = Tool(
                id=tool_id,
                name=tool.get('name'),
                desc=tool.get('desc'),
                code=tool.get('code'),
                user_id=user_id,
                workspace_id=self.data.get('workspace_id'),
                input_field_list=tool.get('input_field_list'),
                init_field_list=tool.get('init_field_list', []),
                folder_id=folder_id,
                scope=scope,
                is_active=False
            )
            tool_model.save()

            # 自动授权给创建者
            UserResourcePermissionSerializer(data={
                'workspace_id': self.data.get('workspace_id'),
                'user_id': self.data.get('user_id'),
                'auth_target_type': AuthTargetType.TOOL.value
            }).auth_resource(str(tool_id))

            return True

    class IconOperate(serializers.Serializer):
        id = serializers.UUIDField(required=True, label=_("function ID"))
        workspace_id = serializers.CharField(required=True, label=_("workspace id"))
        user_id = serializers.UUIDField(required=True, label=_("User ID"))
        image = UploadedImageField(required=True, label=_("picture"))

        def is_valid(self, *, raise_exception=False):
            super().is_valid(raise_exception=True)
            workspace_id = self.data.get('workspace_id')
            query_set = QuerySet(Tool).filter(id=self.data.get('id'))
            if workspace_id:
                query_set = query_set.filter(workspace_id=workspace_id)
            if not query_set.exists():
                raise AppApiException(500, _('Tool id does not exist'))

        def edit(self, with_valid=True):
            if with_valid:
                self.is_valid(raise_exception=True)
            tool = QuerySet(Tool).filter(id=self.data.get('id')).first()
            if tool is None:
                raise AppApiException(500, _('Function does not exist'))
            # 删除旧的图片
            if tool.icon != '':
                QuerySet(File).filter(id=tool.icon.split('/')[-1]).delete()
            if self.data.get('image') is None:
                tool.icon = ''
            else:
                meta = {
                    'debug': False
                }
                file_id = uuid.uuid7()
                file = File(
                    id=file_id,
                    file_name=self.data.get('image').name,
                    source_type=FileSourceType.TOOL,
                    source_id=tool.id,
                    meta=meta
                )
                file.save(self.data.get('image').read())

                tool.icon = f'./oss/file/{file_id}'
            tool.save()

            return tool.icon

    class InternalTool(serializers.Serializer):
        user_id = serializers.UUIDField(required=True, label=_("User ID"))
        name = serializers.CharField(required=False, label=_("tool name"), allow_null=True, allow_blank=True)

        def get_internal_tools(self):
            self.is_valid(raise_exception=True)
            query_set = QuerySet(Tool)

            if self.data.get('name', '') != '':
                query_set = query_set.filter(
                    Q(name__icontains=self.data.get('name')) |
                    Q(desc__icontains=self.data.get('name'))
                )

            query_set = query_set.filter(
                Q(scope=ToolScope.INTERNAL) &
                Q(is_active=True)
            )
            return ToolModelSerializer(query_set, many=True).data

    class AddInternalTool(serializers.Serializer):
        user_id = serializers.UUIDField(required=True, label=_("User ID"))
        workspace_id = serializers.CharField(required=True, label=_("workspace id"))
        tool_id = serializers.UUIDField(required=True, label=_("tool id"))

        def add(self, instance, with_valid=True):
            if with_valid:
                self.is_valid(raise_exception=True)
                AddInternalToolRequest(data=instance).is_valid(raise_exception=True)

            internal_tool = QuerySet(Tool).filter(id=self.data.get('tool_id')).first()
            if internal_tool is None:
                raise AppApiException(500, _('Tool does not exist'))

            tool_id = uuid.uuid7()
            tool = Tool(
                id=tool_id,
                name=instance.get('name', internal_tool.name),
                desc=internal_tool.desc,
                code=internal_tool.code,
                user_id=self.data.get('user_id'),
                icon=internal_tool.icon,
                workspace_id=self.data.get('workspace_id'),
                input_field_list=internal_tool.input_field_list,
                init_field_list=internal_tool.init_field_list,
                scope=ToolScope.WORKSPACE,
                tool_type=ToolType.CUSTOM,
                folder_id=instance.get('folder_id', self.data.get('workspace_id')),
                template_id=internal_tool.id,
                label=internal_tool.label,
                is_active=False
            )
            tool.save()

            # 自动授权给创建者
            UserResourcePermissionSerializer(data={
                'workspace_id': self.data.get('workspace_id'),
                'user_id': self.data.get('user_id'),
                'auth_target_type': AuthTargetType.TOOL.value
            }).auth_resource(str(tool_id))

            return ToolModelSerializer(tool).data

    class StoreTool(serializers.Serializer):
        user_id = serializers.UUIDField(required=True, label=_("User ID"))
        name = serializers.CharField(required=False, label=_("tool name"), allow_null=True, allow_blank=True)

        def get_appstore_tools(self):
            self.is_valid(raise_exception=True)
            # 下载zip文件
            try:
                res = requests.get('https://apps-assets.fit2cloud.com/stable/maxkb.json.zip', timeout=5)
                res.raise_for_status()
                # 创建临时文件保存zip
                with tempfile.NamedTemporaryFile(delete=False, suffix='.zip') as temp_zip:
                    temp_zip.write(res.content)
                    temp_zip_path = temp_zip.name

                try:
                    # 解压zip文件
                    with zipfile.ZipFile(temp_zip_path, 'r') as zip_ref:
                        # 获取zip中的第一个文件（假设只有一个json文件）
                        json_filename = zip_ref.namelist()[0]
                        json_content = zip_ref.read(json_filename)

                    # 将json转换为字典
                    tool_store = json.loads(json_content.decode('utf-8'))
                    tag_dict = {tag['name']: tag['key'] for tag in tool_store['additionalProperties']['tags']}
                    filter_apps = []
                    for tool in tool_store['apps']:
                        if self.data.get('name', '') != '':
                            if self.data.get('name').lower() not in tool.get('name', '').lower():
                                continue
                        versions = tool.get('versions', [])
                        tool['label'] = tag_dict[tool.get('tags')[0]] if tool.get('tags') else ''
                        tool['version'] = next(
                            (version.get('name') for version in versions if
                             version.get('downloadUrl') == tool['downloadUrl']),
                        )
                        filter_apps.append(tool)

                    tool_store['apps'] = filter_apps
                    return tool_store
                finally:
                    # 清理临时文件
                    os.unlink(temp_zip_path)
            except Exception as e:
                maxkb_logger.error(f"fetch appstore tools error: {e}")
                return {'apps': [], 'additionalProperties': {'tags': []}}

    class AddStoreTool(serializers.Serializer):
        user_id = serializers.UUIDField(required=True, label=_("User ID"))
        workspace_id = serializers.CharField(required=True, label=_("workspace id"))
        tool_id = serializers.CharField(required=True, label=_("tool id"))

        def add(self, instance: Dict, with_valid=True):
            if with_valid:
                self.is_valid(raise_exception=True)
                AddInternalToolRequest(data=instance).is_valid(raise_exception=True)

            versions = instance.get('versions', [])
            download_url = instance.get('download_url')
            # 查找匹配的版本名称
            version_name = next(
                (version.get('name') for version in versions if version.get('downloadUrl') == download_url),
            )
            res = requests.get(download_url, timeout=5)
            tool_data = RestrictedUnpickler(io.BytesIO(res.content)).load().tool
            tool_id = uuid.uuid7()
            tool = Tool(
                id=tool_id,
                name=instance.get('name'),
                desc=tool_data.get('desc'),
                code=tool_data.get('code'),
                user_id=self.data.get('user_id'),
                icon=instance.get('icon', ''),
                workspace_id=self.data.get('workspace_id'),
                input_field_list=tool_data.get('input_field_list', []),
                init_field_list=tool_data.get('init_field_list', []),
                scope=ToolScope.WORKSPACE,
                tool_type=ToolType.CUSTOM,
                folder_id=instance.get('folder_id', self.data.get('workspace_id')),
                template_id=self.data.get('tool_id'),
                label=instance.get('label'),
                version=version_name,
                is_active=False
            )
            tool.save()

            # 自动授权给创建者
            UserResourcePermissionSerializer(data={
                'workspace_id': self.data.get('workspace_id'),
                'user_id': self.data.get('user_id'),
                'auth_target_type': AuthTargetType.TOOL.value
            }).auth_resource(str(tool_id))
            try:
                requests.get(instance.get('download_callback_url'), timeout=5)
            except Exception as e:
                maxkb_logger.error(f"callback appstore tool download error: {e}")
            return ToolModelSerializer(tool).data

    class UpdateStoreTool(serializers.Serializer):
        user_id = serializers.UUIDField(required=True, label=_("User ID"))
        workspace_id = serializers.CharField(required=True, label=_("workspace id"))
        tool_id = serializers.UUIDField(required=True, label=_("tool id"))
        download_url = serializers.CharField(required=True, label=_("download url"))
        download_callback_url = serializers.CharField(required=True, label=_("download callback url"))
        icon = serializers.CharField(required=True, label=_("icon"), allow_null=True, allow_blank=True)
        versions = serializers.ListField(required=True, label=_("versions"), child=serializers.DictField())

        def update_tool(self, with_valid=True):
            if with_valid:
                self.is_valid(raise_exception=True)
            tool = QuerySet(Tool).filter(id=self.data.get('tool_id')).first()
            if tool is None:
                raise AppApiException(500, _('Tool does not exist'))
            # 查找匹配的版本名称
            version_name = next(
                (version.get('name') for version in self.data.get('versions') if
                 version.get('downloadUrl') == self.data.get('download_url')),
            )
            res = requests.get(self.data.get('download_url'), timeout=5)
            tool_data = RestrictedUnpickler(io.BytesIO(res.content)).load().tool
            tool.desc = tool_data.get('desc')
            tool.code = tool_data.get('code')
            tool.input_field_list = tool_data.get('input_field_list', [])
            tool.init_field_list = tool_data.get('init_field_list', [])
            tool.icon = self.data.get('icon', tool.icon)
            tool.version = version_name
            # tool.is_active = False
            tool.save()
            try:
                requests.get(self.data.get('download_callback_url'), timeout=5)
            except Exception as e:
                maxkb_logger.error(f"callback appstore tool download error: {e}")
            return ToolModelSerializer(tool).data


class ToolTreeSerializer(serializers.Serializer):
    class Query(serializers.Serializer):
        workspace_id = serializers.CharField(required=True, label=_('workspace id'))
        folder_id = serializers.CharField(required=True, label=_('folder id'))
        name = serializers.CharField(required=False, allow_null=True, allow_blank=True, label=_('tool name'))
        user_id = serializers.UUIDField(required=False, allow_null=True, label=_('user id'))
        scope = serializers.CharField(required=True, label=_('scope'))
        tool_type = serializers.CharField(required=False, label=_('tool type'), allow_null=True, allow_blank=True)
        create_user = serializers.UUIDField(required=False, label=_('create user'), allow_null=True)

        def page_tool(self, current_page: int, page_size: int):
            self.is_valid(raise_exception=True)

            folder_id = self.data.get('folder_id', self.data.get('workspace_id'))
            root = ToolFolder.objects.filter(id=folder_id).first()
            if not root:
                raise serializers.ValidationError(_('Folder not found'))
            # 使用MPTT的get_descendants()方法获取所有相关节点
            all_folders = root.get_descendants(include_self=True)

            if self.data.get('name'):
                tools = QuerySet(Tool).filter(
                    Q(workspace_id=self.data.get('workspace_id')) &
                    Q(folder_id__in=all_folders) &
                    Q(user_id=self.data.get('user_id')) &
                    Q(name__contains=self.data.get('name'))
                )
            else:
                tools = QuerySet(Tool).filter(
                    Q(workspace_id=self.data.get('workspace_id')) &
                    Q(folder_id__in=all_folders) &
                    Q(user_id=self.data.get('user_id'))
                )
            return page_search(current_page, page_size, tools, lambda record: ToolModelSerializer(record).data)

        def get_query_set(self, workspace_manage, is_x_pack_ee):
            tool_query_set = QuerySet(Tool).filter(workspace_id=self.data.get('workspace_id'))
            folder_query_set = QuerySet(ToolFolder)
            default_query_set = QuerySet(Tool)

            workspace_id = self.data.get('workspace_id')
            user_id = self.data.get('user_id')
            scope = self.data.get('scope')
            tool_type = self.data.get('tool_type')
            desc = self.data.get('desc')
            name = self.data.get('name')
            folder_id = self.data.get('folder_id')
            create_user = self.data.get('create_user')

            if workspace_id is not None:
                folder_query_set = folder_query_set.filter(workspace_id=workspace_id)
                default_query_set = default_query_set.filter(workspace_id=workspace_id)
            if folder_id is not None and folder_id != workspace_id:
                folder_query_set = folder_query_set.filter(parent=folder_id)
                default_query_set = default_query_set.filter(folder_id=folder_id)
            if name is not None:
                folder_query_set = folder_query_set.filter(name__icontains=name)
                default_query_set = default_query_set.filter(name__icontains=name)
            if desc is not None:
                folder_query_set = folder_query_set.filter(desc__icontains=desc)
                default_query_set = default_query_set.filter(desc__icontains=desc)
            if create_user is not None:
                tool_query_set = tool_query_set.filter(user_id=create_user)
                folder_query_set = folder_query_set.filter(user_id=create_user)

            default_query_set = default_query_set.order_by("-create_time")

            if scope is not None:
                tool_query_set = tool_query_set.filter(scope=scope)
            if tool_type:
                tool_query_set = tool_query_set.filter(tool_type=tool_type)

            query_set_dict = {
                'tool_query_set': tool_query_set,
                'default_query_set': default_query_set,
            }
            if not workspace_manage:
                query_set_dict['workspace_user_resource_permission_query_set'] = QuerySet(
                    WorkspaceUserResourcePermission).filter(
                    auth_target_type="TOOL",
                    workspace_id=workspace_id,
                    user_id=user_id
                )
            return query_set_dict

        @staticmethod
        def is_x_pack_ee():
            workspace_user_role_mapping_model = DatabaseModelManage.get_model("workspace_user_role_mapping")
            role_permission_mapping_model = DatabaseModelManage.get_model("role_permission_mapping_model")
            return workspace_user_role_mapping_model is not None and role_permission_mapping_model is not None

        def page_tool_with_folders(self, current_page: int, page_size: int):
            self.is_valid(raise_exception=True)

            workspace_manage = is_workspace_manage(self.data.get('user_id'), self.data.get('workspace_id'))
            is_x_pack_ee = self.is_x_pack_ee()

            return native_page_search(
                current_page, page_size, self.get_query_set(workspace_manage, is_x_pack_ee),
                get_file_content(
                    os.path.join(
                        PROJECT_DIR,
                        "apps", "tools", 'sql',
                        'list_tool.sql' if workspace_manage else (
                            'list_tool_user_ee.sql' if is_x_pack_ee else 'list_tool_user.sql'
                        )
                    )
                ),
                post_records_handler=lambda record: {
                    **record,
                    'input_field_list': json.loads(record.get('input_field_list', '[]')),
                    'init_field_list': json.loads(record.get('init_field_list', '[]')),
                },
            )

        def get_tools(self):
            self.is_valid(raise_exception=True)

            workspace_manage = is_workspace_manage(self.data.get('user_id'), self.data.get('workspace_id'))
            is_x_pack_ee = self.is_x_pack_ee()
            results = native_search(
                self.get_query_set(workspace_manage, is_x_pack_ee),
                get_file_content(
                    os.path.join(
                        PROJECT_DIR,
                        "apps", "tools", 'sql',
                        'list_tool.sql' if workspace_manage else (
                            'list_tool_user_ee.sql' if is_x_pack_ee else 'list_tool_user.sql'
                        )
                    )
                ),

            )

            # 返回包含文件夹和工具的结构
            return {
                'folders': [
                    folder for folder in results if folder['resource_type'] == 'folder'
                ],
                'tools': [
                    {
                        **tool,
                        'input_field_list': json.loads(tool.get('input_field_list', '[]')),
                        'init_field_list': json.loads(tool.get('init_field_list', '[]')),
                    } for tool in results if tool['resource_type'] == 'tool'
                ],
            }