From 0e9c87114734ce88c059e9ca76cb6aa872feadaa Mon Sep 17 00:00:00 2001 From: Pankaj Kaushal Date: Wed, 1 Oct 2025 22:04:29 +0530 Subject: [PATCH] obfuscate the secret variable --- api/core/helper/encrypter.py | 58 +++++++++++++++++++++ api/core/workflow/nodes/agent/agent_node.py | 17 +++++- api/core/workflow/nodes/code/code_node.py | 18 ++++++- 3 files changed, 89 insertions(+), 4 deletions(-) diff --git a/api/core/helper/encrypter.py b/api/core/helper/encrypter.py index 17345dc203..181c1b977b 100644 --- a/api/core/helper/encrypter.py +++ b/api/core/helper/encrypter.py @@ -1,4 +1,6 @@ import base64 +from collections.abc import Mapping +from typing import Any, overload from libs import rsa @@ -42,3 +44,59 @@ def get_decrypt_decoding(tenant_id: str): def decrypt_token_with_decoding(token: str, rsa_key, cipher_rsa): return rsa.decrypt_token_with_decoding(base64.b64decode(token), rsa_key, cipher_rsa) + + +# ========================= +# encrypt_secret_keys +# ========================= + + +# Overloads to preserve input type +@overload +def encrypt_secret_keys( + obj: Mapping[str, Any], + secret_variables: list[str] | None = None, + parent_key: str | None = None, +) -> Mapping[str, Any]: ... + + +@overload +def encrypt_secret_keys( + obj: list[Any], + secret_variables: list[str] | None = None, + parent_key: str | None = None, +) -> list[Any]: ... + + +@overload +def encrypt_secret_keys( + obj: Any, + secret_variables: list[str] | None = None, + parent_key: str | None = None, +) -> Any: ... + + +def encrypt_secret_keys( + obj: Any, + secret_variables: list[str] | None = None, + parent_key: str | None = None, +) -> Any: + """ + Recursively obfuscate the value if it belongs to a Secret Variable. + Preserves input type: dict -> dict, list -> list, scalar -> scalar. + """ + secret_variables = secret_variables or [] + + if isinstance(obj, dict): + # recurse into dict + return {key: encrypt_secret_keys(value, secret_variables, key) for key, value in obj.items()} + + elif isinstance(obj, list): + # recurse into all list elements + return [encrypt_secret_keys(value, secret_variables, None) for value in obj] + + else: + # leaf node: obfuscate if parent_key is a secret variable + if parent_key in secret_variables: + return obfuscated_token(str(obj)) + return obj diff --git a/api/core/workflow/nodes/agent/agent_node.py b/api/core/workflow/nodes/agent/agent_node.py index a01686a4b8..c7643854e6 100644 --- a/api/core/workflow/nodes/agent/agent_node.py +++ b/api/core/workflow/nodes/agent/agent_node.py @@ -10,6 +10,7 @@ from sqlalchemy.orm import Session from core.agent.entities import AgentToolEntity from core.agent.plugin_entities import AgentStrategyParameter from core.file import File, FileTransferMethod +from core.helper import encrypter from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance, ModelManager from core.model_runtime.entities.llm_entities import LLMUsage, LLMUsageMetadata @@ -24,6 +25,7 @@ from core.tools.entities.tool_entities import ( ) from core.tools.tool_manager import ToolManager from core.tools.utils.message_transformer import ToolFileMessageTransformer +from core.variables import SecretVariable from core.variables.segments import ArrayFileSegment, StringSegment from core.workflow.entities import VariablePool from core.workflow.enums import ( @@ -139,6 +141,15 @@ class AgentNode(Node): # get conversation id conversation_id = self.graph_runtime_state.variable_pool.get(["sys", SystemVariableKey.CONVERSATION_ID]) + env_vars = self.graph_runtime_state.variable_pool.variable_dictionary.get("env", {}) + + # get secret variables used + secret_variables = [ + var.name + for var in env_vars.values() # iterate over the values directly + if isinstance(var, SecretVariable) + ] + try: message_stream = strategy.invoke( params=parameters, @@ -171,6 +182,7 @@ class AgentNode(Node): node_type=self.node_type, node_id=self._node_id, node_execution_id=self.id, + secret_variables=secret_variables, ) except PluginDaemonClientSideError as e: transform_error = AgentMessageTransformError( @@ -488,6 +500,7 @@ class AgentNode(Node): node_type: NodeType, node_id: str, node_execution_id: str, + secret_variables: list[str] | None = None, ) -> Generator[NodeEventBase, None, None]: """ Convert ToolInvokeMessages into tuple[plain_text, files] @@ -671,9 +684,9 @@ class AgentNode(Node): parent_id=message.message.parent_id, error=message.message.error, status=message.message.status.value, - data=message.message.data, + data=encrypter.encrypt_secret_keys(message.message.data, secret_variables, None), label=message.message.label, - metadata=message.message.metadata, + metadata=encrypter.encrypt_secret_keys(message.message.metadata, secret_variables, None), node_id=node_id, ) diff --git a/api/core/workflow/nodes/code/code_node.py b/api/core/workflow/nodes/code/code_node.py index c87cbf9628..a5ef001c2f 100644 --- a/api/core/workflow/nodes/code/code_node.py +++ b/api/core/workflow/nodes/code/code_node.py @@ -3,10 +3,12 @@ from decimal import Decimal from typing import Any, cast from configs import dify_config +from core.helper import encrypter from core.helper.code_executor.code_executor import CodeExecutionError, CodeExecutor, CodeLanguage from core.helper.code_executor.code_node_provider import CodeNodeProvider from core.helper.code_executor.javascript.javascript_code_provider import JavascriptCodeProvider from core.helper.code_executor.python3.python3_code_provider import Python3CodeProvider +from core.variables import SecretVariable from core.variables.segments import ArrayFileSegment from core.variables.types import SegmentType from core.workflow.enums import ErrorStrategy, NodeType, WorkflowNodeExecutionStatus @@ -73,15 +75,24 @@ class CodeNode(Node): code_language = self._node_data.code_language code = self._node_data.code + # to store secret variables used in the code block. + secret_variables = [] + # Get variables variables = {} for variable_selector in self._node_data.variables: variable_name = variable_selector.variable variable = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector) + if isinstance(variable, SecretVariable): + secret_variables.append(variable_name) + if isinstance(variable, ArrayFileSegment): variables[variable_name] = [v.to_dict() for v in variable.value] if variable.value else None else: variables[variable_name] = variable.to_object() if variable else None + + obfuscated_variables = encrypter.encrypt_secret_keys(variables, secret_variables, None) + # Run code try: result = CodeExecutor.execute_workflow_code_template( @@ -94,10 +105,13 @@ class CodeNode(Node): result = self._transform_result(result=result, output_schema=self._node_data.outputs) except (CodeExecutionError, CodeNodeError) as e: return NodeRunResult( - status=WorkflowNodeExecutionStatus.FAILED, inputs=variables, error=str(e), error_type=type(e).__name__ + status=WorkflowNodeExecutionStatus.FAILED, + inputs=obfuscated_variables, + error=str(e), + error_type=type(e).__name__, ) - return NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, inputs=variables, outputs=result) + return NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED, inputs=obfuscated_variables, outputs=result) def _check_string(self, value: str | None, variable: str) -> str | None: """