Skip to content

Commit 9827e79

Browse files
liuruibinliuruibin
committed
feat: add import and export endpoints to ToolView for workspace tools
1 parent 5570559 commit 9827e79

File tree

7 files changed

+537
-67
lines changed

7 files changed

+537
-67
lines changed

‎apps/common/constants/permission_constants.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,12 @@ class PermissionConstants(Enum):
139139
RoleConstants.USER])
140140
TOOL_DELETE = Permission(group=Group.TOOL, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN,
141141
RoleConstants.USER])
142+
TOOL_DEBUG = Permission(group=Group.TOOL, operate=Operate.USE, role_list=[RoleConstants.ADMIN,
143+
RoleConstants.USER])
144+
TOOL_IMPORT = Permission(group=Group.TOOL, operate=Operate.USE, role_list=[RoleConstants.ADMIN,
145+
RoleConstants.USER])
146+
TOOL_EXPORT = Permission(group=Group.TOOL, operate=Operate.USE, role_list=[RoleConstants.ADMIN,
147+
RoleConstants.USER])
142148

143149
def get_workspace_application_permission(self):
144150
return lambda r, kwargs: Permission(group=self.value.group, operate=self.value.operate,

‎apps/common/utils/tool_code.py‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
# coding=utf-8
2+
3+
import os
4+
import subprocess
5+
import sys
6+
import uuid
7+
from textwrap import dedent
8+
9+
from diskcache import Cache
10+
11+
from maxkb.const import BASE_DIR
12+
from maxkb.const import PROJECT_DIR
13+
14+
python_directory = sys.executable
15+
16+
17+
class ToolExecutor:
18+
def __init__(self, sandbox=False):
19+
self.sandbox = sandbox
20+
if sandbox:
21+
self.sandbox_path = '/opt/maxkb/app/sandbox'
22+
self.user = 'sandbox'
23+
else:
24+
self.sandbox_path = os.path.join(PROJECT_DIR, 'data', 'sandbox')
25+
self.user = None
26+
self._createdir()
27+
if self.sandbox:
28+
os.system(f"chown -R {self.user}:root {self.sandbox_path}")
29+
30+
def _createdir(self):
31+
old_mask = os.umask(0o077)
32+
try:
33+
os.makedirs(self.sandbox_path, 0o700, exist_ok=True)
34+
finally:
35+
os.umask(old_mask)
36+
37+
def exec_code(self, code_str, keywords):
38+
_id = str(uuid.uuid1())
39+
success = '{"code":200,"msg":"成功","data":exec_result}'
40+
err = '{"code":500,"msg":str(e),"data":None}'
41+
path = r'' + self.sandbox_path + ''
42+
_exec_code = f"""
43+
try:
44+
import os
45+
env = dict(os.environ)
46+
for key in list(env.keys()):
47+
if key in os.environ and (key.startswith('MAXKB') or key.startswith('POSTGRES') or key.startswith('PG')):
48+
del os.environ[key]
49+
locals_v={'{}'}
50+
keywords={keywords}
51+
globals_v=globals()
52+
exec({dedent(code_str)!a}, globals_v, locals_v)
53+
f_name, f = locals_v.popitem()
54+
for local in locals_v:
55+
globals_v[local] = locals_v[local]
56+
exec_result=f(**keywords)
57+
from diskcache import Cache
58+
cache = Cache({path!a})
59+
cache.set({_id!a},{success})
60+
except Exception as e:
61+
from diskcache import Cache
62+
cache = Cache({path!a})
63+
cache.set({_id!a},{err})
64+
"""
65+
if self.sandbox:
66+
subprocess_result = self._exec_sandbox(_exec_code, _id)
67+
else:
68+
subprocess_result = self._exec(_exec_code)
69+
if subprocess_result.returncode == 1:
70+
raise Exception(subprocess_result.stderr)
71+
cache = Cache(self.sandbox_path)
72+
result = cache.get(_id)
73+
cache.delete(_id)
74+
if result.get('code') == 200:
75+
return result.get('data')
76+
raise Exception(result.get('msg'))
77+
78+
def _exec_sandbox(self, _code, _id):
79+
exec_python_file = f'{self.sandbox_path}/{_id}.py'
80+
with open(exec_python_file, 'w') as file:
81+
file.write(_code)
82+
os.system(f"chown {self.user}:{self.user} {exec_python_file}")
83+
kwargs = {'cwd': BASE_DIR}
84+
subprocess_result = subprocess.run(
85+
['su', '-s', python_directory, '-c', "exec(open('" + exec_python_file + "').read())", self.user],
86+
text=True,
87+
capture_output=True, **kwargs)
88+
os.remove(exec_python_file)
89+
return subprocess_result
90+
91+
@staticmethod
92+
def _exec(_code):
93+
return subprocess.run([python_directory, '-c', _code], text=True, capture_output=True)

‎apps/modules/views/module.py‎

Lines changed: 42 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,15 @@
1414
class ModuleView(APIView):
1515
authentication_classes = [TokenAuth]
1616

17-
@extend_schema(methods=['POST'],
18-
description=_('Create module'),
19-
operation_id=_('Create module'),
20-
parameters=ModuleCreateAPI.get_parameters(),
21-
request=ModuleCreateAPI.get_request(),
22-
responses=ModuleCreateAPI.get_response(),
23-
tags=[_('Module')])
17+
@extend_schema(
18+
methods=['POST'],
19+
description=_('Create module'),
20+
operation_id=_('Create module'),
21+
parameters=ModuleCreateAPI.get_parameters(),
22+
request=ModuleCreateAPI.get_request(),
23+
responses=ModuleCreateAPI.get_response(),
24+
tags=[_('Module')]
25+
)
2426
@has_permissions(lambda r, kwargs: Permission(group=Group(kwargs.get('source')), operate=Operate.CREATE,
2527
resource_path=f"/WORKSPACE/{kwargs.get('workspace_id')}"))
2628
def post(self, request: Request, workspace_id: str, source: str):
@@ -30,12 +32,14 @@ def post(self, request: Request, workspace_id: str, source: str):
3032
'workspace_id': workspace_id}
3133
).insert(request.data))
3234

