feat: inner api encrypt

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

View File

@ -15,7 +15,7 @@ from core.tools.entities.tool_entities import (
from core.tools.provider.api_tool_provider import ApiToolProviderController
from core.tools.tool_label_manager import ToolLabelManager
from core.tools.tool_manager import ToolManager
from core.tools.utils.configuration import ToolConfigurationManager
from core.tools.utils.configuration import ProviderConfigEncrypter
from core.tools.utils.parser import ApiBasedToolSchemaParser
from extensions.ext_database import db
from models.tools import ApiToolProvider
@ -156,14 +156,14 @@ class ApiToolManageService:
provider_controller.load_bundled_tools(tool_bundles)
# encrypt credentials
tool_configuration = ToolConfigurationManager(
tool_configuration = ProviderConfigEncrypter(
tenant_id=tenant_id,
config=provider_controller.get_credentials_schema(),
provider_type=provider_controller.provider_type.value,
provider_identity=provider_controller.identity.name
)
encrypted_credentials = tool_configuration.encrypt_tool_credentials(credentials)
encrypted_credentials = tool_configuration.encrypt(credentials)
db_provider.credentials_str = json.dumps(encrypted_credentials)
db.session.add(db_provider)
@ -286,21 +286,21 @@ class ApiToolManageService:
provider_controller.load_bundled_tools(tool_bundles)
# get original credentials if exists
tool_configuration = ToolConfigurationManager(
tool_configuration = ProviderConfigEncrypter(
tenant_id=tenant_id,
config=provider_controller.get_credentials_schema(),
provider_type=provider_controller.provider_type.value,
provider_identity=provider_controller.identity.name
)
original_credentials = tool_configuration.decrypt_tool_credentials(provider.credentials)
original_credentials = tool_configuration.decrypt(provider.credentials)
masked_credentials = tool_configuration.mask_tool_credentials(original_credentials)
# check if the credential has changed, save the original credential
for name, value in credentials.items():
if name in masked_credentials and value == masked_credentials[name]:
credentials[name] = original_credentials[name]
credentials = tool_configuration.encrypt_tool_credentials(credentials)
credentials = tool_configuration.encrypt(credentials)
provider.credentials_str = json.dumps(credentials)
db.session.add(provider)
@ -405,13 +405,13 @@ class ApiToolManageService:
# decrypt credentials
if db_provider.id:
tool_configuration = ToolConfigurationManager(
tool_configuration = ProviderConfigEncrypter(
tenant_id=tenant_id,
config=provider_controller.get_credentials_schema(),
provider_type=provider_controller.provider_type.value,
provider_identity=provider_controller.identity.name
)
decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials)
decrypted_credentials = tool_configuration.decrypt(credentials)
# check if the credential has changed, save the original credential
masked_credentials = tool_configuration.mask_tool_credentials(decrypted_credentials)
for name, value in credentials.items():

View File

@ -10,7 +10,7 @@ from core.tools.provider.builtin._positions import BuiltinToolProviderSort
from core.tools.provider.tool_provider import ToolProviderController
from core.tools.tool_label_manager import ToolLabelManager
from core.tools.tool_manager import ToolManager
from core.tools.utils.configuration import ToolConfigurationManager
from core.tools.utils.configuration import ProviderConfigEncrypter
from extensions.ext_database import db
from models.tools import BuiltinToolProvider
from services.tools.tools_transform_service import ToolTransformService
@ -27,7 +27,7 @@ class BuiltinToolManageService:
provider_controller: ToolProviderController = ToolManager.get_builtin_provider(provider)
tools = provider_controller.get_tools()
tool_provider_configurations = ToolConfigurationManager(
tool_provider_configurations = ProviderConfigEncrypter(
tenant_id=tenant_id,
config=provider_controller.get_credentials_schema(),
provider_type=provider_controller.provider_type.value,
@ -47,7 +47,7 @@ class BuiltinToolManageService:
if builtin_provider is not None:
# get credentials
credentials = builtin_provider.credentials
credentials = tool_provider_configurations.decrypt_tool_credentials(credentials)
credentials = tool_provider_configurations.decrypt(credentials)
result = []
for tool in tools:
@ -92,7 +92,7 @@ class BuiltinToolManageService:
provider_controller = ToolManager.get_builtin_provider(provider_name)
if not provider_controller.need_credentials:
raise ValueError(f"provider {provider_name} does not need credentials")
tool_configuration = ToolConfigurationManager(
tool_configuration = ProviderConfigEncrypter(
tenant_id=tenant_id,
config=provider_controller.get_credentials_schema(),
provider_type=provider_controller.provider_type.value,
@ -101,7 +101,7 @@ class BuiltinToolManageService:
# get original credentials if exists
if provider is not None:
original_credentials = tool_configuration.decrypt_tool_credentials(provider.credentials)
original_credentials = tool_configuration.decrypt(provider.credentials)
masked_credentials = tool_configuration.mask_tool_credentials(original_credentials)
# check if the credential has changed, save the original credential
for name, value in credentials.items():
@ -110,7 +110,7 @@ class BuiltinToolManageService:
# validate credentials
provider_controller.validate_credentials(credentials)
# encrypt credentials
credentials = tool_configuration.encrypt_tool_credentials(credentials)
credentials = tool_configuration.encrypt(credentials)
except (ToolProviderNotFoundError, ToolNotFoundError, ToolProviderCredentialValidationError) as e:
raise ValueError(str(e))
@ -154,13 +154,13 @@ class BuiltinToolManageService:
return {}
provider_controller = ToolManager.get_builtin_provider(provider_obj.provider)
tool_configuration = ToolConfigurationManager(
tool_configuration = ProviderConfigEncrypter(
tenant_id=tenant_id,
config=provider_controller.get_credentials_schema(),
provider_type=provider_controller.provider_type.value,
provider_identity=provider_controller.identity.name,
)
credentials = tool_configuration.decrypt_tool_credentials(provider_obj.credentials)
credentials = tool_configuration.decrypt(provider_obj.credentials)
credentials = tool_configuration.mask_tool_credentials(credentials)
return credentials
@ -186,7 +186,7 @@ class BuiltinToolManageService:
# delete cache
provider_controller = ToolManager.get_builtin_provider(provider_name)
tool_configuration = ToolConfigurationManager(
tool_configuration = ProviderConfigEncrypter(
tenant_id=tenant_id,
config=provider_controller.get_credentials_schema(),
provider_type=provider_controller.provider_type.value,

View File

@ -16,7 +16,7 @@ from core.tools.provider.builtin_tool_provider import BuiltinToolProviderControl
from core.tools.provider.workflow_tool_provider import WorkflowToolProviderController
from core.tools.tool.tool import Tool
from core.tools.tool.workflow_tool import WorkflowTool
from core.tools.utils.configuration import ToolConfigurationManager
from core.tools.utils.configuration import ProviderConfigEncrypter
from models.tools import ApiToolProvider, BuiltinToolProvider, WorkflowToolProvider
logger = logging.getLogger(__name__)
@ -107,15 +107,15 @@ class ToolTransformService:
credentials = db_provider.credentials
# init tool configuration
tool_configuration = ToolConfigurationManager(
tool_configuration = ProviderConfigEncrypter(
tenant_id=db_provider.tenant_id,
config=provider_controller.get_credentials_schema(),
provider_type=provider_controller.provider_type.value,
provider_identity=provider_controller.identity.name
)
# decrypt the credentials and mask the credentials
decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials=credentials)
masked_credentials = tool_configuration.mask_tool_credentials(credentials=decrypted_credentials)
decrypted_credentials = tool_configuration.decrypt(data=credentials)
masked_credentials = tool_configuration.mask_tool_credentials(data=decrypted_credentials)
result.masked_credentials = masked_credentials
result.original_credentials = decrypted_credentials
@ -218,7 +218,7 @@ class ToolTransformService:
if decrypt_credentials:
# init tool configuration
tool_configuration = ToolConfigurationManager(
tool_configuration = ProviderConfigEncrypter(
tenant_id=db_provider.tenant_id,
config=provider_controller.get_credentials_schema(),
provider_type=provider_controller.provider_type.value,
@ -226,8 +226,8 @@ class ToolTransformService:
)
# decrypt the credentials and mask the credentials
decrypted_credentials = tool_configuration.decrypt_tool_credentials(credentials=credentials)
masked_credentials = tool_configuration.mask_tool_credentials(credentials=decrypted_credentials)
decrypted_credentials = tool_configuration.decrypt(data=credentials)
masked_credentials = tool_configuration.mask_tool_credentials(data=decrypted_credentials)
result.masked_credentials = masked_credentials