hyper-lab / lib / ssh-manager.js

import { Client } from 'ssh2'
import fs from 'fs'
import chalk from 'chalk'

export class SSHManager {
  constructor(agentConfig) {
    this.config = agentConfig
    this.client = null
    this.connected = false
  }

  async connect() {
    if (this.connected) return

    // Special case: localhost doesn't need SSH
    if (this.config.host === 'localhost' || this.config.host === '127.0.0.1') {
      this.connected = true
      this.isLocalhost = true
      console.log(chalk.green(`✅ Connected to ${this.config.name} (localhost)`))
      return
    }

    this.client = new Client()
    this.isLocalhost = false

    return new Promise((resolve, reject) => {
      const timeout = setTimeout(() => {
        reject(new Error(`Connection timeout for ${this.config.name}`))
      }, 10000)

      this.client.on('ready', () => {
        clearTimeout(timeout)
        this.connected = true
        console.log(chalk.green(`✅ Connected to ${this.config.name}`))
        resolve()
      })

      this.client.on('error', (err) => {
        clearTimeout(timeout)
        reject(new Error(`Connection failed for ${this.config.name}: ${err.message}`))
      })

      // Try to use SSH agent first (for encrypted keys)
      const connectOptions = {
        host: this.config.host,
        username: this.config.user,
        agent: process.env.SSH_AUTH_SOCK, // Use ssh-agent
        agentForward: true
      }

      // If agent not available, try reading key directly
      if (!connectOptions.agent) {
        connectOptions.privateKey = fs.readFileSync(this.config.keyFile)
      }

      this.client.connect(connectOptions)
    })
  }

  async exec(command, options = {}) {
    if (!this.connected) {
      throw new Error('Not connected')
    }

    // Special case: localhost uses child_process
    if (this.isLocalhost) {
      const { exec } = await import('child_process')
      const { promisify } = await import('util')
      const execAsync = promisify(exec)

      try {
        const result = await execAsync(command, { maxBuffer: 10 * 1024 * 1024 })

        if (options.logOutput) {
          console.log(chalk.cyan(`[${this.config.name}] ${result.stdout.trim()}`))
        }

        return {
          stdout: result.stdout,
          stderr: result.stderr,
          exitCode: 0
        }
      } catch (err) {
        if (options.logOutput) {
          console.error(chalk.red(`[${this.config.name}] ${err.stderr || err.message}`))
        }

        return {
          stdout: err.stdout || '',
          stderr: err.stderr || err.message,
          exitCode: err.code || 1
        }
      }
    }

    // SSH execution
    return new Promise((resolve, reject) => {
      this.client.exec(command, (err, stream) => {
        if (err) return reject(err)

        let stdout = ''
        let stderr = ''

        stream.on('data', (data) => {
          stdout += data.toString()
          if (options.logOutput) {
            console.log(chalk.cyan(`[${this.config.name}] ${data.toString().trim()}`))
          }
        })

        stream.stderr.on('data', (data) => {
          stderr += data.toString()
          if (options.logOutput) {
            console.error(chalk.red(`[${this.config.name}] ${data.toString().trim()}`))
          }
        })

        stream.on('close', (code) => {
          resolve({ stdout, stderr, exitCode: code })
        })
      })
    })
  }

  async disconnect() {
    if (this.client) {
      this.client.end()
      this.connected = false
    }
  }
}