♻️ 重构本地数据库-credentials模块

server/app/controller/ssh.js

@@ -1,11 +1,11 @@
-const { readSSHRecord, writeSSHRecord } = require('../utils/storage')
 const { RSADecryptAsync, AESEncryptAsync, AESDecryptAsync } = require('../utils/encrypt')
-const { HostListDB } = require('../utils/db-class')
+const { HostListDB, CredentialsDB } = require('../utils/db-class')
 const hostListDB = new HostListDB().getInstance()
+const credentialsDB = new CredentialsDB().getInstance()
 
 async function getSSHList({ res }) {
   // console.log('get-host-list')
-  let data = await readSSHRecord()
+  let data = await credentialsDB.findAsync({})
   data = data?.map(item => {
     const { name, authType, _id: id, date } = item
     return { id, name, authType, privateKey: '', password: '', date }
@@ -18,8 +18,8 @@ const addSSH = async ({ res, request }) => {
   let { body: { name, authType, password, privateKey, tempKey } } = request
   let record = { name, authType, password, privateKey }
   if (!name || !record[authType]) return res.fail({ data: false, msg: '参数错误' })
-  let sshRecord = await readSSHRecord()
+  let count = await credentialsDB.countAsync({ name })
-  if (sshRecord.some(item => item.name === name)) return res.fail({ data: false, msg: '已存在同名凭证' })
+  if (count > 0) return res.fail({ data: false, msg: '已存在同名凭证' })
 
   const clearTempKey = await RSADecryptAsync(tempKey)
   console.log('clearTempKey:', clearTempKey)
@@ -27,9 +27,7 @@ const addSSH = async ({ res, request }) => {
   // console.log(`${ authType }原密文: `, clearSSHKey)
   record[authType] = await AESEncryptAsync(clearSSHKey)
   // console.log(`${ authType }__commonKey加密存储: `, record[authType])
-
+  await credentialsDB.insertAsync({ ...record, date: Date.now() })
-  sshRecord.push({ ...record, date: Date.now() })
-  await writeSSHRecord(sshRecord)
   consola.info('添加凭证：', name)
   res.success({ data: '保存成功' })
 }
@@ -38,11 +36,8 @@ const updateSSH = async ({ res, request }) => {
   let { body: { id, name, authType, password, privateKey, date, tempKey } } = request
   let record = { name, authType, password, privateKey, date }
   if (!id || !name) return res.fail({ data: false, msg: '请输入凭据名称' })
-  let sshRecord = await readSSHRecord()
+  let oldRecord = await credentialsDB.findOneAsync({ _id: id })
-  let idx = sshRecord.findIndex(item => item._id === id)
+  if (!oldRecord) return res.fail({ data: false, msg: '凭证不存在' })
-  if (sshRecord.some(item => item.name === name && item.date !== date)) return res.fail({ data: false, msg: '已存在同名凭证' })
-  if (idx === -1) res.fail({ data: false, msg: '请输入凭据名称' })
-  const oldRecord = sshRecord[idx]
   // 判断原记录是否存在当前更新记录的认证方式
   if (!oldRecord[authType] && !record[authType]) return res.fail({ data: false, msg: `请输入${ authType === 'password' ? '密码' : '密钥' }` })
   if (!record[authType] && oldRecord[authType]) {
@@ -55,31 +50,30 @@ const updateSSH = async ({ res, request }) => {
     record[authType] = await AESEncryptAsync(clearSSHKey)
     // console.log(`${ authType }__commonKey加密存储: `, record[authType])
   }
-  record._id = sshRecord[idx]._id
+  await credentialsDB.updateAsync({ _id: id }, record)
-  sshRecord.splice(idx, 1, record)
-  await writeSSHRecord(sshRecord)
   consola.info('修改凭证：', name)
   res.success({ data: '保存成功' })
 }
 
 const removeSSH = async ({ res, request }) => {
   let { params: { id } } = request
-  let sshRecord = await readSSHRecord()
+  let count = await credentialsDB.countAsync({ _id: id })
-  let idx = sshRecord.findIndex(item => item._id === id)
+  if (count === 0) return res.fail({ msg: '凭证不存在' })
-  if (idx === -1) return res.fail({ msg: '凭证不存在' })
-  sshRecord.splice(idx, 1)
   // 将删除的凭证id从host中删除
   let hostList = await hostListDB.findAsync({})
   if (Array.isArray(hostList) && hostList.length > 0) {
-    for (let item of hostList) {
+    for (let host of hostList) {
-      if (item.credential === id) {
+      let { credential } = host
-        item.credential = ''
+      credential = await AESDecryptAsync(credential)
-        await hostListDB.updateAsync({ _id: item._id }, item)
+      if (credential === id) {
+        host.credential = ''
+        await hostListDB.updateAsync({ _id: host._id }, host)
       }
     }
   }
+  await hostListDB.compactDatafileAsync()
   consola.info('移除凭证：', id)
-  await writeSSHRecord(sshRecord)
+  await credentialsDB.removeAsync({ _id: id })
   res.success({ data: '移除成功' })
 }
 

server/app/socket/onekey.js

@@ -1,13 +1,13 @@
 const { Server } = require('socket.io')
 const { Client: SSHClient } = require('ssh2')
 const { sendNoticeAsync } = require('../utils/notify')
-const { readSSHRecord } = require('../utils/storage')
 const { verifyAuthSync } = require('../utils/verify-auth')
 const { shellThrottle } = require('../utils/tools')
 const { AESDecryptAsync } = require('../utils/encrypt')
 const { isAllowedIp } = require('../utils/tools')
-const { HostListDB, OnekeyDB } = require('../utils/db-class')
+const { HostListDB, CredentialsDB, OnekeyDB } = require('../utils/db-class')
 const hostListDB = new HostListDB().getInstance()
+const credentialsDB = new CredentialsDB().getInstance()
 const onekeyDB = new OnekeyDB().getInstance()
 
 const execStatusEnum = {
@@ -149,7 +149,7 @@ module.exports = (httpServer) => {
           try {
             if (authType === 'credential') {
               let credentialId = await AESDecryptAsync(hostInfo['credential'])
-              const sshRecordList = await readSSHRecord()
+              const sshRecordList = await credentialsDB.findAsync({})
               const sshRecord = sshRecordList.find(item => item._id === credentialId)
               authInfo.authType = sshRecord.authType
               authInfo[authInfo.authType] = await AESDecryptAsync(sshRecord[authInfo.authType])

server/app/socket/sftp.js

@@ -6,10 +6,10 @@ const { Server } = require('socket.io')
 const { sftpCacheDir } = require('../config')
 const { verifyAuthSync } = require('../utils/verify-auth')
 const { AESDecryptAsync } = require('../utils/encrypt')
-const { readSSHRecord } = require('../utils/storage')
 const { isAllowedIp } = require('../utils/tools')
-const { HostListDB } = require('../utils/db-class')
+const { HostListDB, CredentialsDB } = require('../utils/db-class')
 const hostListDB = new HostListDB().getInstance()
+const credentialsDB = new CredentialsDB().getInstance()
 
 // 读取切片
 const pipeStream = (path, writeStream) => {
@@ -241,7 +241,7 @@ module.exports = (httpServer) => {
       // 解密放到try里面，防止报错【commonKey必须配对, 否则需要重新添加服务器密钥】
       if (authType === 'credential') {
         let credentialId = await AESDecryptAsync(targetHostInfo[authType])
-        const sshRecordList = await readSSHRecord()
+        const sshRecordList = await credentialsDB.findAsync({})
         const sshRecord = sshRecordList.find(item => item._id === credentialId)
         authInfo.authType = sshRecord.authType
         authInfo[authInfo.authType] = await AESDecryptAsync(sshRecord[authInfo.authType])

server/app/socket/terminal.js

@@ -2,11 +2,11 @@ const { Server } = require('socket.io')
 const { Client: SSHClient } = require('ssh2')
 const { verifyAuthSync } = require('../utils/verify-auth')
 const { AESDecryptAsync } = require('../utils/encrypt')
-const { readSSHRecord } = require('../utils/storage')
 const { sendNoticeAsync } = require('../utils/notify')
 const { isAllowedIp, ping } = require('../utils/tools')
-const { HostListDB } = require('../utils/db-class')
+const { HostListDB, CredentialsDB } = require('../utils/db-class')
 const hostListDB = new HostListDB().getInstance()
+const credentialsDB = new CredentialsDB().getInstance()
 
 function createInteractiveShell(socket, sshClient) {
   return new Promise((resolve) => {
@@ -63,8 +63,7 @@ async function createTerminal(hostId, socket, sshClient) {
     // 解密放到try里面，防止报错【commonKey必须配对, 否则需要重新添加服务器密钥】
       if (authType === 'credential') {
         let credentialId = await AESDecryptAsync(targetHostInfo[authType])
-        const sshRecordList = await readSSHRecord()
+        const sshRecord = await credentialsDB.findOneAsync({ _id: credentialId })
-        const sshRecord = sshRecordList.find(item => item._id === credentialId)
         authInfo.authType = sshRecord.authType
         authInfo[authInfo.authType] = await AESDecryptAsync(sshRecord[authInfo.authType])
       } else {

server/app/utils/db-class.js

@@ -35,15 +35,15 @@ module.exports.HostListDB = class HostListDB {
   }
 }
 
-module.exports.SshRecordDB = class SshRecordDB {
+module.exports.CredentialsDB = class CredentialsDB {
   constructor() {
-    if (!SshRecordDB.instance) {
+    if (!CredentialsDB.instance) {
-      SshRecordDB.instance = new Datastore({ filename: credentialsDBPath, autoload: true })
+      CredentialsDB.instance = new Datastore({ filename: credentialsDBPath, autoload: true })
-      // SshRecordDB.instance.setAutocompactionInterval(5000)
+      // CredentialsDB.instance.setAutocompactionInterval(5000)
     }
   }
   getInstance() {
-    return SshRecordDB.instance
+    return CredentialsDB.instance
   }
 }
 

server/app/utils/storage.js

@@ -1,42 +0,0 @@
-const { SshRecordDB } = require('./db-class')
-
-const readSSHRecord = async () => {
-  const sshRecordDB = new SshRecordDB().getInstance()
-  return new Promise((resolve, reject) => {
-    sshRecordDB.find({}, (err, docs) => {
-      if (err) {
-        consola.error('读取ssh-record-db错误: ', err)
-        reject(err)
-      } else {
-        resolve(docs)
-      }
-    })
-  })
-}
-
-const writeSSHRecord = async (record = []) => {
-  return new Promise((resolve, reject) => {
-    const sshRecordDB = new SshRecordDB().getInstance()
-    sshRecordDB.remove({}, { multi: true }, (err) => {
-      if (err) {
-        consola.error('清空SSHRecord出错:', err)
-        reject(err)
-      } else {
-        sshRecordDB.compactDatafile()
-        sshRecordDB.insert(record, (err, newDocs) => {
-          if (err) {
-            consola.error('写入新的ssh记录出错:', err)
-            reject(err)
-          } else {
-            sshRecordDB.compactDatafile()
-            resolve(newDocs)
-          }
-        })
-      }
-    })
-  })
-}
-
-module.exports = {
-  readSSHRecord, writeSSHRecord
-}