Skip to content
12 changes: 7 additions & 5 deletions apps/application/flow/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,11 @@
from models_provider.tools import get_model_credential
from tools.models.tool import Tool

end_nodes = ['ai-chat-node', 'reply-node', 'function-node', 'function-lib-node', 'application-node',
'image-understand-node', 'speech-to-text-node', 'text-to-speech-node', 'image-generate-node',
'variable-assign-node']
END_NODES = frozenset([
'ai-chat-node', 'reply-node', 'function-node', 'function-lib-node', 'application-node',
'image-understand-node', 'speech-to-text-node', 'text-to-speech-node', 'image-generate-node',
'variable-assign-node'
])


class Answer:
Expand Down Expand Up @@ -219,14 +221,14 @@ def is_valid_node(self, node: Node):
for branch in branch_list:
source_anchor_id = f"{node.id}_{branch.get('id')}_right"
edge_list = [edge for edge in self.edges if edge.sourceAnchorId == source_anchor_id]
if len(edge_list) == 0:
if not edge_list:
raise AppApiException(500,
_('The branch {branch} of the {node} node needs to be connected').format(
node=node.properties.get("stepName"), branch=branch.get("type")))

else:
edge_list = [edge for edge in self.edges if edge.sourceNodeId == node.id]
if len(edge_list) == 0 and not end_nodes.__contains__(node.type):
if not edge_list and node.type not in END_NODES:
raise AppApiException(500, _("{node} Nodes cannot be considered as end nodes").format(
node=node.properties.get("stepName")))

Expand Down
10 changes: 6 additions & 4 deletions apps/application/flow/compare/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@
from .start_with import StartWithCompare
from .wildcard_compare import WildcardCompare

from common.utils.logger import maxkb_logger

