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

server/app/controller/user.js

@@ -1,14 +1,14 @@
 const jwt = require('jsonwebtoken')
 const axios = require('axios')
-const { asyncSendNotice } = require('../utils/notify')
-const { readKey, writeKey } = require('../utils/storage')
+const { sendNoticeAsync } = require('../utils/notify')
 const { RSADecryptAsync, AESEncryptAsync, SHA1Encrypt } = require('../utils/encrypt')
 const { getNetIPInfo } = require('../utils/tools')
-const { LogDB } = require('../utils/db-class')
+const { KeyDB, LogDB } = require('../utils/db-class')
+const keyDB = new KeyDB().getInstance()
 const logDB = new LogDB().getInstance()
 
 const getpublicKey = async ({ res }) => {
-  let { publicKey: data } = await readKey()
+  let { publicKey: data } = await keyDB.findOneAsync({})
   if (!data) return res.fail({ msg: 'publicKey not found, Try to restart the server', status: 500 })
   res.success({ data })
 }
@@ -30,7 +30,7 @@ const login = async ({ res, request }) => {
   if (loginErrCount >= allowErrCount) {
     const { ip, country, city } = await getNetIPInfo(clientIp)
     // 异步发送通知&禁止登录
-    asyncSendNotice('err_login', '登录错误提醒', `错误登录次数: ${ loginErrTotal }\n地点：${ country + city }\nIP: ${ ip }`)
+    sendNoticeAsync('err_login', '登录错误提醒', `错误登录次数: ${ loginErrTotal }\n地点：${ country + city }\nIP: ${ ip }`)
     forbidLogin = true
     loginErrCount = 0
 
@@ -56,7 +56,7 @@ const login = async ({ res, request }) => {
     // console.log('ciphertext', ciphertext)
     let loginPwd = await RSADecryptAsync(ciphertext)
     // console.log('Decrypt解密password:', loginPwd)
-    let { user, pwd } = await readKey()
+    let { user, pwd } = await keyDB.findOneAsync({})
     if (loginName === user && loginPwd === 'admin' && pwd === 'admin') {
       const token = await beforeLoginHandler(clientIp, jwtExpires)
       return res.success({ data: { token, jwtExpires }, msg: '登录成功，请及时修改默认用户名和密码' })
@@ -76,7 +76,7 @@ const beforeLoginHandler = async (clientIp, jwtExpires) => {
 
   // consola.success('登录成功, 准备生成token', new Date())
   // 生产token
-  let { commonKey } = await readKey()
+  let { commonKey } = await keyDB.findOneAsync({})
   let token = jwt.sign({ date: Date.now() }, commonKey, { expiresIn: jwtExpires }) // 生成token
   token = await AESEncryptAsync(token) // 对称加密token后再传输给前端
 
@@ -86,7 +86,7 @@ const beforeLoginHandler = async (clientIp, jwtExpires) => {
   consola.info('登录成功:', new Date(), { ip, country, city })
 
   // 邮件登录通知
-  asyncSendNotice('login', '登录提醒', `地点：${ country + city }\nIP: ${ ip }`)
+  sendNoticeAsync('login', '登录提醒', `地点：${ country + city }\nIP: ${ ip }`)
 
   await logDB.insertAsync({ ip, country, city, date: Date.now(), type: 'login' })
   return token
@@ -96,17 +96,15 @@ const updatePwd = async ({ res, request }) => {
   let { body: { oldLoginName, oldPwd, newLoginName, newPwd } } = request
   let rsaOldPwd = await RSADecryptAsync(oldPwd)
   oldPwd = rsaOldPwd === 'admin' ? 'admin' : SHA1Encrypt(rsaOldPwd)
-  let keyObj = await readKey()
+  let keyObj = await keyDB.findOneAsync({})
   let { user, pwd } = keyObj
   if (oldLoginName !== user || oldPwd !== pwd) return res.fail({ data: false, msg: '原用户名或密码校验失败' })
   // 旧密钥校验通过，加密保存新密码
   newPwd = await RSADecryptAsync(newPwd) === 'admin' ? 'admin' : SHA1Encrypt(await RSADecryptAsync(newPwd))
   keyObj.user = newLoginName
   keyObj.pwd = newPwd
-  await writeKey(keyObj)
-
-  asyncSendNotice('updatePwd', '用户密码修改提醒', `原用户名：${ user }\n更新用户名: ${ newLoginName }`)
-
+  await keyDB.updateAsync({}, keyObj)
+  sendNoticeAsync('updatePwd', '用户密码修改提醒', `原用户名：${ user }\n更新用户名: ${ newLoginName }`)
   res.success({ data: true, msg: 'success' })
 }
 

server/app/db.js

@@ -1,29 +1,28 @@
-const { writeKey } = require('./utils/storage')
+const NodeRSA = require('node-rsa')
+const { randomStr } = require('./utils/tools')
+const { AESEncryptAsync } = require('./utils/encrypt')
 const { KeyDB, GroupDB, NotifyDB, NotifyConfigDB } = require('./utils/db-class')
 
-function initKeyDB() {
-  return new Promise((resolve, reject) => {
-    const keyDB = new KeyDB().getInstance()
-    keyDB.count({}, async (err, count) => {
-      if (err) {
-        consola.log('初始化keyDB错误:', err)
-        reject(err)
-      } else {
-        if (count === 0) {
-          consola.log('初始化keyDB✔')
-          const defaultData = {
-            user: 'admin',
-            pwd: 'admin',
-            commonKey: '',
-            publicKey: '',
-            privateKey: ''
-          }
-          await writeKey(defaultData)
-        }
-      }
-      resolve()
-    })
-  })
+async function initKeyDB() {
+  const keyDB = new KeyDB().getInstance()
+  let count = await keyDB.countAsync({})
+  if (count !== 0) return consola.info('公私钥已存在[重新生成会导致已保存的ssh密钥信息失效]')
+  let newConfig = {
+    user: 'admin',
+    pwd: 'admin',
+    commonKey: randomStr(16),
+    publicKey: '',
+    privateKey: ''
+  }
+  await keyDB.insertAsync(newConfig)
+  let key = new NodeRSA({ b: 1024 })
+  key.setOptions({ encryptionScheme: 'pkcs1', environment: 'browser' })
+  let privateKey = key.exportKey('pkcs1-private-pem')
+  let publicKey = key.exportKey('pkcs8-public-pem')
+  newConfig.privateKey = await AESEncryptAsync(privateKey, newConfig.commonKey) // 加密私钥
+  newConfig.publicKey = publicKey // 公开公钥
+  await keyDB.updateAsync({}, { $set: newConfig }, { upsert: true })
+  consola.info('Task: 已生成新的非对称加密公私钥')
 }
 
 async function initGroupDB() {

server/app/init.js

@@ -1,33 +0,0 @@
-const NodeRSA = require('node-rsa')
-const { readKey, writeKey } = require('./utils/storage')
-const { randomStr } = require('./utils/tools')
-const { AESEncryptAsync } = require('./utils/encrypt')
-
-// 初始化公私钥, 供登录、保存ssh密钥/密码等加解密
-async function initRsa() {
-  let keyObj = await readKey()
-  if (keyObj.privateKey && keyObj.publicKey) return consola.info('公私钥已存在[重新生成会导致已保存的ssh密钥信息失效]')
-  let key = new NodeRSA({ b: 1024 })
-  key.setOptions({ encryptionScheme: 'pkcs1', environment: 'browser' })
-  let privateKey = key.exportKey('pkcs1-private-pem')
-  let publicKey = key.exportKey('pkcs8-public-pem')
-  keyObj.privateKey = await AESEncryptAsync(privateKey) // 加密私钥
-  keyObj.publicKey = publicKey // 公开公钥
-  await writeKey(keyObj)
-  consola.info('Task: 已生成新的非对称加密公私钥')
-}
-
-// 随机的commonKey secret
-async function randomJWTSecret() {
-  let keyObj = await readKey()
-  if (keyObj?.commonKey) return consola.info('commonKey密钥已存在')
-
-  keyObj.commonKey = randomStr(16)
-  await writeKey(keyObj)
-  consola.info('Task: 已生成新的随机commonKey密钥')
-}
-
-module.exports = async () => {
-  await randomJWTSecret() // 全局密钥
-  await initRsa() // 全局公钥密钥
-}

server/app/main.js

@@ -2,12 +2,10 @@ const consola = require('consola')
 global.consola = consola
 const { httpServer } = require('./server')
 const initDB = require('./db')
-const initEncryptConf = require('./init')
 const scheduleJob = require('./schedule')
 
 async function main() {
   await initDB()
-  await initEncryptConf()
   httpServer()
   scheduleJob()
 }

server/app/schedule/expired-notify.js

@@ -1,5 +1,5 @@
 const schedule = require('node-schedule')
-const { asyncSendNotice } = require('../utils/notify')
+const { sendNoticeAsync } = require('../utils/notify')
 const { formatTimestamp } = require('../utils/tools')
 const { HostListDB } = require('../utils/db-class')
 const hostListDB = new HostListDB().getInstance()
@@ -16,13 +16,13 @@ const expiredNotifyJob = async () => {
     let content = `别名: ${ name }\nIP: ${ host }\n到期时间：${ formatTimestamp(expired, 'week') }\n控制台: ${ consoleUrl || '未填写' }`
     if (0 <= restDay && restDay <= 1) {
       let temp = '有服务器将在一天后到期，请关注\n'
-      asyncSendNotice('host_expired', title, temp + content)
+      sendNoticeAsync('host_expired', title, temp + content)
     } else if (3 <= restDay && restDay < 4) {
       let temp = '有服务器将在三天后到期，请关注\n'
-      asyncSendNotice('host_expired', title, temp + content)
+      sendNoticeAsync('host_expired', title, temp + content)
     } else if (7 <= restDay && restDay < 8) {
       let temp = '有服务器将在七天后到期，请关注\n'
-      asyncSendNotice('host_expired', title, temp + content)
+      sendNoticeAsync('host_expired', title, temp + content)
     }
   }
 }

server/app/socket/onekey.js

@@ -1,6 +1,6 @@
 const { Server } = require('socket.io')
 const { Client: SSHClient } = require('ssh2')
-const { asyncSendNotice } = require('../utils/notify')
+const { sendNoticeAsync } = require('../utils/notify')
 const { readSSHRecord } = require('../utils/storage')
 const { verifyAuthSync } = require('../utils/verify-auth')
 const { shellThrottle } = require('../utils/tools')
@@ -124,7 +124,7 @@ module.exports = (httpServer) => {
           }
         })
         let reason = `执行超时,已强制终止执行 - 超时时间${ timeout }秒`
-        asyncSendNotice('onekey_complete', '批量指令执行超时', reason)
+        sendNoticeAsync('onekey_complete', '批量指令执行超时', reason)
         socket.emit('timeout', { reason, result: execResult })
         socket.disconnect()
         disconnectAllExecClient()
@@ -191,7 +191,7 @@ module.exports = (httpServer) => {
         await Promise.all(execPromise)
         consola.success('onekey执行完成')
         socket.emit('exec_complete')
-        asyncSendNotice('onekey_complete', '批量指令执行完成', '请登录面板查看执行结果')
+        sendNoticeAsync('onekey_complete', '批量指令执行完成', '请登录面板查看执行结果')
         socket.disconnect()
       } catch (error) {
         consola.error('onekey执行失败', error)

server/app/socket/terminal.js

@@ -3,7 +3,7 @@ const { Client: SSHClient } = require('ssh2')
 const { verifyAuthSync } = require('../utils/verify-auth')
 const { AESDecryptAsync } = require('../utils/encrypt')
 const { readSSHRecord } = require('../utils/storage')
-const { asyncSendNotice } = require('../utils/notify')
+const { sendNoticeAsync } = require('../utils/notify')
 const { isAllowedIp, ping } = require('../utils/tools')
 const { HostListDB } = require('../utils/db-class')
 const hostListDB = new HostListDB().getInstance()
@@ -75,7 +75,7 @@ async function createTerminal(hostId, socket, sshClient) {
       consola.log('连接信息', { username, port, authType })
       sshClient
         .on('ready', async() => {
-          asyncSendNotice('host_login', '终端登录', `别名: ${ name } \n IP：${ host } \n 端口：${ port } \n 状态: 登录成功`)
+          sendNoticeAsync('host_login', '终端登录', `别名: ${ name } \n IP：${ host } \n 端口：${ port } \n 状态: 登录成功`)
           consola.success('终端连接成功：', host)
           socket.emit('connect_terminal_success', `终端连接成功：${ host }`)
           let stream = await createInteractiveShell(socket, sshClient)
@@ -92,7 +92,7 @@ async function createTerminal(hostId, socket, sshClient) {
         })
         .on('error', (err) => {
           consola.log(err)
-          asyncSendNotice('host_login', '终端登录', `别名: ${ name } \n IP：${ host } \n 端口：${ port } \n 状态: 登录失败`)
+          sendNoticeAsync('host_login', '终端登录', `别名: ${ name } \n IP：${ host } \n 端口：${ port } \n 状态: 登录失败`)
           consola.error('连接终端失败:', host, err.message)
           socket.emit('connect_fail', err.message)
         })

server/app/utils/encrypt.js

@@ -1,12 +1,13 @@
 const CryptoJS = require('crypto-js')
 const rawCrypto = require('crypto')
 const NodeRSA = require('node-rsa')
-const { readKey } = require('./storage.js')
+const { KeyDB } = require('./db-class')
+const keyDB = new KeyDB().getInstance()
 
 // rsa非对称 私钥解密
 const RSADecryptAsync = async (ciphertext) => {
   if (!ciphertext) return
-  let { privateKey } = await readKey()
+  let { privateKey } = await keyDB.findOneAsync({})
   privateKey = await AESDecryptAsync(privateKey) // 先解密私钥
   const rsakey = new NodeRSA(privateKey)
   rsakey.setOptions({ encryptionScheme: 'pkcs1', environment: 'browser' }) // Must Set It When Frontend Use jsencrypt
@@ -17,7 +18,7 @@ const RSADecryptAsync = async (ciphertext) => {
 // aes对称 加密(default commonKey)
 const AESEncryptAsync = async (text, key) => {
   if (!text) return
-  let { commonKey } = await readKey()
+  let { commonKey } = await keyDB.findOneAsync({})
   let ciphertext = CryptoJS.AES.encrypt(text, key || commonKey).toString()
   return ciphertext
 }
@@ -25,7 +26,7 @@ const AESEncryptAsync = async (text, key) => {
 // aes对称 解密(default commonKey)
 const AESDecryptAsync = async (ciphertext, key) => {
   if (!ciphertext) return
-  let { commonKey } = await readKey()
+  let { commonKey } = await keyDB.findOneAsync({})
   let bytes = CryptoJS.AES.decrypt(ciphertext, key || commonKey)
   let originalText = bytes.toString(CryptoJS.enc.Utf8)
   return originalText

server/app/utils/notify.js

@@ -56,7 +56,7 @@ function sendEmail({ service, user, pass }, title, content) {
 }
 
 // 异步发送通知
-async function asyncSendNotice(noticeAction, title, content) {
+async function sendNoticeAsync(noticeAction, title, content) {
   try {
     let notifyList = await notifyDB.findAsync({})
     let { sw } = notifyList.find((item) => item.type === noticeAction) // 获取对应动作的通知开关
@@ -87,7 +87,7 @@ async function asyncSendNotice(noticeAction, title, content) {
 }
 
 module.exports = {
-  asyncSendNotice,
+  sendNoticeAsync,
   sendServerChan,
   sendEmail
 }

server/app/utils/storage.js

@@ -1,31 +1,4 @@
-const { KeyDB, SshRecordDB, OnekeyDB } = require('./db-class')
-
-const readKey = async () => {
-  return new Promise((resolve, reject) => {
-    const keyDB = new KeyDB().getInstance()
-    keyDB.findOne({}, (err, doc) => {
-      if (err) {
-        reject(err)
-      } else {
-        resolve(doc)
-      }
-    })
-  })
-}
-
-const writeKey = async (keyObj = {}) => {
-  const keyDB = new KeyDB().getInstance()
-  return new Promise((resolve, reject) => {
-    keyDB.update({}, { $set: keyObj }, { upsert: true }, (err, numReplaced) => {
-      if (err) {
-        reject(err)
-      } else {
-        keyDB.compactDatafile()
-        resolve(numReplaced)
-      }
-    })
-  })
-}
+const { SshRecordDB } = require('./db-class')
 
 const readSSHRecord = async () => {
   const sshRecordDB = new SshRecordDB().getInstance()
@@ -65,6 +38,5 @@ const writeSSHRecord = async (record = []) => {
 }
 
 module.exports = {
-  readSSHRecord, writeSSHRecord,
-  readKey, writeKey
+  readSSHRecord, writeSSHRecord
 }

server/app/utils/verify-auth.js

@@ -1,7 +1,8 @@
 
 const { AESDecryptAsync } = require('./encrypt')
-const { readKey } = require('./storage')
 const jwt = require('jsonwebtoken')
+const { KeyDB } = require('./db-class')
+const keyDB = new KeyDB().getInstance()
 
 const enumLoginCode = {
   SUCCESS: 1,
@@ -14,7 +15,7 @@ const verifyAuthSync = async (token, clientIp) => {
   consola.info('verifyAuthSync IP：', clientIp)
   try {
     token = await AESDecryptAsync(token) // 先aes解密
-    const { commonKey } = await readKey()
+    const { commonKey } = await keyDB.findOneAsync({})
     const { exp } = jwt.verify(token, commonKey)
     if (Date.now() > (exp * 1000)) return { code: -1, msg: 'token expires' } // 过期
     return { code: enumLoginCode.SUCCESS, msg: 'success' } // 验证成功