youngkingdom/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-05-02 16:38:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: inner api encrypt

Browse Source
This commit is contained in:
Yeuoly
2024-08-30 21:25:58 +08:00
parent 60e75dc748
commit de01ca8d55
8 changed files with 88 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
api/core/tools/tool_manager.py
View File

@ -24,7 +24,7 @@ from core.tools.tool.builtin_tool import BuiltinTool
from core.tools.tool.tool import Tool
from core.tools.tool.workflow_tool import WorkflowTool
from core.tools.tool_label_manager import ToolLabelManager
from core.tools.utils.configuration import ToolConfigurationManager, ToolParameterConfigurationManager
from core.tools.utils.configuration import ProviderConfigEncrypter, ToolParameterConfigurationManager
from core.tools.utils.tool_parameter_converter import ToolParameterConverter
from core.workflow.nodes.tool.entities import ToolEntity
from extensions.ext_database import db
@ -116,14 +116,14 @@ class ToolManager:
# decrypt the credentials
credentials = builtin_provider.credentials
controller = cls.get_builtin_provider(provider_id)
tool_configuration = ToolConfigurationManager(
tool_configuration = ProviderConfigEncrypter(
tenant_id=tenant_id,
config=controller.get_credentials_schema(),
provider_type=controller.provider_type.value,
provider_identity=controller.identity.name
)
decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials)
decrypted_credentials = tool_configuration.decrypt(credentials)
return cast(BuiltinTool, builtin_tool.fork_tool_runtime(runtime={
'tenant_id': tenant_id,
@ -140,13 +140,13 @@ class ToolManager:
api_provider, credentials = cls.get_api_provider_controller(tenant_id, provider_id)
# decrypt the credentials
tool_configuration = ToolConfigurationManager(
tool_configuration = ProviderConfigEncrypter(
tenant_id=tenant_id,
config=api_provider.get_credentials_schema(),
provider_type=api_provider.provider_type.value,
provider_identity=api_provider.identity.name
)
decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials)
decrypted_credentials = tool_configuration.decrypt(credentials)
return cast(ApiTool, api_provider.get_tool(tool_name).fork_tool_runtime(runtime={
'tenant_id': tenant_id,
@ -523,14 +523,14 @@ class ToolManager:
provider_obj, ApiProviderAuthType.API_KEY if credentials['auth_type'] == 'api_key' else ApiProviderAuthType.NONE
)
# init tool configuration
tool_configuration = ToolConfigurationManager(
tool_configuration = ProviderConfigEncrypter(
tenant_id=tenant_id,
config=controller.get_credentials_schema(),
provider_type=controller.provider_type.value,
provider_identity=controller.identity.name
)
decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials)
decrypted_credentials = tool_configuration.decrypt(credentials)
masked_credentials = tool_configuration.mask_tool_credentials(decrypted_credentials)
try:

52
api/core/tools/utils/configuration.py
View File

@ -15,60 +15,60 @@ from core.tools.entities.tool_entities import (
from core.tools.tool.tool import Tool
class ToolConfigurationManager(BaseModel):
class ProviderConfigEncrypter(BaseModel):
tenant_id: str
config: Mapping[str, BasicProviderConfig]
provider_type: str
provider_identity: str
def _deep_copy(self, credentials: dict[str, str]) -> dict[str, str]:
def _deep_copy(self, data: dict[str, str]) -> dict[str, str]:
"""
deep copy credentials
deep copy data
"""
return deepcopy(credentials)
return deepcopy(data)
def encrypt_tool_credentials(self, credentials: dict[str, str]) -> dict[str, str]:
def encrypt(self, data: dict[str, str]) -> Mapping[str, str]:
"""
encrypt tool credentials with tenant id
return a deep copy of credentials with encrypted values
"""
credentials = self._deep_copy(credentials)
data = self._deep_copy(data)
# get fields need to be decrypted
fields = self.config
for field_name, field in fields.items():
if field.type == BasicProviderConfig.Type.SECRET_INPUT:
if field_name in credentials:
encrypted = encrypter.encrypt_token(self.tenant_id, credentials[field_name])
credentials[field_name] = encrypted
if field_name in data:
encrypted = encrypter.encrypt_token(self.tenant_id, data[field_name])
data[field_name] = encrypted
return credentials
return data
def mask_tool_credentials(self, credentials: dict[str, Any]) -> dict[str, Any]:
def mask_tool_credentials(self, data: dict[str, Any]) -> Mapping[str, Any]:
"""
mask tool credentials
return a deep copy of credentials with masked values
"""
credentials = self._deep_copy(credentials)
data = self._deep_copy(data)
# get fields need to be decrypted
fields = self.config
for field_name, field in fields.items():
if field.type == BasicProviderConfig.Type.SECRET_INPUT:
if field_name in credentials:
if len(credentials[field_name]) > 6:
credentials[field_name] = \
credentials[field_name][:2] + \
'*' * (len(credentials[field_name]) - 4) + \
credentials[field_name][-2:]
if field_name in data:
if len(data[field_name]) > 6:
data[field_name] = \
data[field_name][:2] + \
'*' * (len(data[field_name]) - 4) + \
data[field_name][-2:]
else:
credentials[field_name] = '*' * len(credentials[field_name])
data[field_name] = '*' * len(data[field_name])
return credentials
return data
def decrypt_tool_credentials(self, credentials: dict[str, str]) -> dict[str, str]:
def decrypt(self, data: dict[str, str]) -> Mapping[str, str]:
"""
decrypt tool credentials with tenant id
@ -82,19 +82,19 @@ class ToolConfigurationManager(BaseModel):
cached_credentials = cache.get()
if cached_credentials:
return cached_credentials
credentials = self._deep_copy(credentials)
data = self._deep_copy(data)
# get fields need to be decrypted
fields = self.config
for field_name, field in fields.items():
if field.type == BasicProviderConfig.Type.SECRET_INPUT:
if field_name in credentials:
if field_name in data:
try:
credentials[field_name] = encrypter.decrypt_token(self.tenant_id, credentials[field_name])
data[field_name] = encrypter.decrypt_token(self.tenant_id, data[field_name])
except:
pass
cache.set(credentials)
return credentials
cache.set(data)
return data
def delete_tool_credentials_cache(self):
cache = ToolProviderCredentialsCache(