feat: expose ssh port forward as event based API

shellhub-io · Oct 3, 2024 · 46f3cb4 · 46f3cb4
1 parent 702a9e3
commit 46f3cb4
Show file tree

Hide file tree

Showing 4 changed files with 242 additions and 132 deletions.
diff --git a/package.json b/package.json
@@ -32,7 +32,6 @@
     "@vueuse/core": "^10.7.0",
     "axios": "^1.6.2",
     "electron-updater": "^6.1.1",
-    "node-ssh": "^13.2.0",
     "pinia": "^2.1.7",
     "sass": "^1.69.5",
     "socksv5": "^0.0.6",
@@ -48,7 +47,7 @@
     "@electron-toolkit/eslint-config-ts": "^1.0.0",
     "@electron-toolkit/tsconfig": "^1.0.1",
     "@rushstack/eslint-patch": "^1.3.3",
-    "@types/node": "^18.17.5",
+    "@types/node": "^22.7.4",
     "@vitejs/plugin-vue": "^4.3.1",
     "@vue/eslint-config-prettier": "^8.0.0",
     "@vue/eslint-config-typescript": "^11.0.3",

diff --git a/src/preload/index.d.ts b/src/preload/index.d.ts
@@ -1,4 +1,201 @@
 import { ElectronAPI } from '@electron-toolkit/preload'