33-
@extend_schema(methods=['GET'],
34-
description=_('Get module tree'),
35-
operation_id=_('Get module tree'),
36-
parameters=ModuleTreeReadAPI.get_parameters(),
37-
responses=ModuleTreeReadAPI.get_response(),
38-
tags=[_('Module')])
35+
@extend_schema(
36+
methods=['GET'],
37+
description=_('Get module tree'),
38+
operation_id=_('Get module tree'),
39+
parameters=ModuleTreeReadAPI.get_parameters(),
40+
responses=ModuleTreeReadAPI.get_response(),
41+
tags=[_('Module')]
42+
)
3943
@has_permissions(lambda r, kwargs: Permission(group=Group(kwargs.get('source')), operate=Operate.READ,
4044
resource_path=f"/WORKSPACE/{kwargs.get('workspace_id')}"))
4145
def get(self, request: Request, workspace_id: str, source: str):
@@ -46,39 +50,45 @@ def get(self, request: Request, workspace_id: str, source: str):
4650
class Operate(APIView):
4751
authentication_classes = [TokenAuth]
4852

49-
@extend_schema(methods=['PUT'],
50-
description=_('Update module'),
51-
operation_id=_('Update module'),
52-
parameters=ModuleEditAPI.get_parameters(),
53-
request=ModuleEditAPI.get_request(),
54-
responses=ModuleEditAPI.get_response(),
55-
tags=[_('Module')])
53+
@extend_schema(
54+
methods=['PUT'],
55+
description=_('Update module'),
56+
operation_id=_('Update module'),
57+
parameters=ModuleEditAPI.get_parameters(),
58+
request=ModuleEditAPI.get_request(),
59+
responses=ModuleEditAPI.get_response(),
60+
tags=[_('Module')]
61+
)
5662
@has_permissions(lambda r, kwargs: Permission(group=Group(kwargs.get('source')), operate=Operate.EDIT,
5763
resource_path=f"/WORKSPACE/{kwargs.get('workspace_id')}"))
5864
def put(self, request: Request, workspace_id: str, source: str, module_id: str):
5965
return result.success(ModuleSerializer.Operate(
6066
data={'id': module_id, 'workspace_id': workspace_id, 'source': source}
6167
).edit(request.data))
6268

63-
@extend_schema(methods=['GET'],
64-
description=_('Get module'),
65-
operation_id=_('Get module'),
66-
parameters=ModuleReadAPI.get_parameters(),
67-
responses=ModuleReadAPI.get_response(),
68-
tags=[_('Module')])
69+
@extend_schema(
70+
methods=['GET'],
71+
description=_('Get module'),
72+
operation_id=_('Get module'),
73+
parameters=ModuleReadAPI.get_parameters(),
74+
responses=ModuleReadAPI.get_response(),
75+
tags=[_('Module')]
76+
)
6977
@has_permissions(lambda r, kwargs: Permission(group=Group(kwargs.get('source')), operate=Operate.READ,
7078
resource_path=f"/WORKSPACE/{kwargs.get('workspace_id')}"))
7179
def get(self, request: Request, workspace_id: str, source: str, module_id: str):
7280
return result.success(ModuleSerializer.Operate(
7381
data={'id': module_id, 'workspace_id': workspace_id, 'source': source}
7482
).one())
7583

