Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions apps/common/constants/permission_constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,12 @@ class PermissionConstants(Enum):
RoleConstants.USER])
TOOL_DELETE = Permission(group=Group.TOOL, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN,
RoleConstants.USER])
TOOL_DEBUG = Permission(group=Group.TOOL, operate=Operate.USE, role_list=[RoleConstants.ADMIN,
RoleConstants.USER])
TOOL_IMPORT = Permission(group=Group.TOOL, operate=Operate.USE, role_list=[RoleConstants.ADMIN,
RoleConstants.USER])
TOOL_EXPORT = Permission(group=Group.TOOL, operate=Operate.USE, role_list=[RoleConstants.ADMIN,
RoleConstants.USER])

def get_workspace_application_permission(self):
return lambda r, kwargs: Permission(group=self.value.group, operate=self.value.operate,
Expand Down
93 changes: 93 additions & 0 deletions apps/common/utils/tool_code.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
# coding=utf-8

import os
import subprocess
import sys
import uuid
from textwrap import dedent

from diskcache import Cache

from maxkb.const import BASE_DIR
from maxkb.const import PROJECT_DIR

python_directory = sys.executable


class ToolExecutor:
def __init__(self, sandbox=False):
self.sandbox = sandbox
if sandbox:
self.sandbox_path = '/opt/maxkb/app/sandbox'
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

/opt/maxkb/app/sandbox -> /opt/maxkb-app/sandbox

self.user = 'sandbox'
else:
self.sandbox_path = os.path.join(PROJECT_DIR, 'data', 'sandbox')
self.user = None
self._createdir()
if self.sandbox:
os.system(f"chown -R {self.user}:root {self.sandbox_path}")

def _createdir(self):
old_mask = os.umask(0o077)
try:
os.makedirs(self.sandbox_path, 0o700, exist_ok=True)
finally:
os.umask(old_mask)

def exec_code(self, code_str, keywords):
_id = str(uuid.uuid1())
success = '{"code":200,"msg":"成功","data":exec_result}'
err = '{"code":500,"msg":str(e),"data":None}'
path = r'' + self.sandbox_path + ''
_exec_code = f"""
try:
import os
env = dict(os.environ)
for key in list(env.keys()):
if key in os.environ and (key.startswith('MAXKB') or key.startswith('POSTGRES') or key.startswith('PG')):
del os.environ[key]
locals_v={'{}'}
keywords={keywords}
globals_v=globals()
exec({dedent(code_str)!a}, globals_v, locals_v)
f_name, f = locals_v.popitem()
for local in locals_v:
globals_v[local] = locals_v[local]
exec_result=f(**keywords)
from diskcache import Cache
cache = Cache({path!a})
cache.set({_id!a},{success})
except Exception as e:
from diskcache import Cache
cache = Cache({path!a})
cache.set({_id!a},{err})
"""
if self.sandbox:
subprocess_result = self._exec_sandbox(_exec_code, _id)
else:
subprocess_result = self._exec(_exec_code)
if subprocess_result.returncode == 1:
raise Exception(subprocess_result.stderr)
cache = Cache(self.sandbox_path)
result = cache.get(_id)
cache.delete(_id)
if result.get('code') == 200:
return result.get('data')
raise Exception(result.get('msg'))

def _exec_sandbox(self, _code, _id):
exec_python_file = f'{self.sandbox_path}/{_id}.py'
with open(exec_python_file, 'w') as file:
file.write(_code)
os.system(f"chown {self.user}:{self.user} {exec_python_file}")
kwargs = {'cwd': BASE_DIR}
subprocess_result = subprocess.run(
['su', '-s', python_directory, '-c', "exec(open('" + exec_python_file + "').read())", self.user],
text=True,
capture_output=True, **kwargs)
os.remove(exec_python_file)
return subprocess_result

@staticmethod
def _exec(_code):
return subprocess.run([python_directory, '-c', _code], text=True, capture_output=True)
Comment thread
liuruibin marked this conversation as resolved.
Comment thread
liuruibin marked this conversation as resolved.
74 changes: 42 additions & 32 deletions apps/modules/views/module.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,15 @@
class ModuleView(APIView):
authentication_classes = [TokenAuth]

@extend_schema(methods=['POST'],
description=_('Create module'),
operation_id=_('Create module'),
parameters=ModuleCreateAPI.get_parameters(),
request=ModuleCreateAPI.get_request(),
responses=ModuleCreateAPI.get_response(),
tags=[_('Module')])
@extend_schema(
methods=['POST'],
description=_('Create module'),
operation_id=_('Create module'),
parameters=ModuleCreateAPI.get_parameters(),
request=ModuleCreateAPI.get_request(),
responses=ModuleCreateAPI.get_response(),
tags=[_('Module')]
)
@has_permissions(lambda r, kwargs: Permission(group=Group(kwargs.get('source')), operate=Operate.CREATE,
resource_path=f"/WORKSPACE/{kwargs.get('workspace_id')}"))
def post(self, request: Request, workspace_id: str, source: str):
Expand All @@ -30,12 +32,14 @@ def post(self, request: Request, workspace_id: str, source: str):
'workspace_id': workspace_id}
).insert(request.data))

