ssh-bridge-ai/src/ssh-key-manager.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 crypto = require('crypto');
const { execSync, spawn } = require('child_process');
const { SSHClient } = require('./ssh');
const { Config } = require('./config');
const logger = require('./utils/logger');

/**
 * SSH Key Manager - Handles key generation, copying, and management
 */
class SSHKeyManager {
  constructor() {
    this.sshDir = path.join(process.env.HOME || os.homedir(), '.ssh');
    this.config = new Config();
  }

  /**
   * Generate a new SSH key pair
   */
  async generateKey(options = {}) {
    const {
      type = 'ed25519',
      bits = '4096',
      output = '',
      email = '',
      passphrase = false,
      force = false
    } = options;

    // Validate key type
    const validTypes = ['ed25519', 'rsa', 'ecdsa', 'dsa'];
    if (!validTypes.includes(type)) {
      throw new Error(`Invalid key type: ${type}. Supported types: ${validTypes.join(', ')}`);
    }

    // Determine output path
    let keyPath = output;
    if (!keyPath) {
      const defaultName = `id_${type}`;
      keyPath = path.join(this.sshDir, defaultName);
    } else {
      // Expand ~ to home directory
      if (keyPath.startsWith('~')) {
        keyPath = path.join(process.env.HOME || os.homedir(), keyPath.slice(1));
      }
      // Ensure .ssh directory exists
      const keyDir = path.dirname(keyPath);
      if (keyDir !== this.sshDir) {
        await fs.mkdir(keyDir, { recursive: true, mode: 0o700 });
      }
    }

    // Check if key already exists
    if (fsSync.existsSync(keyPath) && !force) {
      throw new Error(`Key already exists: ${keyPath}. Use --force to overwrite.`);
    }

    // Ensure .ssh directory exists with correct permissions
    await this.ensureSSHDirectory();

    // Build ssh-keygen command
    const args = ['-t', type];
    
    if (type === 'rsa' && bits) {
      args.push('-b', bits);
    }
    
    if (email) {
      args.push('-C', email);
    }
    
    if (passphrase) {
      args.push('-N', passphrase); // Use provided passphrase
    } else {
      args.push('-N', '""'); // No passphrase (empty string)
    }
    
    // Use quiet mode for clean generation
    args.push('-q');
    args.push('-f', keyPath);

    try {
      // Generate the key
      console.log(`Generating ${type.toUpperCase()} key...`);
      
      // Use spawn for better error handling
      return new Promise((resolve, reject) => {
        const child = spawn('ssh-keygen', args, { 
          stdio: 'pipe'
        });
        
        let stdout = '';
        let stderr = '';
        
        child.stdout.on('data', (data) => {
          stdout += data.toString();
        });
        
        child.stderr.on('data', (data) => {
          stderr += data.toString();
        });
        
        child.on('close', async (code) => {
          if (code === 0) {
            try {
              // Set correct permissions
              await fs.chmod(keyPath, 0o600);
              await fs.chmod(`${keyPath}.pub`, 0o644);

              // Verify key was created
              if (!fsSync.existsSync(keyPath) || !fsSync.existsSync(`${keyPath}.pub`)) {
                reject(new Error('Key generation failed - files not created'));
                return;
              }

              // Get key fingerprint and randomart after generation
              const fingerprint = await this.getKeyFingerprint(`${keyPath}.pub`);
              const randomart = await this.getKeyRandomart(keyPath);

              logger.info('SSH key generated successfully', { 
                type, 
                path: keyPath,
                publicKey: `${keyPath}.pub`,
                randomart: randomart ? 'generated' : 'not found'
              });

              resolve({
                privateKeyPath: keyPath,
                publicKeyPath: `${keyPath}.pub`,
                type: type,
                email: email,
                randomart: randomart,
                fingerprint: fingerprint
              });
            } catch (error) {
              reject(error);
            }
          } else {
            reject(new Error(`ssh-keygen failed with exit code ${code}: ${stderr}`));
          }
        });
        
        child.on('error', (error) => {
          reject(new Error(`Failed to spawn ssh-keygen: ${error.message}`));
        });
      });

    } catch (error) {
      // Clean up partial files
      try {
        if (fsSync.existsSync(keyPath)) await fs.unlink(keyPath);
        if (fsSync.existsSync(`${keyPath}.pub`)) await fs.unlink(`${keyPath}.pub`);
      } catch (cleanupError) {
        logger.warn('Failed to cleanup partial key files', { error: cleanupError.message });
      }

      // ssh-keygen returns status 1 for various reasons, not just user cancellation
      // Let's provide a more helpful error message
      if (error.status === 1) {
        throw new Error(`Key generation failed with exit code 1: ${error.message}`);
      }
      
      throw new Error(`Key generation failed: ${error.message}`);
    }
  }

