ssh-bridge-ai/src/ssh-enhanced.js

AI-Powered SSH Tool with Bulletproof Connections & Enterprise Sandbox Security + Cursor-like Confirmation - Enable AI assistants to securely SSH into your servers with persistent sessions, keepalive, automatic recovery, sandbox command testing, and user c

const fs = require('fs').promises;
const fsSync = require('fs');
const path = require('path');
const os = require('os');
const EventEmitter = require('events');

// Try to import node-ssh first, fallback to ssh2 if needed
let Client;
let SSH2Client;
try {
  const nodeSSH = require('node-ssh');
  Client = nodeSSH.NodeSSH;
  console.log('Using node-ssh for SSH connections');
} catch (error) {
  console.log('node-ssh not available, falling back to ssh2');
  try {
    const ssh2 = require('ssh2');
    SSH2Client = ssh2.Client;
    console.log('Using ssh2 for SSH connections');
  } catch (ssh2Error) {
    throw new Error('No SSH client available. Please install node-ssh or ssh2: npm install node-ssh ssh2');
  }
}

const { Config } = require('./config');
const { ValidationUtils } = require('./utils/validation');
const logger = require('./utils/logger');
const { SECURITY, NETWORK, FILESYSTEM, ERROR_CODES } = require('./utils/constants');

/**
 * Enhanced SSH Client with robust connection strategies
 * Implements keepalive, fallback shells, session recovery, and comprehensive error handling
 */
class EnhancedSSHClient extends EventEmitter {
  constructor(connection, options = {}) {
    super();
    
    this.connection = connection;
    this.options = this.buildDefaultOptions(options);
    
    // Initialize appropriate SSH client
    if (Client) {
      this.ssh = new Client();
      this.clientType = 'node-ssh';
    } else if (SSH2Client) {
      this.ssh = new SSH2Client();
      this.clientType = 'ssh2';
    } else {
      throw new Error('No SSH client available');
    }
    
    this.config = new Config();
    this.secureBuffers = [];
    this.failedAttempts = 0;
    
    // Connection state
    this.isConnected = false;
    this.connectionStartTime = null;
    this.lastActivityTime = null;
    
    // Parse connection string
    this.parseConnection(connection);
    
    // Initialize components
    this.healthMonitor = new ConnectionHealthMonitor(this);
    this.sessionRecovery = new SessionRecovery(this);
    this.progressiveFallback = new ProgressiveFallback(this);
    
    // Bind event handlers
    this.bindEventHandlers();
  }

  buildDefaultOptions(options) {
    return {
      // Connection settings
      port: 22,
      readyTimeout: 60000,
      
      // Keepalive settings
      keepalive: true,
      keepaliveInterval: 30,
      keepaliveCountMax: 3,
      serverAliveInterval: 30,
      serverAliveCountMax: 3,
      clientAliveInterval: 60,
      clientAliveCountMax: 3,
      keepaliveInitialDelay: 1000,
      
      // Fallback settings
      maxReconnectAttempts: 5,
      reconnectDelay: 2000,
      fallbackShells: ['bash', 'sh', 'dash'],
      useTemporaryHome: true,
      
      // Session recovery
      sessionRecovery: true,
      saveSessionState: true,
      
      // Error handling
      progressiveFallback: true,
      maxFallbackLevel: 5,
      
      // Environment overrides
      environmentOverrides: {
        HOME: '/tmp',
        SHELL: '/bin/bash',
        TERM: 'xterm-256color',
        PATH: '/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin',
        USER: null,
        LOGNAME: null,
        PWD: '/tmp'
      },
      
      ...options
    };
  }

  parseConnection(connection) {
    if (!connection || typeof connection !== 'string') {
      throw new Error('Connection string is required');
    }
    
    const parts = connection.split('@');
    if (parts.length !== 2) {
      throw new Error('Invalid connection format. Use: user@host');
    }
    
    const username = parts[0].trim();
    const host = parts[1].trim();
    
    if (!username || username.length === 0) {
      throw new Error('Username cannot be empty');
    }
    
    if (!ValidationUtils.validateHostname(host)) {
      throw new Error(`Invalid hostname: ${host}`);
    }
    
    this.username = username;
    this.host = host;
    this.connectionString = connection;
    
    logger.debug('Connection parsed', { username, host });
  }