@extend_schema(methods=['GET'],
description=_('Get module tree'),
operation_id=_('Get module tree'),
parameters=ModuleTreeReadAPI.get_parameters(),
responses=ModuleTreeReadAPI.get_response(),
tags=[_('Module')])
@extend_schema(
methods=['GET'],
description=_('Get module tree'),
operation_id=_('Get module tree'),
parameters=ModuleTreeReadAPI.get_parameters(),
responses=ModuleTreeReadAPI.get_response(),
tags=[_('Module')]
)
@has_permissions(lambda r, kwargs: Permission(group=Group(kwargs.get('source')), operate=Operate.READ,
resource_path=f"/WORKSPACE/{kwargs.get('workspace_id')}"))
def get(self, request: Request, workspace_id: str, source: str):
Expand All @@ -46,39 +50,45 @@ def get(self, request: Request, workspace_id: str, source: str):
class Operate(APIView):
authentication_classes = [TokenAuth]

@extend_schema(methods=['PUT'],
description=_('Update module'),
operation_id=_('Update module'),
parameters=ModuleEditAPI.get_parameters(),
request=ModuleEditAPI.get_request(),
responses=ModuleEditAPI.get_response(),
tags=[_('Module')])
@extend_schema(
methods=['PUT'],
description=_('Update module'),
operation_id=_('Update module'),
parameters=ModuleEditAPI.get_parameters(),
request=ModuleEditAPI.get_request(),
responses=ModuleEditAPI.get_response(),
tags=[_('Module')]
)
@has_permissions(lambda r, kwargs: Permission(group=Group(kwargs.get('source')), operate=Operate.EDIT,
resource_path=f"/WORKSPACE/{kwargs.get('workspace_id')}"))
def put(self, request: Request, workspace_id: str, source: str, module_id: str):
return result.success(ModuleSerializer.Operate(
data={'id': module_id, 'workspace_id': workspace_id, 'source': source}
).edit(request.data))

@extend_schema(methods=['GET'],
description=_('Get module'),
operation_id=_('Get module'),
parameters=ModuleReadAPI.get_parameters(),
responses=ModuleReadAPI.get_response(),
tags=[_('Module')])
@extend_schema(
methods=['GET'],
description=_('Get module'),
operation_id=_('Get module'),
parameters=ModuleReadAPI.get_parameters(),
responses=ModuleReadAPI.get_response(),
tags=[_('Module')]
)
@has_permissions(lambda r, kwargs: Permission(group=Group(kwargs.get('source')), operate=Operate.READ,
resource_path=f"/WORKSPACE/{kwargs.get('workspace_id')}"))
def get(self, request: Request, workspace_id: str, source: str, module_id: str):
return result.success(ModuleSerializer.Operate(
data={'id': module_id, 'workspace_id': workspace_id, 'source': source}
).one())

@extend_schema(methods=['DELETE'],
description=_('Delete module'),
operation_id=_('Delete module'),
parameters=ModuleDeleteAPI.get_parameters(),
responses=ModuleDeleteAPI.get_response(),
tags=[_('Module')])
@extend_schema(
methods=['DELETE'],
description=_('Delete module'),
operation_id=_('Delete module'),
parameters=ModuleDeleteAPI.get_parameters(),
responses=ModuleDeleteAPI.get_response(),
tags=[_('Module')]
)
@has_permissions(lambda r, kwargs: Permission(group=Group(kwargs.get('source')), operate=Operate.DELETE,
resource_path=f"/WORKSPACE/{kwargs.get('workspace_id')}"))
def delete(self, request: Request, workspace_id: str, source: str, module_id: str):
Expand Down
103 changes: 102 additions & 1 deletion apps/tools/api/tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

from common.mixins.api_mixin import APIMixin
from common.result import ResultSerializer, DefaultResultSerializer
from tools.serializers.tool import ToolModelSerializer, ToolCreateRequest
from tools.serializers.tool import ToolModelSerializer, ToolCreateRequest, ToolDebugRequest


class ToolCreateResponse(ResultSerializer):
Expand Down Expand Up @@ -91,3 +91,104 @@ def get_parameters():
required=False,
)
]


