mirror of
https://github.com/YFGaia/dify-plus.git
synced 2026-06-12 18:11:42 +08:00
feat: enable dsl export encrypt dataset id or not (#25102)
Signed-off-by: kenwoodjw <blackxin55+@gmail.com>
This commit is contained in:
@@ -17,6 +17,7 @@ from pydantic import BaseModel, Field
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from configs import dify_config
|
||||
from core.helper import ssrf_proxy
|
||||
from core.model_runtime.utils.encoders import jsonable_encoder
|
||||
from core.plugin.entities.plugin import PluginDependency
|
||||
@@ -786,7 +787,10 @@ class AppDslService:
|
||||
|
||||
@classmethod
|
||||
def encrypt_dataset_id(cls, dataset_id: str, tenant_id: str) -> str:
|
||||
"""Encrypt dataset_id using AES-CBC mode"""
|
||||
"""Encrypt dataset_id using AES-CBC mode or return plain text based on configuration"""
|
||||
if not dify_config.DSL_EXPORT_ENCRYPT_DATASET_ID:
|
||||
return dataset_id
|
||||
|
||||
key = cls._generate_aes_key(tenant_id)
|
||||
iv = key[:16]
|
||||
cipher = AES.new(key, AES.MODE_CBC, iv)
|
||||
@@ -795,12 +799,34 @@ class AppDslService:
|
||||
|
||||
@classmethod
|
||||
def decrypt_dataset_id(cls, encrypted_data: str, tenant_id: str) -> str | None:
|
||||
"""AES decryption"""
|
||||
"""AES decryption with fallback to plain text UUID"""
|
||||
# First, check if it's already a plain UUID (not encrypted)
|
||||
if cls._is_valid_uuid(encrypted_data):
|
||||
return encrypted_data
|
||||
|
||||
# If it's not a UUID, try to decrypt it
|
||||
try:
|
||||
key = cls._generate_aes_key(tenant_id)
|
||||
iv = key[:16]
|
||||
cipher = AES.new(key, AES.MODE_CBC, iv)
|
||||
pt = unpad(cipher.decrypt(base64.b64decode(encrypted_data)), AES.block_size)
|
||||
return pt.decode()
|
||||
decrypted_text = pt.decode()
|
||||
|
||||
# Validate that the decrypted result is a valid UUID
|
||||
if cls._is_valid_uuid(decrypted_text):
|
||||
return decrypted_text
|
||||
else:
|
||||
# If decrypted result is not a valid UUID, it's probably not our encrypted data
|
||||
return None
|
||||
except Exception:
|
||||
# If decryption fails completely, return None
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _is_valid_uuid(value: str) -> bool:
|
||||
"""Check if string is a valid UUID format"""
|
||||
try:
|
||||
uuid.UUID(value)
|
||||
return True
|
||||
except (ValueError, TypeError):
|
||||
return False
|
||||
|
||||
Reference in New Issue
Block a user