From dfc1c12428f2db582896c6f9132c595689c46a76 Mon Sep 17 00:00:00 2001
From: fit2cloud-chenyw <yawen.chen@fit2cloud.com>
Date: Thu, 28 Aug 2025 15:59:47 +0800
Subject: [PATCH] perf: Embedded Data Source Interface Supports AES Encryption

---
 backend/apps/system/api/assistant.py  |  2 ++
 backend/apps/system/crud/assistant.py | 23 ++++++++++++++++-------
 backend/common/utils/aes_crypto.py    | 16 ++++++++++++++++
 3 files changed, 34 insertions(+), 7 deletions(-)
 create mode 100644 backend/common/utils/aes_crypto.py

diff --git a/backend/apps/system/api/assistant.py b/backend/apps/system/api/assistant.py
index c765b8dc..b64f56a1 100644
--- a/backend/apps/system/api/assistant.py
+++ b/backend/apps/system/api/assistant.py
@@ -31,6 +31,8 @@ async def info(request: Request, response: Response, session: SessionDep, trans:
     db_model = AssistantModel.model_validate(db_model)
     response.headers["Access-Control-Allow-Origin"] = db_model.domain
     origin = request.headers.get("origin") or get_origin_from_referer(request)
+    if not origin:
+        raise RuntimeError(trans('i18n_embedded.invalid_origin', origin = origin or ''))
     origin = origin.rstrip('/')
     if origin != db_model.domain:
         raise RuntimeError(trans('i18n_embedded.invalid_origin', origin = origin or ''))
diff --git a/backend/apps/system/crud/assistant.py b/backend/apps/system/crud/assistant.py
index 7b2c6fd8..1a49f28a 100644
--- a/backend/apps/system/crud/assistant.py
+++ b/backend/apps/system/crud/assistant.py
@@ -15,9 +15,9 @@ from apps.system.schemas.system_schema import AssistantHeader, AssistantOutDsSch
 from common.core.config import settings
 from common.core.db import engine
 from common.core.sqlbot_cache import cache
+from common.utils.aes_crypto import simple_aes_decrypt
 from common.utils.utils import string_to_numeric_hash
 
-
 @cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="assistant_id")
 async def get_assistant_info(*, session: Session, assistant_id: int) -> AssistantModel | None:
     db_model = session.get(AssistantModel, assistant_id)
@@ -117,13 +117,13 @@ class AssistantOutDs:
         res = requests.get(url=endpoint, params=param, headers=header, cookies=cookies, timeout=10)
         if res.status_code == 200:
             result_json: dict[any] = json.loads(res.text)
-            if result_json.get('code') == 0:
+            if result_json.get('code') == 0 or result_json.get('code') == 200:
                 temp_list = result_json.get('data', [])
-                self.ds_list = [
-                    self.convert2schema(item)
+                temp_ds_list = [
+                    self.convert2schema(item, config)
                     for item in temp_list
                 ]
-
+                self.ds_list = temp_ds_list
                 return self.ds_list
             else:
                 raise Exception(f"Failed to get datasource list from {endpoint}, error: {result_json.get('message')}")
@@ -169,9 +169,19 @@ class AssistantOutDs:
             raise Exception("Datasource list is not found.")
         raise Exception(f"Datasource with id {ds_id} not found.")
 
-    def convert2schema(self, ds_dict: dict) -> AssistantOutDsSchema:
+    def convert2schema(self, ds_dict: dict, config: dict[any]) -> AssistantOutDsSchema:
         id_marker: str = ''
         attr_list = ['name', 'type', 'host', 'port', 'user', 'dataBase', 'schema']
+        if config.get('encrypt', True):
+            key = config.get('aes_key', None)
+            iv = config.get('aes_iv', None)
+            aes_attrs = ['host', 'user', 'password', 'dataBase', 'db_schema']
+            for attr in aes_attrs:
+                if attr in ds_dict and ds_dict[attr]:
+                    try:
+                        ds_dict[attr] = simple_aes_decrypt(ds_dict[attr], key, iv)
+                    except Exception as e:
+                        raise Exception(f"Failed to encrypt {attr} for datasource {ds_dict.get('name')}, error: {str(e)}")
         for attr in attr_list:
             if attr in ds_dict:
                 id_marker += str(ds_dict.get(attr, '')) + '--sqlbot--'
@@ -180,7 +190,6 @@ class AssistantOutDs:
         ds_dict.pop("schema", None)
         return AssistantOutDsSchema(**{**ds_dict, "id": id, "db_schema": db_schema})
 
-
 class AssistantOutDsFactory:
     @staticmethod
     def get_instance(assistant: AssistantHeader) -> AssistantOutDs:
diff --git a/backend/common/utils/aes_crypto.py b/backend/common/utils/aes_crypto.py
new file mode 100644
index 00000000..e75ab5a6
--- /dev/null
+++ b/backend/common/utils/aes_crypto.py
@@ -0,0 +1,16 @@
+from typing import Optional
+from common.core.config import settings
+from sqlbot_xpack.aes_utils import SecureEncryption
+
+simple_aes_iv_text = 'sqlbot_em_aes_iv'
+def sqlbot_aes_encrypt(text: str, key: Optional[str] = None) -> str:
+    return SecureEncryption.encrypt_to_single_string(text, key or settings.SECRET_KEY)
+
+def sqlbot_aes_decrypt(text: str, key: Optional[str] = None) -> str:
+    return SecureEncryption.decrypt_from_single_string(text, key or settings.SECRET_KEY)
+
+def simple_aes_encrypt(text: str, key: Optional[str] = None, ivtext: Optional[str] = None) -> str:
+    return SecureEncryption.simple_aes_encrypt(text, key or settings.SECRET_KEY[:32], ivtext or simple_aes_iv_text)
+
+def simple_aes_decrypt(text: str, key: Optional[str] = None, ivtext: Optional[str] = None) -> str:
+    return SecureEncryption.simple_aes_decrypt(text, key or settings.SECRET_KEY[:32], ivtext or simple_aes_iv_text)
\ No newline at end of file