alwaysai/src/util/secure-tunnel-ssh.ts

The alwaysAI command-line interface (CLI)

import { CliTerseError } from '@alwaysai/alwayscli';
import { spawn } from 'child_process';
import { ALWAYSAI_OS_PLATFORM } from '../environment';
import { PRIVATE_KEY_FILE_PATH } from '../paths';
import { logger } from './logger';
import rKill = require('tree-kill');

interface SecureTunnelSshBase {
  processArgs(args: string[]): string[];
  runInteractiveSshAsync(args: string[]): void;
}

export class SecureTunnelInteractiveSsh implements SecureTunnelSshBase {
  private targetHost: string;
  private sshPort: number;

  constructor(opts: { targetHost: string; sshPort: number }) {
    const { targetHost, sshPort } = opts;
    this.targetHost = targetHost;
    this.sshPort = sshPort;
  }

  /**
   * processArgs - process arguments for the ssh
   * @param {string[]} args - arguments to process
   * @returns {string[]} - processed arguments
   */
  public processArgs(args: string[]): string[] {
    logger.debug(`-> SecureTunnelInteractiveSsh:processArgs(args = ${args})`);
    const processedArgs: string[] = [
      '-i',
      PRIVATE_KEY_FILE_PATH,
      '-o',
      'StrictHostKeyChecking=no',
      '-t',
      ...(this.sshPort > 0 ? ['-p', `${this.sshPort}`] : []),
      ...args,
      this.targetHost
    ];
    logger.debug(
      `<- SecureTunnelInteractiveSsh:processArgs(processedArgs = ${processedArgs})`
    );
    return processedArgs;
  }

  /**
   * runInteractiveSshAsync - run the ssh command in interactive mode
   * @param {string[]} args - arguments to use with ssh
   */
  public async runInteractiveSshAsync(args: string[]) {
    logger.debug(
      `-> SecureTunnelInteractiveSsh:runInteractiveSshAsync(args = ${args})`
    );
    const sshCommand = `ssh ${args.join(' ')}`;

    const sshSession = spawn(sshCommand, {
      stdio: [process.stdin, process.stdout, process.stderr],
      shell: true
    });

    sshSession.on('close', (code) => {
      console.log(`SSH session closed with code ${code}`);
    });

    await new Promise<void>((resolve) => {
      sshSession.on('exit', () => {
        if (ALWAYSAI_OS_PLATFORM === 'darwin') {
          process.on('SIGINT', () => {
            if (sshSession.pid) {
              rKill(sshSession.pid, 'SIGINT', (error?: Error | undefined) => {
                if (error) {
                  throw new CliTerseError(
                    `Failed to kill child process ${error.name} ${error.message}`
                  );
                }
              });
            }
          });
        }
        resolve();
      });
    });

    logger.debug('<- SecureTunnelInteractiveSsh:runInteractiveSshAsync()');
    return sshSession;
  }
}