  bindEventHandlers() {
    if (this.ssh) {
      this.ssh.on('ready', () => {
        this.isConnected = true;
        this.connectionStartTime = Date.now();
        this.lastActivityTime = Date.now();
        this.emit('connected');
        this.emit('ready');
        
        // Start health monitoring
        this.healthMonitor.startMonitoring();
        
        // Apply environment overrides
        this.applyEnvironmentOverrides();
        
        console.log(`✅ Connected to ${this.host}:${this.options.port} as ${this.username}`);
      });

      this.ssh.on('error', (error) => {
        this.isConnected = false;
        this.emit('error', error);
        this.handleConnectionError(error);
      });

      this.ssh.on('close', () => {
        this.isConnected = false;
        this.emit('disconnected');
        this.handleDisconnection();
      });

      this.ssh.on('end', () => {
        this.isConnected = false;
        this.emit('ended');
        this.handleDisconnection();
      });
    }
  }

  async connect() {
    try {
      if (this.isConnected) {
        logger.info('Already connected', { connection: this.sanitizeConnectionString() });
        return true;
      }

      if (this.failedAttempts >= SECURITY.MAX_FAILED_ATTEMPTS) {
        const error = new Error('Too many failed connection attempts. Please wait before retrying.');
        error.code = ERROR_CODES.SSH_CONNECTION_FAILED;
        throw error;
      }

      // Try progressive fallback if enabled
      if (this.options.progressiveFallback) {
        return await this.progressiveFallback.attemptConnection();
      } else {
        return await this.standardConnection();
      }
      
    } catch (error) {
      this.failedAttempts++;
      this.secureCleanup();
      throw error;
    }
  }

  async standardConnection() {
    const connectOptions = {
      host: this.host,
      username: this.username,
      port: parseInt(this.options.port || NETWORK.DEFAULT_SSH_PORT.toString()),
      
      // Keepalive settings
      keepalive: this.options.keepalive,
      keepaliveInterval: this.options.keepaliveInterval,
      keepaliveCountMax: this.options.keepaliveCountMax,
      serverAliveInterval: this.options.serverAliveInterval,
      serverAliveCountMax: this.options.serverAliveCountMax,
      clientAliveInterval: this.options.clientAliveInterval,
      clientAliveCountMax: this.options.clientAliveCountMax,
      keepaliveInitialDelay: this.options.keepaliveInitialDelay,
      
      // Timeout settings
      readyTimeout: this.options.readyTimeout
    };

    // Add authentication
    if (this.options.password) {
      connectOptions.password = this.options.password;
    } else if (this.options.key) {
      const keyContent = await this.loadAndValidateKey(this.options.key);
      connectOptions.privateKey = keyContent;
    } else {
      // Try to find default keys
      const keyPath = this.config.getDefaultSSHKey();
      if (keyPath && fsSync.existsSync(keyPath)) {
        const keyContent = await this.loadAndValidateKey(keyPath);
        connectOptions.privateKey = keyContent;
      }
    }

    // Validate connection parameters
    if (!ValidationUtils.validateHostname(this.host)) {
      throw new Error(`Invalid hostname: ${this.host}`);
    }

    if (!ValidationUtils.validatePort(connectOptions.port)) {
      throw new Error(`Invalid port number: ${connectOptions.port}`);
    }

    // Connect
    await this.ssh.connect(connectOptions);
    
    // Save session state if enabled
    if (this.options.saveSessionState) {
      this.sessionRecovery.saveSessionState();
    }
    
    return true;
  }

  async connectWithFallbackShell(shellStrategy) {
    const connectOptions = {
      host: this.host,
      username: this.username,
      port: parseInt(this.options.port || NETWORK.DEFAULT_SSH_PORT.toString()),
      
      // Keepalive settings
      keepalive: this.options.keepalive,
      keepaliveInterval: this.options.keepaliveInterval,
      keepaliveCountMax: this.options.keepaliveCountMax,
      
      // Shell strategy
      command: shellStrategy.command,
      args: shellStrategy.args
    };

    // Add authentication
    if (this.options.password) {
      connectOptions.password = this.options.password;
    } else if (this.options.key) {
      const keyContent = await this.loadAndValidateKey(this.options.key);
      connectOptions.privateKey = keyContent;
    }

    // Connect with shell strategy
    await this.ssh.connect(connectOptions);
    return true;
  }