_compare_handler_dict = {
'is_null': IsNullCompare(),
'is_not_null': IsNotNullCompare(),
Expand Down Expand Up @@ -65,13 +67,13 @@ def _compare(source_value, compare, target_value):
def _assertion(workflow_manage, field_list: List[str], compare: str, value):
try:
value = workflow_manage.generate_prompt(value)
except Exception:
pass
except Exception as e:
maxkb_logger.debug(f"Failed to generate field value for comparison: {e}")
field_value = None
try:
field_value = workflow_manage.get_reference_field(field_list[0], field_list[1:])
except Exception:
pass
except Exception as e:
maxkb_logger.debug(f"Failed to get reference field for comparison: {e}")
return _compare(field_value, compare, value)


Expand Down
87 changes: 87 additions & 0 deletions apps/application/serializers/application.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,76 @@
from application.serializers.common import update_resource_mapping_by_application


def get_bound_tool_ids(instance: Dict) -> List[str]:
"""
收集应用配置(含工作流节点)中引用的所有工具id,用于绑定前的权限校验
"""
tool_ids = set()
for key in ("tool_ids", "skill_tool_ids", "mcp_tool_ids"):
for tool_id in (instance.get(key) or []):
tool_ids.add(str(tool_id))
if instance.get("mcp_tool_id"):
tool_ids.add(str(instance.get("mcp_tool_id")))

def walk(work_flow):
if not work_flow:
return
for node in work_flow.get("nodes", []) or []:
node_data = (node.get("properties") or {}).get("node_data") or {}
for key in ("tool_lib_id", "mcp_tool_id"):
if node_data.get(key):
tool_ids.add(str(node_data.get(key)))
for key in ("mcp_tool_ids", "tool_ids", "skill_tool_ids"):
for tool_id in (node_data.get(key) or []):
tool_ids.add(str(tool_id))
if node.get("type") == "loop-node":
walk(node_data.get("loop_body"))

walk(instance.get("work_flow"))
return list(tool_ids)


def get_authorized_tool_ids(user_id: str, workspace_id: str, tool_ids: List[str]) -> List[str]:
"""
返回 tool_ids 中当前用户被授权绑定/使用的工具id。
工作空间管理员默认拥有全部工具权限;其他用户必须在 workspace_user_resource_permission
中存在针对该工具的显式授权记录(默认拒绝)。
"""
if not tool_ids:
return []
tool_ids = list({str(t) for t in tool_ids})
if is_workspace_manage(user_id, workspace_id):
return tool_ids
granted_tool_ids = {
str(permission.target)
for permission in QuerySet(WorkspaceUserResourcePermission).filter(
workspace_id=workspace_id,
user_id=user_id,
auth_target_type=AuthTargetType.TOOL.value,
target__in=tool_ids,
)
if "VIEW" in permission.permission_list or "ROLE" in permission.permission_list
}
return [tool_id for tool_id in tool_ids if tool_id in granted_tool_ids]


def validate_bound_tool_permissions(user_id: str, workspace_id: str, instance: Dict):
"""
校验应用/工作流中绑定的工具,当前用户是否都有权限使用,防止低权限成员
绑定自己被禁止访问的工具,并通过应用/工作流执行绕过工具的单独授权控制。
"""
tool_ids = get_bound_tool_ids(instance)
if not tool_ids:
return
authorized_tool_ids = set(get_authorized_tool_ids(user_id, workspace_id, tool_ids))
unauthorized_tool_ids = [tool_id for tool_id in tool_ids if tool_id not in authorized_tool_ids]
if unauthorized_tool_ids:
message = lazy_format(
_("No permission to use tool(s): {tool_ids}"), tool_ids=", ".join(unauthorized_tool_ids)
)
raise AppApiException(403, str(message))


def get_base_node_work_flow(work_flow):
node_list = work_flow.get("nodes")
base_node_list = [node for node in node_list if node.get("id") == "base-node"]
Expand Down Expand Up @@ -623,6 +693,7 @@ def insert_workflow(self, instance: Dict):
workspace_id = self.data.get("workspace_id")
wq = ApplicationCreateSerializer.WorkflowRequest(data=instance)
wq.is_valid(raise_exception=True)
validate_bound_tool_permissions(user_id, workspace_id, instance)
application_model = wq.to_application_model(user_id, workspace_id, instance)
application_model.save()
# 插入认证信息
Expand Down Expand Up @@ -703,6 +774,20 @@ def import_(self, instance: dict, is_import_tool, with_valid=True):
if not exits_tool_id_list.__contains__(tool.get("id"))
and not exits_tool_id_list.__contains__(generate_uuid((tool.get("id") + workspace_id or "")))
]
# 导入包内新建的工具由导入者本人持有,无需校验;仅需校验绑定到已存在工具的引用
existing_bound_tool_ids = [
tool_id for tool_id in get_bound_tool_ids(application) if tool_id not in update_tool_map
]
if existing_bound_tool_ids:
authorized_tool_ids = set(get_authorized_tool_ids(user_id, workspace_id, existing_bound_tool_ids))
unauthorized_tool_ids = [
tool_id for tool_id in existing_bound_tool_ids if tool_id not in authorized_tool_ids
]
if unauthorized_tool_ids:
message = lazy_format(
_("No permission to use tool(s): {tool_ids}"), tool_ids=", ".join(unauthorized_tool_ids)
)
raise AppApiException(403, str(message))
application_model = self.to_application(application, workspace_id, user_id, update_tool_map, folder_id)
tool_model_list = [self.to_tool(f, workspace_id, user_id) for f in tool_list]
application_model.save()
Expand Down Expand Up @@ -1190,6 +1275,8 @@ def edit(self, instance: Dict, with_valid=True):
if "work_flow_template" in instance:
return self.update_template_workflow(instance, application)

validate_bound_tool_permissions(self.data.get("user_id"), self.data.get("workspace_id"), instance)

if instance.get("model_id") is None or len(instance.get("model_id")) == 0:
application.model_id = None
else:
Expand Down
3 changes: 3 additions & 0 deletions apps/chat/mcp/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

import uuid_utils.compat as uuid
from django.db.models import QuerySet
from django.utils import timezone

from application.models import ApplicationApiKey, Application, ChatUserType, ChatSourceChoices
from chat.serializers.chat import ChatSerializers
Expand All @@ -13,6 +14,8 @@ def __init__(self, auth_header):
app_key = QuerySet(ApplicationApiKey).filter(secret_key=auth_header, is_active=True).first()
if not app_key:
raise PermissionError("Invalid API Key")
if app_key.is_permanent is False and app_key.expire_time < timezone.now():
raise PermissionError("API Key is expired")

self.application = QuerySet(Application).filter(id=app_key.application_id, is_publish=True).first()
if not self.application:
Expand Down
Loading
Loading