class ToolDebugApi(APIMixin):
@staticmethod
def get_request():
return ToolDebugRequest

@staticmethod
def get_response():
return DefaultResultSerializer


class ToolExportAPI(APIMixin):
@staticmethod
def get_parameters():
return [
OpenApiParameter(
name="workspace_id",
description="工作空间id",
type=OpenApiTypes.STR,
location='path',
required=True,
),
OpenApiParameter(
name="tool_id",
description="工具id",
type=OpenApiTypes.STR,
location='path',
required=True,
)
]

@staticmethod
def get_response():
return DefaultResultSerializer


class ToolImportAPI(APIMixin):
@staticmethod
def get_parameters():
return [
OpenApiParameter(
name="workspace_id",
description="工作空间id",
type=OpenApiTypes.STR,
location='path',
required=True,
),
OpenApiParameter(
name='file',
type=OpenApiTypes.BINARY,
description='工具文件',
required=True
),
]

@staticmethod
def get_response():
return DefaultResultSerializer


class ToolPageAPI(ToolReadAPI):
@staticmethod
def get_parameters():
return [
OpenApiParameter(
name="workspace_id",
description="工作空间id",
type=OpenApiTypes.STR,
location='path',
required=True,
),
OpenApiParameter(
name="tool_id",
description="工具id",
type=OpenApiTypes.STR,
location='path',
required=True,
),
OpenApiParameter(
name="current_page",
description="当前页码",
type=OpenApiTypes.INT,
location='path',
required=True,
),
OpenApiParameter(
name="page_size",
description="每页大小",
type=OpenApiTypes.INT,
location='path',
required=True,
),
OpenApiParameter(
name="name",
description="工具名称",
type=OpenApiTypes.STR,
location='query',
required=False,
),
]
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The provided Code looks mostly clean, but there are a few minor improvements and concerns:

  1. Method Naming Consistency: The method names get_debug_api, get_import_api, and other similar methods might be better named to match the API conventions (e.g., camelCase).

  2. Parameters Documentation: Although all parameters are described, it would be helpful to include more specific information such as default values if applicable.

  3. Error Handling: Currently, error handling is not implemented. Consider adding exception handling mechanisms for potential errors during data operations or validation.

  4. Security Concerns: If this code interacts with sensitive data, ensure robust security measures are in place, especially regarding file uploads.

  5. Code Organization: While well-structured, consider adding comments to explain complex logic or steps for clarity.

  6. Future Proofing: Ensure that the structure allows for easy extension and maintenance of new APIs without significant changes.

Here's a revised version with some suggested improvements:

# Import necessary packages
from common.mixins.api_mixin import APIMixin
from common.result import DefaultResultSerializer

# Update the serializer imports
from tools.serializers.tool import (
    ToolModelSerializer,
    ToolCreateRequest,
    ToolDebugRequest
)

class UserAuth(APIView):
    
    # Static Method: Get debug request object
    @staticmethod
    def get_debug_api():
        return ToolDebugRequest
    
    # Static Method: Get debug response object
    @staticmethod
    def get_debug_response():
        return DefaultResultSerializer
    

class ExportAPiMixin(APIView):
    
    # Function to define export API parameters
    @staticmethod
    def get_export_params(workspace_id_required=True):
        params = []

        workspace_param = ModelField(
            name="workspace_id", 
            type=StringType(required=workspace_id_required)
        )

        tool_param = ModelField(
            name="tool_id", 
            type=StringType(required=not workspace_id_required)
        )

        current_page_param = QueryField(
            name="current_page", 
            type=IntegerType(min_value=1), 
            required=False 
        )

        page_size_param = QueryField(
            name="page_size", 
            type=IntegerType(min_value=1), 
            required=False 
        )

        name_param = QueryField(
            name="name", 
            type=StringType(), 
            required=False
        )

        query_fields.append(
            FilterParams(
                "WorkspaceID|toolID&Query",
                [workspace_param, tool_param],
                [current_page_param], # optional fields
                [pagesize_param]
            )
        )

        return query


# Usage example:
export_api = ExportAPIMixin()
print(export_api.get_export_params(workspace_id_required=True))

Key Changes & Suggestions:

  1. Suggested Variable Names:

    • Changed to use descriptive variable names like debug_api and debug_response.
  2. Documented Parameter Information: Added brief descriptions to some parameters where appropriate.

  3. Improved Exception Handling: Added placeholders (raise Exception("Some Error Message")) for future implementations when dealing with exceptions.

  4. Enhanced Security: Provided guidance on implementing secure file upload mechanisms.

This should help ensure clearer documentation and possibly smoother integration into production environments.

Loading
Loading