diff --git a/server/app/controller/ssh.js b/server/app/controller/ssh.js index 8b52654..4a794c6 100644 --- a/server/app/controller/ssh.js +++ b/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() - if (sshRecord.some(item => item.name === name)) return res.fail({ data: false, msg: '已存在同名凭证' }) + let count = await credentialsDB.countAsync({ name }) + 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]) - - sshRecord.push({ ...record, date: Date.now() }) - await writeSSHRecord(sshRecord) + await credentialsDB.insertAsync({ ...record, date: Date.now() }) 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 idx = sshRecord.findIndex(item => item._id === id) - 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] + let oldRecord = await credentialsDB.findOneAsync({ _id: id }) + if (!oldRecord) return res.fail({ data: false, msg: '凭证不存在' }) // 判断原记录是否存在当前更新记录的认证方式 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 - sshRecord.splice(idx, 1, record) - await writeSSHRecord(sshRecord) + await credentialsDB.updateAsync({ _id: id }, record) consola.info('修改凭证:', name) res.success({ data: '保存成功' }) } const removeSSH = async ({ res, request }) => { let { params: { id } } = request - let sshRecord = await readSSHRecord() - let idx = sshRecord.findIndex(item => item._id === id) - if (idx === -1) return res.fail({ msg: '凭证不存在' }) - sshRecord.splice(idx, 1) + let count = await credentialsDB.countAsync({ _id: id }) + if (count === 0) return res.fail({ msg: '凭证不存在' }) // 将删除的凭证id从host中删除 let hostList = await hostListDB.findAsync({}) if (Array.isArray(hostList) && hostList.length > 0) { - for (let item of hostList) { - if (item.credential === id) { - item.credential = '' - await hostListDB.updateAsync({ _id: item._id }, item) + for (let host of hostList) { + let { credential } = host + credential = await AESDecryptAsync(credential) + 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: '移除成功' }) } diff --git a/server/app/socket/onekey.js b/server/app/socket/onekey.js index a1e3ecb..17daf32 100644 --- a/server/app/socket/onekey.js +++ b/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]) diff --git a/server/app/socket/sftp.js b/server/app/socket/sftp.js index c7e203e..444c09f 100644 --- a/server/app/socket/sftp.js +++ b/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]) diff --git a/server/app/socket/terminal.js b/server/app/socket/terminal.js index 9956bd9..13eb9f2 100644 --- a/server/app/socket/terminal.js +++ b/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 = sshRecordList.find(item => item._id === credentialId) + const sshRecord = await credentialsDB.findOneAsync({ _id: credentialId }) authInfo.authType = sshRecord.authType authInfo[authInfo.authType] = await AESDecryptAsync(sshRecord[authInfo.authType]) } else { diff --git a/server/app/utils/db-class.js b/server/app/utils/db-class.js index 35876d4..6c11aba 100644 --- a/server/app/utils/db-class.js +++ b/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) { - SshRecordDB.instance = new Datastore({ filename: credentialsDBPath, autoload: true }) - // SshRecordDB.instance.setAutocompactionInterval(5000) + if (!CredentialsDB.instance) { + CredentialsDB.instance = new Datastore({ filename: credentialsDBPath, autoload: true }) + // CredentialsDB.instance.setAutocompactionInterval(5000) } } getInstance() { - return SshRecordDB.instance + return CredentialsDB.instance } } diff --git a/server/app/utils/storage.js b/server/app/utils/storage.js deleted file mode 100644 index 85921c2..0000000 --- a/server/app/utils/storage.js +++ /dev/null @@ -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 -}