+import EventEmitter from 'events'
+import net from 'net'
+import socks from 'socksv5'
+import ssh2 from 'ssh2'
+
+export enum SSHEvent {
+  Auth = 'auth',
+  Connect = 'connect',
+  Error = 'error',
+  Disconnect = 'disconnect'
+}
+
+export class SSHEmitter extends EventEmitter {}
+
+export type SSHConnectionAuth = {
+  host: string
+  username: string
+  password: string
+  namespace: string
+  device: string
+}
+
+export type SSHLocalPortForwardSettings = {
+  sourceAddr: string
+  sourcePort: number
+  destinationAddr: string
+  destinationPort: number
+}
+
+export type SSHDynamicPortForwardSettings = {
+  destinationAddr: string
+  destinationPort: number
+}
+
+export interface SSHConnection {
+  eventEmitter: SSHEmitter
+  client: ssh2.Client
+  connect(auth: SSHConnectionAuth, settings: any): void
+  disconnect(): void
+  onAuth(callback: any): void
+  onListen(callback: any): void
+  onError(callback: any): void
+  onDisconnect(callback: any): void
+}
+
+export class SSHConnectionLocalPortForward implements SSHConnection {
+  eventEmitter = new SSHEmitter()
+  client = new ssh2.Client()
+
+  connect(auth: SSHConnectionAuth, settings: any) {
+    this.client.on('ready', () => {
+      console.info('SSH connection established')
+      this.eventEmitter.emit(SSHEvent.Auth)
+
+      this.client.forwardOut(
+        settings.sourceAddr,
+        settings.sourcePort,
+        settings.destinationAddr,
+        settings.destinationPort,
+        (err, stream) => {
+          if (err) {
+            console.error('SSH forwarding error:', err)
+            this.eventEmitter.emit(SSHEvent.Error, err)
+            return
+          }
+
+          const server = net.createServer((client) => {
+            client.pipe(stream).pipe(client)
+          })
+
+          server.on('close', () => {
+            this.eventEmitter.emit(SSHEvent.Disconnect)
+          })
+
+          server.listen(settings.sourcePort, settings.sourceAddr, () => {
+            console.log(
+              `Local port forward started from ${settings.sourceAddr}:${settings.sourcePort} to ${settings.destinationAddr}:${settings.destinationPort}.`
+            )
+            this.eventEmitter.emit(SSHEvent.Connect, settings.sourcePort, settings.destinationAddr)
+          })
+        }
+      )
+    })
+
+    this.client.on('error', (err) => {
+      console.error('SSH connection error:', err)
+      this.eventEmitter.emit(SSHEvent.Error, err)
+    })
+
+    this.client.connect({
+      host: auth.host,
+      username: `${auth.username}@${auth.namespace}.${auth.device}`,
+      password: auth.password
+    })
+  }
+
+  disconnect() {
+    if (this.client) {
+      this.client.end()
+      this.eventEmitter.emit(SSHEvent.Disconnect)
+    }
+  }
+
+  onAuth(callback: any) {
+    this.eventEmitter.on(SSHEvent.Auth, callback)
+  }
+
+  onConnect(callback: any) {
+    this.eventEmitter.on(SSHEvent.Connect, callback)
+  }
+
+  onError(callback: any) {
+    this.eventEmitter.on(SSHEvent.Error, callback)
+  }
+
+  onDisconnect(callback: any) {
+    this.eventEmitter.on(SSHEvent.Disconnect, callback)
+  }
+}
+
+export class SSHConnectionDynamicPortForward implements SSHConnection {
+  eventEmitter = new SSHEmitter()
+  client = new ssh2.Client()
+
+  connect(auth: SSHConnectionAuth, settings: any) {
+    this.client.on('ready', () => {
+      console.info('SSH connection established')
+      this.eventEmitter.emit(SSHEvent.Auth)
+
+      const server = socks.createServer((info, accept, deny) => {
+        this.client.forwardOut(
+          info.srcAddr,
+          info.srcPort,
+          info.dstAddr,
+          info.dstPort,
+          (err, stream) => {
+            if (err) {
+              console.error('SSH forwarding error:', err)
+              this.eventEmitter.emit(SSHEvent.Error, err)
+              deny()
+              return
+            }
+
+            const client = accept(true)
+            stream.pipe(client).pipe(stream)
+          }
+        )
+      })
+
+      server.on('error', (err) => {
+        console.log('Server error:', err)
+        this.eventEmitter.emit(SSHEvent.Error, err)
+      })
+
+      server.useAuth(socks.auth.None())
+
+      server.listen(settings.destinationPort, settings.destinationAddr, () => {
+        console.log('Server listening on', settings.destinationPort, settings.destinationAddr)
+        this.eventEmitter.emit(SSHEvent.Connect, settings.destinationPort, settings.destinationAddr)
+      })
+    })
+
+    this.client.on('error', (err) => {
+      console.error('SSH connection error:', err)
+      this.eventEmitter.emit(SSHEvent.Error, err)
+    })
+
+    this.client.connect({
+      host: auth.host,
+      username: `${auth.username}@${auth.namespace}.${auth.device}`,
+      password: auth.password
+    })
+  }
+
+  disconnect() {
+    if (this.client) {
+      this.client.end()
+      this.eventEmitter.emit(SSHEvent.Disconnect)
+    }
+  }
+
+  onAuth(callback: any) {
+    this.eventEmitter.on(SSHEvent.Auth, callback)
+  }
+
+  onConnect(callback: any) {
+    this.eventEmitter.on(SSHEvent.Cnnect, callback)
+  }
+
+  onError(callback: any) {
+    this.eventEmitter.on(SSHEvent.Error, callback)
+  }
+
+  onDisconnect(callback: any) {
+    this.eventEmitter.on(SSHEvent.Disconnect, callback)
+  }
+}
 
 declare global {
   interface Window {

diff --git a/src/preload/index.ts b/src/preload/index.ts
@@ -1,104 +1,42 @@
 import { contextBridge } from 'electron'
 import { electronAPI } from '@electron-toolkit/preload'
-import net from 'net'
-import { Channel } from 'ssh2'
-import { NodeSSH } from 'node-ssh'
-import socks from 'socksv5'
+import { SSHConnectionLocalPortForward, SSHConnectionDynamicPortForward } from './index.d'
 
 // Custom APIs for renderer
 const api = {}
 
-type SSHAuth = {
-  host: string
-  namespace: string
-  device: string
-  username: string
-  password: string
-}
-
-type SSHLocalPortForwardSettings = {
-  sourceIP: string
-  sourcePort: number
-  destinationIP: string
-  destinationPort: number
-}
-
-type SSHDynamicPortForwardSettings = {
-  destinationIP: string
-  destinationPort: number
-}
-
 const ssh = {
-  localPortForward: (auth: SSHAuth, settings: SSHLocalPortForwardSettings) => {
-    new NodeSSH()
-      .connect({
-        host: auth.host,
-        username: `${auth.username}@${auth.namespace}.${auth.device}`,
-        password: auth.password
-      })
-      .then((ssh: NodeSSH) => {
-        console.info('SSH connection stablished')
-
-        return ssh.forwardOut(
-          settings.sourceIP,
-          settings.sourcePort,
-          settings.destinationIP,
-          settings.destinationPort
-        )
-      })
-      .then((connection: Channel) => {
-        console.info('Local port forward started')
+  localPortForward: () => {
+    const localPortForwardInstance = new SSHConnectionLocalPortForward()
+
+    const localPortForwardObject = {
+      eventEmitter: localPortForwardInstance.eventEmitter,
+      client: localPortForwardInstance.client,
+      connect: localPortForwardInstance.connect.bind(localPortForwardInstance),
+      disconnect: localPortForwardInstance.disconnect.bind(localPortForwardInstance),
+      onAuth: localPortForwardInstance.onAuth.bind(localPortForwardInstance),
+      onConnect: localPortForwardInstance.onConnect.bind(localPortForwardInstance),
+      onError: localPortForwardInstance.onError.bind(localPortForwardInstance),
+      onDisconnect: localPortForwardInstance.onDisconnect.bind(localPortForwardInstance)
+    }
 
-        return net
-          .createServer((client: net.Socket) => {
-            client.pipe(connection).pipe(client)
-          })
-          .listen(settings.sourcePort, settings.sourceIP, () => {
-            console.log(
-              `Local port forward started from ${settings.sourceIP}:${settings.sourcePort} to ${settings.destinationIP}:${settings.destinationPort}.`
-            )
-          })
-      })
-      .catch((e) => {
-        console.error('SSH connection error:', e)
-      })
+    return localPortForwardObject
   },
-  dynamicPortForward: (auth: SSHAuth, settings: SSHDynamicPortForwardSettings) => {
-    try {
-      const client = new NodeSSH().connect({
-        host: auth.host,
-        username: `${auth.username}@${auth.namespace}.${auth.device}`,
-        password: auth.password
-      })
-
-      const server = socks.createServer((info, accept, _) => {
-        client.then((connection: NodeSSH) => {
-          connection
-            .forwardOut(info.srcAddr, info.srcPort, info.dstAddr, info.dstPort)
-            .then((channel: Channel) => {
-              console.info(info)
-              const client = accept(true)
-
-              channel.pipe(client).pipe(channel)
-            })
-            .catch((e) => {
-              console.error(e)
-            })
-        })
-      })
-
-      server.on('error', (err) => {
-        console.log('Server error:', err)
-      })
-
-      server.useAuth(socks.auth.None())
-
-      server.listen(settings.destinationPort, settings.destinationIP, () => {
-        console.log('Server listening on', settings.destinationPort, settings.destinationIP)
-      })
-    } catch (e) {
-      console.log('Failed to create server', e)
+  dynamicPortForward: () => {
+    const dynamicPortForwardInstance = new SSHConnectionDynamicPortForward()
+
+    const dynamicPortForwardObject = {
+      eventEmitter: dynamicPortForwardInstance.eventEmitter,
+      client: dynamicPortForwardInstance.client,
+      connect: dynamicPortForwardInstance.connect.bind(dynamicPortForwardInstance),
+      disconnect: dynamicPortForwardInstance.disconnect.bind(dynamicPortForwardInstance),
+      onAuth: dynamicPortForwardInstance.onAuth.bind(dynamicPortForwardInstance),
+      onConnect: dynamicPortForwardInstance.onConnect.bind(dynamicPortForwardInstance),
+      onError: dynamicPortForwardInstance.onError.bind(dynamicPortForwardInstance),
+      onDisconnect: dynamicPortForwardInstance.onDisconnect.bind(dynamicPortForwardInstance)
     }
+
+    return dynamicPortForwardObject
   }
 }