  /**
   * Extract randomart from ssh-keygen output
   */
  extractRandomart(output) {
    try {
      // Look for the randomart pattern in the output
      // Pattern for ssh-keygen -lv output - matches both formats:
      // +---[RSA 2048]----+ and +--[ED25519 256]--+
      const randomartMatch = output.match(/(\+-+\[[^\]]+\]-+\+[\s\S]*?\+----\[[^\]]+\]-----)/);
      if (randomartMatch && randomartMatch[1]) {
        return randomartMatch[1].trim();
      }
      
      // Alternative pattern for some ssh-keygen versions
      const altMatch = output.match(/randomart image is:\s*\n([\s\S]*?)(?=\n\n|\n$|$)/);
      if (altMatch && altMatch[1]) {
        return altMatch[1].trim();
      }
      
      return null;
    } catch (error) {
      logger.warn('Failed to extract randomart', { error: error.message });
      return null;
    }
  }

  /**
   * Get key fingerprint using ssh-keygen
   */
  async getKeyFingerprint(publicKeyPath) {
    try {
      const { execSync } = require('child_process');
      const output = execSync(`ssh-keygen -lf "${publicKeyPath}"`, { encoding: 'utf8' });
      const match = output.match(/^\d+\s+([a-f0-9:]+)\s+/);
      return match ? match[1] : null;
    } catch (error) {
      logger.warn('Failed to get key fingerprint', { error: error.message });
      return null;
    }
  }

  /**
   * Copy SSH public key to remote server
   */
  async copyKeyToServer(server, options = {}) {
    const {
      key = '',
      port = '22',
      force = false,
      output = '',
      passphrase = null
    } = options;

    // Parse server string
    let username, hostname;
    if (server.includes('@')) {
      [username, hostname] = server.split('@');
    } else {
      // Try to get from config
      const config = this.config.getConfig();
      const serverConfig = config[server];
      if (!serverConfig) {
        throw new Error(`Server '${server}' not found in config. Use user@hostname format or add to config.`);
      }
      username = serverConfig.username;
      hostname = serverConfig.hostname;
    }

    // Determine which key to use
    let keyPath = key;
    if (!keyPath) {
      // Auto-detect best available key
      const keys = await this.listKeys();
      if (keys.length === 0) {
        throw new Error('No SSH keys found. Generate one with: sshbridge keygen');
      }
      keyPath = keys[0].path; // Use first available key
    }

    // Validate key exists
    if (!fsSync.existsSync(keyPath)) {
      throw new Error(`SSH key not found: ${keyPath}`);
    }

    // Read public key
    const publicKeyPath = `${keyPath}.pub`;
    if (!fsSync.existsSync(publicKeyPath)) {
      throw new Error(`Public key not found: ${publicKeyPath}`);
    }

    const publicKey = await fs.readFile(publicKeyPath, 'utf8');
    const publicKeyContent = publicKey.trim();

    // Connect to server and copy key
    const connection = `${username}@${hostname}`;
    const sshOptions = { port: parseInt(port) };
    
    // Add key and passphrase if provided
    if (keyPath) {
      sshOptions.key = keyPath;
      if (passphrase) {
        sshOptions.passphrase = passphrase;
        logger.debug('Passphrase provided for encrypted key', { keyPath });
      } else {
        logger.debug('No passphrase provided, key may be unencrypted', { keyPath });
      }
    }
    
    logger.debug('SSH connection options', { 
      connection, 
      port: sshOptions.port, 
      hasKey: !!sshOptions.key, 
      hasPassphrase: !!sshOptions.passphrase 
    });
    
    // Test key decryption before connection
    if (keyPath && passphrase) {
      try {
        logger.debug('Testing key decryption before connection...');
        // Test key loading without actual connection
        const testKey = await fs.readFile(keyPath);
        logger.debug('Key decryption test passed');
      } catch (keyError) {
        logger.error('Key decryption test failed', { error: keyError.message });
        if (keyError.message.includes('integrity check failed') || keyError.message.includes('bad passphrase')) {
          throw new Error(`SSH key decryption failed: ${keyError.message}. Please verify your passphrase is correct.`);
        }
        throw keyError;
      }
    }
    
    const ssh = new SSHClient(connection, sshOptions);

    try {
      await ssh.connect();

      // Determine authorized_keys path
      const authKeysPath = output || '~/.ssh/authorized_keys';
      
      // Check if authorized_keys exists
      const checkResult = await ssh.exec(`test -f ${authKeysPath} && echo "exists" || echo "missing"`);
      const authKeysExists = checkResult.stdout.trim() === 'exists';

      if (authKeysExists && !force) {
        // Check if key already exists
        const grepResult = await ssh.exec(`grep -F "${publicKeyContent}" ${authKeysPath} || echo "not_found"`);
        if (grepResult.stdout.trim() !== 'not_found') {
          return {
            server: connection,
            keyPath: keyPath,
            status: 'Key already exists on server',
            action: 'none'
          };
        }
      }

      // Create .ssh directory if it doesn't exist
      await ssh.exec(`mkdir -p ~/.ssh`);
      await ssh.exec(`chmod 700 ~/.ssh`);

      // Add key to authorized_keys
      if (authKeysExists) {
        // Append to existing file
        await ssh.exec(`echo "${publicKeyContent}" >> ${authKeysPath}`);
      } else {
        // Create new file
        await ssh.exec(`echo "${publicKeyContent}" > ${authKeysPath}`);
      }

      // Set correct permissions
      await ssh.exec(`chmod 600 ${authKeysPath}`);

      logger.info('SSH key copied to server successfully', { 
        server: connection, 
        keyPath,
        authKeysPath 
      });

      return {
        server: connection,
        keyPath: keyPath,
        status: 'Key copied successfully',
        action: authKeysExists ? 'appended' : 'created'
      };

    } catch (error) {
      // Log the actual error for debugging
      logger.error('SSH connection failed in copyKeyToServer', { 
        error: error.message, 
        stack: error.stack,
        connection,
        hasPassphrase: !!passphrase 
      });
      throw new Error(`Failed to copy key to server: ${error.message}`);
    } finally {
      await ssh.dispose();
    }
  }

  /**
   * List all available SSH keys
   */
  async listKeys() {
    const keys = [];
    const defaultKey = this.config.getDefaultSSHKey();

    try {
      // Ensure .ssh directory exists
      if (!fsSync.existsSync(this.sshDir)) {
        return keys;
      }

      const files = await fs.readdir(this.sshDir);
      
      for (const file of files) {
        if (file.endsWith('.pub')) continue; // Skip public keys
        if (file.includes('.pub.')) continue; // Skip backup files
        
        const keyPath = path.join(this.sshDir, file);
        const publicKeyPath = `${keyPath}.pub`;
        
        // Check if this is a valid key pair
        if (!fsSync.existsSync(publicKeyPath)) continue;
        
        try {
          const stats = await fs.stat(keyPath);
          const publicKeyStats = await fs.stat(publicKeyPath);
          
          // Check permissions (private key should be 600, public key 644)
          const privateMode = stats.mode & 0o777;
          const publicMode = publicKeyStats.mode & 0o777;
          
          if (privateMode !== 0o600 || publicMode !== 0o644) continue;
          
          // Get key type and fingerprint
          const keyInfo = await this.getKeyInfo(keyPath);
          
          keys.push({
            name: file,
            path: keyPath,
            type: keyInfo.type,
            bits: keyInfo.bits,
            fingerprint: keyInfo.fingerprint,
            isDefault: keyPath === defaultKey,
            size: stats.size,
            modified: stats.mtime
          });
          
        } catch (error) {
          logger.debug(`Skipping invalid key file: ${file}`, { error: error.message });
          continue;
        }
      }

      // Sort keys: default first, then by type (ed25519 preferred), then by name
      keys.sort((a, b) => {
        if (a.isDefault && !b.isDefault) return -1;
        if (!a.isDefault && b.isDefault) return 1;
        
        const typePriority = { 'ed25519': 1, 'ecdsa': 2, 'rsa': 3, 'dsa': 4 };
        const aPriority = typePriority[a.type] || 5;
        const bPriority = typePriority[b.type] || 5;
        
        if (aPriority !== bPriority) return aPriority - bPriority;
        return a.name.localeCompare(b.name);
      });

    } catch (error) {
      logger.error('Failed to list SSH keys', { error: error.message });
    }

    return keys;
  }

  /**
   * Get key information (type, bits, fingerprint)
   */
  async getKeyInfo(keyPath) {
    try {
      // Try to get fingerprint using ssh-keygen
      const fingerprint = execSync(`ssh-keygen -lf "${keyPath}"`, { 
        encoding: 'utf8',
        stdio: ['pipe', 'pipe', 'pipe']
      }).trim();

      // Parse fingerprint output: "2048 SHA256:abc123... user@host (RSA)"
      const match = fingerprint.match(/^(\d+)\s+(\w+):([a-f0-9]+)\s+(.+?)\s+\((\w+)\)$/);
      
      if (match) {
        return {
          bits: parseInt(match[1]),
          hash: match[2],
          fingerprint: match[3],
          comment: match[4],
          type: match[5].toLowerCase()
        };
      } else {
        // Fallback: just return the fingerprint
        return {
          type: 'unknown',
          bits: null,
          fingerprint: fingerprint.split(/\s+/)[1] || fingerprint
        };
      }
      
    } catch (error) {
      logger.warn('Failed to get key info using ssh-keygen', { error: error.message });
      
      // Fallback: try to determine type from file content
      try {
        const keyContent = await fs.readFile(keyPath, 'utf8');
        
        if (keyContent.includes('BEGIN OPENSSH PRIVATE KEY')) {
          return { type: 'openssh', bits: null, fingerprint: 'unknown' };
        } else if (keyContent.includes('BEGIN RSA PRIVATE KEY')) {
          return { type: 'rsa', bits: null, fingerprint: 'unknown' };
        } else if (keyContent.includes('BEGIN EC PRIVATE KEY')) {
          return { type: 'ecdsa', bits: null, fingerprint: 'unknown' };
        } else if (keyContent.includes('BEGIN DSA PRIVATE KEY')) {
          return { type: 'dsa', bits: null, fingerprint: 'unknown' };
        } else {
          return { type: 'unknown', bits: null, fingerprint: 'unknown' };
        }
      } catch (readError) {
        return { type: 'unknown', bits: null, fingerprint: 'unknown' };
      }
    }
  }

  /**
   * Get key fingerprint
   */
  async getKeyFingerprint(keyPath) {
    const keyInfo = await this.getKeyInfo(keyPath);
    return keyInfo.fingerprint;
  }

  /**
   * Remove SSH key
   */
  async removeKey(keyPath) {
    // Validate key path
    if (!fsSync.existsSync(keyPath)) {
      throw new Error(`SSH key not found: ${keyPath}`);
    }

    // Check if it's a private key
    if (!keyPath.endsWith('.pub') && !keyPath.includes('id_')) {
      throw new Error(`Invalid SSH key path: ${keyPath}`);
    }

    // Determine public key path
    const publicKeyPath = keyPath.endsWith('.pub') ? keyPath : `${keyPath}.pub`;
    const privateKeyPath = keyPath.endsWith('.pub') ? keyPath.replace('.pub', '') : keyPath;

    // Remove both files
    const filesToRemove = [];
    if (fsSync.existsSync(privateKeyPath)) filesToRemove.push(privateKeyPath);
    if (fsSync.existsSync(publicKeyPath)) filesToRemove.push(publicKeyPath);

    for (const file of filesToRemove) {
      await fs.unlink(file);
      logger.info('SSH key file removed', { path: file });
    }

    return {
      removed: filesToRemove,
      message: `Removed ${filesToRemove.length} key file(s)`
    };
  }

  /**
   * Ensure .ssh directory exists with correct permissions
   */
  async ensureSSHDirectory() {
    if (!fsSync.existsSync(this.sshDir)) {
      await fs.mkdir(this.sshDir, { recursive: true, mode: 0o700 });
    } else {
      // Check and fix permissions
      const stats = await fs.stat(this.sshDir);
      const mode = stats.mode & 0o777;
      if (mode !== 0o700) {
        await fs.chmod(this.sshDir, 0o700);
        logger.info('Fixed .ssh directory permissions', { oldMode: mode.toString(8), newMode: '700' });
      }
    }
  }

  /**
   * Get randomart for an existing SSH key
   */
  async getKeyRandomart(keyPath) {
    try {
      const { execSync } = require('child_process');
      
      // Determine if it's a public or private key
      const isPublic = keyPath.endsWith('.pub');
      const publicKeyPath = isPublic ? keyPath : `${keyPath}.pub`;
      
      if (!fsSync.existsSync(publicKeyPath)) {
        throw new Error(`Public key not found: ${publicKeyPath}`);
      }
      
      // Use ssh-keygen -lv to get verbose output including randomart
      const output = execSync(`ssh-keygen -lv -f "${publicKeyPath}"`, { encoding: 'utf8' });
      
      // Debug: log the output
      logger.debug('ssh-keygen -lv output', { 
        keyPath: publicKeyPath, 
        output: output.substring(0, 200) + '...' 
      });
      

      
      // Extract randomart from the output
      const randomart = this.extractRandomart(output);
      
      if (!randomart) {
        logger.warn('No randomart found in ssh-keygen output', { 
          keyPath: publicKeyPath,
          outputLength: output.length 
        });
      }
      
      return randomart;
      
    } catch (error) {
      logger.warn('Failed to get key randomart', { error: error.message });
      return null;
    }
  }

  /**
   * Test SSH key authentication
   */
  async testKeyAuthentication(server, keyPath) {
    try {
      const ssh = new SSHClient(server, { key: keyPath });
      await ssh.connect();
      await ssh.dispose();
      return { success: true, message: 'Authentication successful' };
    } catch (error) {
      return { success: false, error: error.message };
    }
  }
}

module.exports = { SSHKeyManager };