76-
@extend_schema(methods=['DELETE'],
77-
description=_('Delete module'),
78-
operation_id=_('Delete module'),
79-
parameters=ModuleDeleteAPI.get_parameters(),
80-
responses=ModuleDeleteAPI.get_response(),
81-
tags=[_('Module')])
84+
@extend_schema(
85+
methods=['DELETE'],
86+
description=_('Delete module'),
87+
operation_id=_('Delete module'),
88+
parameters=ModuleDeleteAPI.get_parameters(),
89+
responses=ModuleDeleteAPI.get_response(),
90+
tags=[_('Module')]
91+
)
8292
@has_permissions(lambda r, kwargs: Permission(group=Group(kwargs.get('source')), operate=Operate.DELETE,
8393
resource_path=f"/WORKSPACE/{kwargs.get('workspace_id')}"))
8494
def delete(self, request: Request, workspace_id: str, source: str, module_id: str):

‎apps/tools/api/tool.py‎

Lines changed: 102 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
from common.mixins.api_mixin import APIMixin
66
from common.result import ResultSerializer, DefaultResultSerializer
7-
from tools.serializers.tool import ToolModelSerializer, ToolCreateRequest
7+
from tools.serializers.tool import ToolModelSerializer, ToolCreateRequest, ToolDebugRequest
88

99

1010
class ToolCreateResponse(ResultSerializer):
@@ -91,3 +91,104 @@ def get_parameters():
9191
required=False,
9292
)
9393
]
94+
95+
96+
class ToolDebugApi(APIMixin):
97+
@staticmethod
98+
def get_request():
99+
return ToolDebugRequest
100+
101+
@staticmethod
102+
def get_response():
103+
return DefaultResultSerializer
104+
105+
106+
class ToolExportAPI(APIMixin):
107+
@staticmethod
108+
def get_parameters():
109+
return [
110+
OpenApiParameter(
111+
name="workspace_id",
112+
description="工作空间id",
113+
type=OpenApiTypes.STR,
114+
location='path',
115+
required=True,
116+
),
117+
OpenApiParameter(
118+
name="tool_id",
119+
description="工具id",
120+
type=OpenApiTypes.STR,
121+
location='path',
122+
required=True,
123+
)
124+
]
125+
126+
@staticmethod
127+
def get_response():
128+
return DefaultResultSerializer
129+
130+
131+
class ToolImportAPI(APIMixin):
132+
@staticmethod
133+
def get_parameters():
134+
return [
135+
OpenApiParameter(
136+
name="workspace_id",
137+
description="工作空间id",
138+
type=OpenApiTypes.STR,
139+
location='path',
140+
required=True,
141+
),
142+
OpenApiParameter(
143+
name='file',
144+
type=OpenApiTypes.BINARY,
145+
description='工具文件',
146+
required=True
147+
),
148+
]
149+
150+
@staticmethod
151+
def get_response():
152+
return DefaultResultSerializer
153+
154+
155+
class ToolPageAPI(ToolReadAPI):
156+
@staticmethod
157+
def get_parameters():
158+
return [
159+
OpenApiParameter(
160+
name="workspace_id",
161+
description="工作空间id",
162+
type=OpenApiTypes.STR,
163+
location='path',
164+
required=True,
165+
),
166+
OpenApiParameter(
167+
name="tool_id",
168+
description="工具id",
169+
type=OpenApiTypes.STR,
170+
location='path',
171+
required=True,
172+
),
173+
OpenApiParameter(
174+
name="current_page",
175+
description="当前页码",
176+
type=OpenApiTypes.INT,
177+
location='path',
178+
required=True,
179+
),
180+
OpenApiParameter(
181+
name="page_size",
182+
description="每页大小",
183+
type=OpenApiTypes.INT,
184+
location='path',
185+
required=True,
186+
),
187+
OpenApiParameter(
188+
name="name",
189+
description="工具名称",
190+
type=OpenApiTypes.STR,
191+
location='query',
192+
required=False,
193+
),
194+
]

0 commit comments

Comments
 (0)