  applyEnvironmentOverrides() {
    if (this.clientType === 'ssh2' && this.ssh.shell) {
      // For ssh2, we need to set environment variables in the shell
      Object.entries(this.options.environmentOverrides).forEach(([key, value]) => {
        if (value !== null) {
          this.ssh.shell.env[key] = value;
        }
      });
    }
  }

  async exec(command) {
    await this.connect();

    try {
      // Update activity time
      this.lastActivityTime = Date.now();
      
      const result = await this.ssh.execCommand(command);
      
      return {
        stdout: result.stdout,
        stderr: result.stderr,
        code: result.code,
        exitCode: result.code
      };
    } catch (error) {
      // Try to recover from execution errors
      if (this.options.sessionRecovery) {
        const recovered = await this.sessionRecovery.handleExecutionError(error);
        if (recovered) {
          // Retry the command
          return await this.exec(command);
        }
      }
      throw error;
    }
  }

  async execWithKeepalive(command, keepaliveInterval = 30000) {
    await this.connect();

    return new Promise((resolve, reject) => {
      let keepaliveTimer;
      let commandCompleted = false;

      const cleanup = () => {
        if (keepaliveTimer) clearInterval(keepaliveTimer);
        commandCompleted = true;
      };

      // Set up keepalive
      keepaliveTimer = setInterval(() => {
        if (!commandCompleted && this.isConnected) {
          this.ssh.execCommand('echo "keepalive"', (err) => {
            if (err) {
              cleanup();
              reject(new Error(`Keepalive failed: ${err.message}`));
            }
          });
        }
      }, keepaliveInterval);

      // Execute command
      this.ssh.execCommand(command, (err, stream) => {
        if (err) {
          cleanup();
          reject(err);
          return;
        }

        let stdout = '';
        let stderr = '';

        stream.on('data', (data) => {
          stdout += data.toString();
        });

        stream.stderr.on('data', (data) => {
          stderr += data.toString();
        });

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

  async testConnectionStability(duration = 30000) {
    const startTime = Date.now();
    const testInterval = 5000;
    const tests = [];
    
    return new Promise((resolve) => {
      const testTimer = setInterval(async () => {
        try {
          const testResult = await this.exec('echo "stability_test"');
          tests.push({
            timestamp: Date.now() - startTime,
            success: true,
            response: testResult.stdout.trim()
          });
          
          if (Date.now() - startTime >= duration) {
            clearInterval(testTimer);
            resolve({
              success: true,
              duration: Date.now() - startTime,
              tests: tests,
              averageResponseTime: this.calculateAverageResponseTime(tests)
            });
          }
        } catch (error) {
          tests.push({
            timestamp: Date.now() - startTime,
            success: false,
            error: error.message
          });
          
          clearInterval(testTimer);
          resolve({
            success: false,
            duration: Date.now() - startTime,
            tests: tests,
            failureReason: error.message
          });
        }
      }, testInterval);
    });
  }

  calculateAverageResponseTime(tests) {
    const successfulTests = tests.filter(t => t.success);
    if (successfulTests.length === 0) return 0;
    
    const totalTime = successfulTests.reduce((sum, test) => sum + test.timestamp, 0);
    return totalTime / successfulTests.length;
  }

  async loadAndValidateKey(keyPath) {
    try {
      if (!ValidationUtils.validateFilePath(keyPath)) {
        throw new Error('Invalid SSH key file path');
      }
      
      const isValidKeyFile = await ValidationUtils.validateSSHKeyFile(keyPath);
      if (!isValidKeyFile) {
        throw new Error('SSH key file validation failed');
      }

      const keyContent = await fs.readFile(keyPath, 'utf8');
      
      const buffer = Buffer.from(keyContent, 'utf8');
      this.secureBuffers.push(buffer);
      
      return keyContent;
    } catch (error) {
      if (error.code === 'ENOENT') {
        throw new Error(`SSH key file not found: ${keyPath}`);
      } else if (error.code === 'EACCES') {
        throw new Error(`Permission denied reading SSH key: ${keyPath}`);
      } else {
        throw new Error(`Error reading SSH key: ${error.message}`);
      }
    }
  }

  handleConnectionError(error) {
    logger.error('SSH connection error', { 
      error: error.message, 
      host: this.host,
      username: this.username 
    });
    
    // Try to recover if session recovery is enabled
    if (this.options.sessionRecovery) {
      this.sessionRecovery.handleConnectionError(error);
    }
  }

  async handleDisconnection() {
    logger.info('SSH connection disconnected', { 
      host: this.host,
      username: this.username,
      sessionDuration: this.connectionStartTime ? Date.now() - this.connectionStartTime : 0
    });
    
    // Stop health monitoring
    this.healthMonitor.stopMonitoring();
    
    // Try to recover if session recovery is enabled
    if (this.options.sessionRecovery) {
      await this.sessionRecovery.handleDisconnection();
    }
  }

  secureCleanup() {
    this.secureBuffers.forEach(buffer => {
      buffer.fill(0);
    });
    this.secureBuffers = [];
  }

  sanitizeConnectionString() {
    if (!this.connectionString) return '';
    const parts = this.connectionString.split('@');
    if (parts.length === 2) {
      return `${parts[0]}@${parts[1]}`;
    }
    return this.connectionString;
  }

  async dispose() {
    this.healthMonitor.stopMonitoring();
    
    if (this.ssh) {
      await this.ssh.dispose();
    }
    
    this.secureCleanup();
  }
}

/**
 * Connection Health Monitor
 * Monitors connection health and sends keepalive signals
 */
class ConnectionHealthMonitor {
  constructor(sshClient) {
    this.sshClient = sshClient;
    this.keepaliveTimer = null;
    this.healthCheckInterval = 30000;
    this.maxReconnectAttempts = 5;
    this.reconnectDelay = 2000;
  }
  
  startMonitoring() {
    this.keepaliveTimer = setInterval(() => {
      this.sendKeepalive();
    }, this.healthCheckInterval);
    
    logger.debug('Connection health monitoring started');
  }
  
  stopMonitoring() {
    if (this.keepaliveTimer) {
      clearInterval(this.keepaliveTimer);
      this.keepaliveTimer = null;
      logger.debug('Connection health monitoring stopped');
    }
  }
  
  async sendKeepalive() {
    if (this.sshClient.ssh && this.sshClient.isConnected) {
      try {
        await this.sshClient.ssh.execCommand('echo "keepalive"');
        logger.debug('Keepalive sent successfully');
      } catch (error) {
        logger.warn('Keepalive failed', { error: error.message });
        this.handleConnectionFailure();
      }
    }
  }
  
  handleConnectionFailure() {
    if (this.sshClient.options.sessionRecovery) {
      this.sshClient.sessionRecovery.handleConnectionFailure();
    }
  }
}

/**
 * Session Recovery
 * Handles session recovery and reconnection
 */
class SessionRecovery {
  constructor(sshClient) {
    this.sshClient = sshClient;
    this.reconnectAttempts = 0;
    this.maxReconnectAttempts = 10;
    this.reconnectDelay = 1000;
    this.sessionState = {};
  }
  
  async handleDisconnection() {
    if (this.reconnectAttempts < this.maxReconnectAttempts) {
      this.reconnectAttempts++;
      logger.info(`Connection lost. Attempting reconnection ${this.reconnectAttempts}/${this.maxReconnectAttempts}`);
      
      await this.delay(this.reconnectDelay * this.reconnectAttempts);
      
      const newConnection = await this.reconnect();
      if (newConnection) {
        this.restoreSessionState(newConnection);
        return newConnection;
      }
    }
    
    throw new Error('Max reconnection attempts exceeded');
  }
  
  async handleConnectionError(error) {
    logger.warn('Connection error, attempting recovery', { error: error.message });
    
    if (this.reconnectAttempts < this.maxReconnectAttempts) {
      await this.handleDisconnection();
    }
  }
  
  async handleExecutionError(error) {
    logger.warn('Execution error, attempting recovery', { error: error.message });
    
    // Check if it's a connection issue
    if (error.message.includes('connection') || error.message.includes('disconnected')) {
      return await this.handleDisconnection();
    }
    
    return false;
  }
  
  async handleConnectionFailure() {
    logger.warn('Connection failure detected, attempting recovery');
    return await this.handleDisconnection();
  }
  
  async reconnect() {
    try {
      await this.sshClient.connect();
      return this.sshClient;
    } catch (error) {
      logger.error('Reconnection failed', { error: error.message });
      return null;
    }
  }
  
  saveSessionState() {
    this.sessionState = {
      timestamp: Date.now(),
      host: this.sshClient.host,
      username: this.sshClient.username
    };
  }
  
  restoreSessionState(newConnection) {
    logger.info('Session state restored', { 
      host: this.sessionState.host,
      username: this.sessionState.username 
    });
  }
  
  delay(ms) {
    return new Promise(resolve => setTimeout(resolve, ms));
  }
}

/**
 * Progressive Fallback System
 * Implements multiple connection strategies with fallbacks
 */
class ProgressiveFallback {
  constructor(sshClient) {
    this.sshClient = sshClient;
    this.fallbackLevel = 0;
    this.maxFallbackLevel = 5;
  }
  
  async attemptConnection() {
    while (this.fallbackLevel <= this.maxFallbackLevel) {
      try {
        const connection = await this.tryConnectionMethod(this.fallbackLevel);
        if (connection && this.sshClient.isConnected) {
          return connection;
        }
      } catch (error) {
        logger.warn(`Fallback level ${this.fallbackLevel} failed`, { error: error.message });
        this.fallbackLevel++;
      }
    }
    
    throw new Error('All connection methods failed');
  }
  
  async tryConnectionMethod(level) {
    switch (level) {
      case 0: return await this.standardConnection();
      case 1: return await this.keepaliveConnection();
      case 2: return await this.alternativeShellConnection();
      case 3: return await this.temporaryHomeConnection();
      case 4: return await this.minimalConnection();
      case 5: return await this.emergencyConnection();
      default: throw new Error('Invalid fallback level');
    }
  }
  
  async standardConnection() {
    return await this.sshClient.standardConnection();
  }
  
  async keepaliveConnection() {
    // Use enhanced keepalive settings
    this.sshClient.options.keepaliveInterval = 15;
    this.sshClient.options.serverAliveInterval = 15;
    return await this.sshClient.standardConnection();
  }
  
  async alternativeShellConnection() {
    // Try different shell strategies
    const shellStrategies = [
      { command: 'bash', args: ['--noprofile', '--norc'] },
      { command: 'sh', args: [] },
      { command: 'dash', args: [] },
      { command: 'bash', args: ['-i'] }
    ];
    
    for (const strategy of shellStrategies) {
      try {
        return await this.sshClient.connectWithFallbackShell(strategy);
      } catch (error) {
        logger.debug(`Shell strategy ${strategy.command} failed`, { error: error.message });
        continue;
      }
    }
    
    throw new Error('All shell strategies failed');
  }
  
  async temporaryHomeConnection() {
    // Override environment to use temporary home
    this.sshClient.options.environmentOverrides.HOME = '/tmp';
    this.sshClient.options.environmentOverrides.PWD = '/tmp';
    return await this.sshClient.standardConnection();
  }
  
  async minimalConnection() {
    // Minimal connection with basic settings
    this.sshClient.options.keepalive = false;
    this.sshClient.options.readyTimeout = 30000;
    return await this.sshClient.standardConnection();
  }
  
  async emergencyConnection() {
    // Last resort - very basic connection
    this.sshClient.options.keepalive = false;
    this.sshClient.options.readyTimeout = 15000;
    this.sshClient.options.serverAliveInterval = 0;
    this.sshClient.options.clientAliveInterval = 0;
    return await this.sshClient.standardConnection();
  }
}

module.exports = { EnhancedSSHClient };