diff --git a/apps/common/utils/tool_code.py b/apps/common/utils/tool_code.py index 02da9c1b8..13e894c3a 100644 --- a/apps/common/utils/tool_code.py +++ b/apps/common/utils/tool_code.py @@ -38,7 +38,7 @@ class ToolExecutor: os.umask(old_mask) def exec_code(self, code_str, keywords): - self.validateBannedKeywords(code_str) + self.validate_banned_keywords(code_str) _id = str(uuid.uuid7()) success = '{"code":200,"msg":"成功","data":exec_result}' err = '{"code":500,"msg":str(e),"data":None}' @@ -96,7 +96,7 @@ except Exception as e: os.remove(exec_python_file) return subprocess_result - def validateBannedKeywords(self, code_str): + def validate_banned_keywords(self, code_str): matched = next((bad for bad in self.banned_keywords if bad in code_str), None) if matched: raise Exception(f"keyword '{matched}' is banned in the tool.") diff --git a/apps/tools/serializers/tool.py b/apps/tools/serializers/tool.py index 4f37b99c3..323326bfb 100644 --- a/apps/tools/serializers/tool.py +++ b/apps/tools/serializers/tool.py @@ -200,6 +200,8 @@ class ToolSerializer(serializers.Serializer): if with_valid: self.is_valid(raise_exception=True) ToolCreateRequest(data=instance).is_valid(raise_exception=True) + # 校验代码是否包括禁止的关键字 + ToolExecutor().validate_banned_keywords(instance.get('code', '')) tool_id = uuid.uuid7() Tool( id=tool_id, @@ -312,6 +314,8 @@ class ToolSerializer(serializers.Serializer): 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 not QuerySet(Tool).filter(id=self.data.get('id')).exists(): raise serializers.ValidationError(_('Tool not found'))