feat: datasource config add aes encrypt

Author: XiaJunjie2020

backend/apps/datasource/api/datasource.py

@@ -10,8 +10,8 @@ def datasource_list(session: SessionDep):
     return get_datasource_list(session=session)
 
 @router.post("/check")
-def check(session: SessionDep, conf: DatasourceConf):
-    return check_status(session, conf)
+def check(session: SessionDep, ds: CoreDatasource):
+    return check_status(session, ds)
 
 @router.post("/add", response_model=CoreDatasource)
 def add(session: SessionDep, ds: CoreDatasource):

backend/apps/datasource/crud/datasource.py

@@ -3,13 +3,17 @@
 import pyodbc
 import datetime
 from common.core.deps import SessionDep
+import json
+from ..utils.utils import aes_decrypt
+
 
 def get_datasource_list(session: SessionDep) -> CoreDatasource:
     statement = select(CoreDatasource).order_by(CoreDatasource.create_time.desc())
     datasource_list = session.exec(statement).fetchall()
     return datasource_list
 
-def check_status(session: SessionDep, conf: DatasourceConf):
+def check_status(session: SessionDep, ds: CoreDatasource):
+    conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
     conn_str = (
         "DRIVER={MySQL ODBC 9.3 Unicode Driver};" # todo driver config
         f"SERVER={conf.host};"
@@ -32,20 +36,20 @@ def check_status(session: SessionDep, conf: DatasourceConf):
             conn.close()
 
 def create_ds(session: SessionDep, ds: CoreDatasource):
+    ds.create_time = datetime.datetime.now()
+    ds.status = "Success"  # todo check status
     record = CoreDatasource(**ds.model_dump())
-    record.create_time = datetime.datetime.now()
-    record.status = "Success" # todo check status
     session.add(record)
     session.commit()
-    return record
+    return ds
 
 def update_ds(session: SessionDep, ds: CoreDatasource):
     ds.id = int(ds.id)
-    term = session.exec(select(CoreDatasource).where(CoreDatasource.id == ds.id)).first()
+    record = session.exec(select(CoreDatasource).where(CoreDatasource.id == ds.id)).first()
     update_data = ds.model_dump(exclude_unset=True)
     for field, value in update_data.items():
-        setattr(term, field, value)
-    session.add(term)
+        setattr(record, field, value)
+    session.add(record)
     session.commit()
     return ds
 

backend/apps/datasource/utils/__init__.py

@@ -0,0 +1,19 @@
+from Crypto.Cipher import AES
+from Crypto.Util.Padding import unpad, pad
+import base64
+
+key = b'SQLBot1234567890'
+
+def aes_encrypt(data):
+    data = bytes(data,'utf-8')
+    cipher = AES.new(key, AES.MODE_ECB)
+    data = pad(data, AES.block_size)
+    encrypt = cipher.encrypt(data)
+    return base64.b64encode(encrypt)
+
+def aes_decrypt(encrypted_data):
+    encrypted_data = base64.b64decode(encrypted_data)
+    cipher = AES.new(key, AES.MODE_ECB)
+    text = cipher.decrypt(encrypted_data)
+    decrypted_text = unpad(text, AES.block_size)
+    return decrypted_text.decode('utf-8')

backend/apps/datasource/utils/utils.py

@@ -19,6 +19,7 @@ dependencies = [
     "sentry-sdk[fastapi]<2.0.0,>=1.40.6",
     "pyjwt<3.0.0,>=2.8.0",
     "pyodbc (>=5.2.0,<5.2.1)",
+    "pycryptodome (>=3.22.0,<4.0.0)",
 ]
 [[tool.uv.index]]
 url = "https://pypi.tuna.tsinghua.edu.cn/simple"

backend/pyproject.toml

@@ -19,6 +19,7 @@
     "@vitejs/plugin-vue": "^5.2.2",
     "@vue/tsconfig": "^0.7.0",
     "axios": "^1.8.4",
+    "crypto-js": "^4.2.0",
     "element-plus": "^2.9.7",
     "less": "^4.3.0",
     "pinia": "^3.0.2",

frontend/package.json

@@ -22,19 +22,19 @@
         </el-select>
       </el-form-item>
       <el-form-item label="Host/Ip">
-        <el-input v-model="form.host" />
+        <el-input v-model="config.host" />
       </el-form-item>
       <el-form-item label="Port">
-        <el-input v-model="form.port" />
+        <el-input v-model="config.port" />
       </el-form-item>
       <el-form-item label="Username">
-        <el-input v-model="form.username" />
+        <el-input v-model="config.username" />
       </el-form-item>
       <el-form-item label="Password">
-        <el-input v-model="form.password" />
+        <el-input v-model="config.password" />
       </el-form-item>
       <el-form-item label="Database">
-        <el-input v-model="form.database" />
+        <el-input v-model="config.database" />
       </el-form-item>
 
       <div style="display: flex;justify-content: flex-end;">
@@ -48,6 +48,7 @@
 <script lang="ts" setup>
 import { ref } from 'vue'
 import { datasourceApi } from '@/api/datasource'
+import { encrypted, decrypted } from './js/aes'
 
 const emit = defineEmits(['refresh'])
 
@@ -59,13 +60,15 @@ const form = ref<any>({
   name:'',
   description:'',
   type:'mysql',
+  configuration: ''
+})
+const config = ref<any>({
   driver:'',
   host:'',
   port:0,
   username:'',
   password:'',
   database:'',
-  configuration: ''
 })
 
 const close = () => {
@@ -78,24 +81,43 @@ const open = (item: any) => {
     form.value.name = item.name
     form.value.description = item.description
     form.value.type = item.type
-    const configuration = JSON.parse(item.configuration)
-    form.value.host = configuration.host
-    form.value.port = configuration.port
-    form.value.username = configuration.username
-    form.value.password = configuration.password
-    form.value.database = configuration.database
+    if(item.configuration) {
+      const configuration = JSON.parse(decrypted(item.configuration))
+      config.value.host = configuration.host
+      config.value.port = configuration.port
+      config.value.username = configuration.username
+      config.value.password = configuration.password
+      config.value.database = configuration.database
+    }
+  } else {
+    form.value = {
+      driver:'',
+      host:'',
+      port:0,
+      username:'',
+      password:'',
+      database:'',
+    }
+    config.value = {
+      driver:'',
+      host:'',
+      port:0,
+      username:'',
+      password:'',
+      database:'',
+    }
   }
   dialogVisible.value = true
 }
 
 const save = () => {
-  form.value.configuration = JSON.stringify({
-    host:form.value.host,
-    port:form.value.port,
-    username:form.value.username,
-    password:form.value.password,
-    database:form.value.database
-  })
+  form.value.configuration = encrypted(JSON.stringify({
+    host:config.value.host,
+    port:config.value.port,
+    username:config.value.username,
+    password:config.value.password,
+    database:config.value.database
+  }))
   if (form.value.id) {
     datasourceApi.update(form.value).then((res) => {
       console.log(res)
@@ -112,6 +134,13 @@ const save = () => {
 }
 
 const check = () => {
+  form.value.configuration = encrypted(JSON.stringify({
+    host:config.value.host,
+    port:config.value.port,
+    username:config.value.username,
+    password:config.value.password,
+    database:config.value.database
+  }))
   datasourceApi.check(form.value).then((res: any) => {
     console.log(res)
   })
@@ -120,4 +149,4 @@ const check = () => {
 defineExpose({ open })
 </script>
 <style lang="less" scoped>
-</style>
+</style>./js/aes

frontend/src/views/ds/form.vue

@@ -14,7 +14,7 @@
         </el-input>
       </div>
 
-      <el-button type="primary" :icon="IconOpeAdd" @click="editDs()">Add Datasource</el-button>
+      <el-button type="primary" :icon="IconOpeAdd" @click="editDs(undefined)">Add Datasource</el-button>
     </div>
 
     <div class="connections-container">
@@ -28,7 +28,7 @@
         <div class="connection-details">
           <div class="connection-name">{{ ds.name }}</div>
           <div class="connection-type">{{ ds.type }}</div>
-          <div class="connection-host">{{ ds.host }}</div>
+          <div class="connection-host">{{ ds.description }}</div>
           <div class="connection-last">{{ datetimeFormat(ds.create_time) }}</div>
         </div>
         <div class="connection-status" :class="`${getStatus(ds.status)}`">{{ ds.status }}</div>
@@ -82,7 +82,7 @@ const list = () => {
   })
 }
 
-const editDs = (item: any = {}) => {
+const editDs = (item: any) => {
   dsForm.value.open(item)
 }
 

frontend/src/views/ds/index.vue

@@ -0,0 +1,18 @@
+import CryptoJS from 'crypto-js'
+
+const key = CryptoJS.enc.Utf8.parse('SQLBot1234567890')
+
+export const encrypted = (str:string) => {
+  return CryptoJS.AES.encrypt(str, key, {
+    mode: CryptoJS.mode.ECB,
+    padding: CryptoJS.pad.Pkcs7
+  }).toString()
+}
+
+export const decrypted = (str:string) => {
+  const bytes = CryptoJS.AES.decrypt(str, key, {
+    mode: CryptoJS.mode.ECB,
+    padding: CryptoJS.pad.Pkcs7
+  })
+  return bytes.toString(CryptoJS.enc.Utf8)
+}