要解决Langflow中PythonCodeTool的安全问题，可以在使用exec之前实现验证，以防止执行恶意代码。以下是操作步骤：

1. 使用validate_code函数:此函数将代码解析为抽象语法树(AST),并检查导入错误和函数定义错误。将此验证集成到PythonCodeTool类中。

2. 修改build方法:确保在执行代码之前进行验证。

以下是如何实现这个过程的示例：

步骤1:实现validate_code函数

创建一个名为validate.py的文件，内容如下：

import ast
import importlib

def add_type_ignores():
if not hasattr(ast, "TypeIgnore"):
class TypeIgnore(ast.AST):
_fields = ()
ast.TypeIgnore = TypeIgnore

def validate_code(code):
errors = {"imports": {"errors": []}, "function": {"errors": []}}

try:
    tree = ast.parse(code)
except Exception as e:
    errors["function"]["errors"].append(str(e))
    return errors

add_type_ignores()
tree.type_ignores = []

for node in tree.body:
    if isinstance(node, ast.Import):
        for alias in node.names:
            try:
                importlib.import_module(alias.name)
            except ModuleNotFoundError as e:
                errors["imports"]["errors"].append(str(e))

for node in tree.body:
    if isinstance(node, ast.FunctionDef):
        code_obj = compile(ast.Module(body=[node], type_ignores=[]), "<string>", "exec")
        try:
            exec(code_obj)
        except Exception as e:
            errors["function"]["errors"].append(str(e))

return errors

步骤2:将验证集成到PythonCodeTool类中

修改PythonCodeTool类，在执行任何代码之前使用validate_code函数：

from typing import Callable, Optional, List
from pydantic.v1 import BaseModel, validator
from langflow.utils import validate_code
from langchain.agents.tools import Tool

class Function(BaseModel):
code: str
function: Optional[Callable] = None
imports: Optional[str] = None

def __init__(self, **data):
    super().__init__(**data)

@validator("code")
def validate_func(cls, v):
    errors = validate_code(v)
    if errors["imports"]["errors"] or errors["function"]["errors"]:
        raise ValueError(f"Code validation errors: {errors}")
    return v

def get_function(self):
    function_name = self.extract_function_name(self.code)
    return self.create_function(self.code, function_name)

class PythonCodeTool(Function, Tool):
name: str = "Custom Tool"
description: str
code: str

def __init__(self, name: str, description: str, code: str):
    self.name = name
    self.description = description
    self.code = code
    self.func = self.get_function()
    super().__init__(name=name, description=description, func=self.func)

async def build(
    self,
    tool_code: str,
    name: str,
    description: str,
    tool_function: List[str],
    return_direct: bool,
    tool_class: Optional[List[str]] = None,
) -> Tool:
    errors = validate_code(tool_code)
    if errors["imports"]["errors"] or errors["function"]["errors"]:
        raise ValueError(f"Code validation errors: {errors}")

    local_namespace = {}
    exec(tool_code, globals(), local_namespace)

    func = local_namespace[tool_function]
    _class = None

    if tool_class:
        _class = local_namespace[tool_class]

    tool = StructuredTool.from_function(
        func=func, args_schema=_class, name=name, description=description, return_direct=return_direct
    )
    return tool

通过集成validate_code函数，可以确保在执行代码之前进行验证，从而防止执行恶意代码1(https://github.com/logspace-ai/langflow/blob/dev/src/backend/base/langflow/utils/validate.py)2(https://github.com/logspace-ai/langflow/blob/dev/src/backend/langflow/utils/validate.py)3(https://github.com/langflow-ai/langflow/blob/dev/src/backend/base/langflow/utils/validate.py)4(https://github.com/langflow-ai/langflow/blob/dev/src/backend/base/langflow/interface/tools/custom.py)5(https://github.com/logspace-ai/langflow/blob/dev/src/backend/langflow/interface/tools/custom.py)。

• 要继续与Dosu交流，请提及@dosu。*

我同意这个观点。使用exec是非常危险的，我在开发过程中对此感到担忧。我现在正在进行一项新设计。

预览

$x_{1}a_{0}b_{1}x$

LangChain的结构化工具中有一些选项，比如提取签名的选项。我正在考虑一个重用这些选项的格式，并且不使用代码字段。如果你有任何改进的建议，我会考虑它们。

为了避免 exec 的风险，我们可以：

• 在下面的代码中限制 find_imports 函数之后的导入。例如 os、subprocess 以及子模块如 langflow.utils.validate.importlib.resources.os。
• 由于有些模块无法避免，尝试在 docker 环境中执行代码。

你好@ogabrielluiz,
langflow的安全策略是什么？我想知道维护者是否会修补这个漏洞并请求CVE。
谢谢。