From 90ee38ff442843b89212c91494cb8c0e93bf787c Mon Sep 17 00:00:00 2001
From: chaos-zhu <ggchaos1022@gmail.com>
Date: Tue, 22 Oct 2024 23:00:12 +0800
Subject: [PATCH] =?UTF-8?q?:recycle:=20=E9=87=8D=E6=9E=84=E6=9C=AC?=
 =?UTF-8?q?=E5=9C=B0=E6=95=B0=E6=8D=AE=E5=BA=93-onekey=E6=A8=A1=E5=9D=97?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server/app/controller/onekey.js | 12 ++++----
 server/app/socket/onekey.js     |  7 +++--
 server/app/utils/storage.js     | 50 ++-------------------------------
 3 files changed, 11 insertions(+), 58 deletions(-)

diff --git a/server/app/controller/onekey.js b/server/app/controller/onekey.js
index 8fa378a..05a4795 100644
--- a/server/app/controller/onekey.js
+++ b/server/app/controller/onekey.js
@@ -1,7 +1,8 @@
-const { readOneKeyRecord, deleteOneKeyRecord } = require('../utils/storage')
+const { OnekeyDB } = require('../utils/db-class')
+const onekeyDB = new OnekeyDB().getInstance()
 
 async function getOnekeyRecord({ res }) {
-  let data = await readOneKeyRecord()
+  let data = await onekeyDB.findAsync({})
   data = data.map(item => {
     return { ...item, id: item._id }
   })
@@ -11,14 +12,11 @@ async function getOnekeyRecord({ res }) {
 
 const removeOnekeyRecord = async ({ res, request }) => {
   let { body: { ids } } = request
-  let onekeyRecord = await readOneKeyRecord()
   if (ids === 'ALL') {
-    ids = onekeyRecord.map(item => item._id)
-    await deleteOneKeyRecord(ids)
+    await onekeyDB.removeAsync({}, { multi: true })
     res.success({ data: '移除全部成功' })
   } else {
-    if (!onekeyRecord.some(item => ids.includes(item._id))) return res.fail({ msg: '批量指令记录ID不存在' })
-    await deleteOneKeyRecord(ids)
+    await onekeyDB.removeAsync({ _id: { $in: ids } })
     res.success({ data: '移除成功' })
   }
 }
diff --git a/server/app/socket/onekey.js b/server/app/socket/onekey.js
index a366338..57a3d4a 100644
--- a/server/app/socket/onekey.js
+++ b/server/app/socket/onekey.js
@@ -1,13 +1,14 @@
 const { Server } = require('socket.io')
 const { Client: SSHClient } = require('ssh2')
 const { asyncSendNotice } = require('../utils/notify')
-const { readSSHRecord, writeOneKeyRecord } = require('../utils/storage')
+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 } = require('../utils/db-class')
+const { HostListDB, OnekeyDB } = require('../utils/db-class')
 const hostListDB = new HostListDB().getInstance()
+const onekeyDB = new OnekeyDB().getInstance()
 
 const execStatusEnum = {
   connecting: '连接中',
@@ -207,7 +208,7 @@ module.exports = (httpServer) => {
           item.status = execStatusEnum.socketInterrupt
         }
       })
-      await writeOneKeyRecord(execResult)
+      await onekeyDB.insertAsync(execResult)
       isExecuting = false
       execResult = []
       execClient = []
diff --git a/server/app/utils/storage.js b/server/app/utils/storage.js
index d1ae48b..a12acce 100644
--- a/server/app/utils/storage.js
+++ b/server/app/utils/storage.js
@@ -1,4 +1,4 @@
-const { KeyDB, SshRecordDB, ScriptsDB, OnekeyDB } = require('./db-class')
+const { KeyDB, SshRecordDB, OnekeyDB } = require('./db-class')
 
 const readKey = async () => {
   return new Promise((resolve, reject) => {
@@ -64,53 +64,7 @@ const writeSSHRecord = async (record = []) => {
   })
 }
 
-const readOneKeyRecord = async () => {
-  return new Promise((resolve, reject) => {
-    const onekeyDB = new OnekeyDB().getInstance()
-    onekeyDB.find({}, (err, docs) => {
-      if (err) {
-        consola.error('读取onekey record错误: ', err)
-        reject(err)
-      } else {
-        onekeyDB.compactDatafile()
-        resolve(docs)
-      }
-    })
-  })
-}
-
-const writeOneKeyRecord = async (records =[]) => {
-  return new Promise((resolve, reject) => {
-    const onekeyDB = new OnekeyDB().getInstance()
-    onekeyDB.insert(records, (err, newDocs) => {
-      if (err) {
-        consola.error('写入新的onekey记录出错:', err)
-        reject(err)
-      } else {
-        onekeyDB.compactDatafile()
-        resolve(newDocs)
-      }
-    })
-  })
-}
-
-const deleteOneKeyRecord = async (ids =[]) => {
-  return new Promise((resolve, reject) => {
-    const onekeyDB = new OnekeyDB().getInstance()
-    onekeyDB.remove({ _id: { $in: ids } }, { multi: true }, function (err, numRemoved) {
-      if (err) {
-        consola.error('Error deleting onekey record(s):', err)
-        reject(err)
-      } else {
-        onekeyDB.compactDatafile()
-        resolve(numRemoved)
-      }
-    })
-  })
-}
-
 module.exports = {
   readSSHRecord, writeSSHRecord,
-  readKey, writeKey,
-  readOneKeyRecord, writeOneKeyRecord, deleteOneKeyRecord
+  readKey, writeKey
 }