Prototype SSH version of ZX

ssh.js

// @ts-check
import { readFile } from 'fs/promises';
import { Client } from 'ssh2';
import { ProcessPromise, ProcessOutput, $ } from 'zx';
import chalk from 'chalk';

/**
 * @typedef {import('ssh2/lib/Channel.js').Channel} Channel)}
 * @typedef {import('./ssh-types').HostConfig} HostConfig
 * @typedef {import('./ssh-types').SshZx} SshZx
 */

/**
 * Example usage:
 *
 * const fredAtExample = {
 *   user: 'fred',
 *   identityFile: '/path/to/id_rsa',
 *   hostName: 'example.com',
 *   port: 22,
 * };
 *
 * await ssh(fredAtExample, $ => {
 *   $`ls -la`;
 * });
 * @param {HostConfig} hostConfig
 * @param {($: SshZx) => Promise<void>} runner
 */
export async function ssh(hostConfig, runner) {
  const {
    user: username,
    identityFile,
    hostName: host = 'localhost',
    port = 22,
    passphrase,
  } = hostConfig;

  const conn = new Client();

  // @ts-expect-error I'm not sure what's going on here
  const sshZx = createSshZx(conn, host);

  const { promise, resolve, reject } = (() => {
    let resolve = /** @type {(_value: any) => void} */ _value => {};
    let reject = /** @type {(ex: any) => void} */ () => {};
    const promise = new Promise((...args) => ([resolve, reject] = args));
    return { promise, resolve, reject };
  })();

  conn.on('ready', () => {
    runner(sshZx)
      .then(() => {
        resolve(undefined);
      })
      .catch(ex => {
        reject(ex);
      })
      .finally(() => {
        conn.end();
      });
  });

  const privateKey = await readFile(identityFile);
  conn.connect({ host, port, username, privateKey, passphrase });

  return promise;
}

/**
 * @param {Client} client
 * @param {string} hostname
 * @returns {SshZx}
 */
function createSshZx(client, hostname) {
  return (pieces, ...args) => {
    const { verbose, prefix } = $;
    const __from = /** @type {string} */ (new Error().stack)
      .split(/^\s*at\s/m)[2]
      .trim();

    let cmd = pieces[0];
    let i = 0;

    while (i < args.length) {
      /** @type {(x: any) => string} */
      const quoteSubst = x => $.quote(substitute(x));
      /** @type {string} */
      const s = Array.isArray(args[i])
        ? args[i].map(quoteSubst).join(' ')
        : quoteSubst(args[i]);
      cmd += s + pieces[++i];
    }

    const { promise, resolve, reject } = (() => {
      let resolve = /** @type {(_value: any) => void} */ _value => {};
      let reject = /** @type {(ex: any) => void} */ () => {};
      // @ts-expect-error Exported zx types don't match the actual types
      const promise = new ProcessPromise(
        // @ts-expect-error Exported zx types don't match the actual types
        (...args) => ([resolve, reject] = args)
      );
      return { promise, resolve, reject };
    })();

    if (resolve == null || reject == null) {
      throw new Error('resolve or reject is null');
    }

    promise._run = () => {
      if (promise.client != null) {
        return;
      }
      if (promise._prerun != null) {
        promise._prerun();
      }
      if (verbose) {
        printCmd(cmd, hostname);
      }

      if (promise._piped) {
        throw new Error('piped is not supported');
      }

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

      /**
       * @param {Error} err
       * @param {Channel} channel
       */
      function handleExec(err, channel) {
        if (err) {
          reject(err);
          return;
        }

        /**
         * @param {number} code
         * @param {number | undefined} _signal
         */
        function handleClose(code, _signal) {
          const message =
            `${stderr || '\n'}    at ${__from}\n    exit code: ${code}` +
            (exitCodeInfo(code) ? ` (${exitCodeInfo(code)})` : '');

          // @ts-expect-error Exported zx types don't match the actual types
          const output = new ProcessOutput({
            code,
            stdout,
            stderr,
            combined,
            message,
          });

          if (code === 0 || promise._nothrow) {
            resolve(output);
          } else {
            reject(output);
          }
          promise._resolved = true;
        }
        channel.on('close', handleClose);

        /** @param {string} data */
        function handleStdoutData(data) {
          if (verbose) {
            process.stdout.write(data);
          }
          stdout += data;
          combined += data;
        }
        channel.on('data', handleStdoutData);

        /** @param {string} data */
        function handleStderrData(data) {
          if (verbose) {
            process.stderr.write(data);
          }
          stderr += data;
          combined += data;
        }
        channel.stderr.on('data', handleStderrData);
      }

      // @ts-expect-error The exec method does exist because this works
      client.exec(prefix + cmd, handleExec);

      promise.client = client;
      if (promise._postrun) {
        promise._postrun();
      }
    };

    // eslint-disable-next-line @typescript-eslint/no-implied-eval, @typescript-eslint/no-unsafe-argument
    setTimeout(promise._run, 0); // Make sure all subprocesses started.

    return promise;
  };
}

/**
 * @param {string} cmd
 * @param {string} hostname
 */
function printCmd(cmd, hostname) {
  if (/\n/.test(cmd)) {
    console.log(
      cmd
        .split('\n')
        .map((line, i) => (i === 0 ? '$' : '>') + ' ' + colorize(line))
        .join('\n')
    );
  } else {
    console.log(hostname, '$', colorize(cmd));
  }
}

/**
 * @param {string} cmd
 */
function colorize(cmd) {
  return cmd.replace(/^[\w_.-]+(\s|$)/, substr => {
    return chalk.greenBright(substr);
  });
}

/**
 * @param {any} arg
 */
function substitute(arg) {
  if (arg instanceof ProcessOutput) {
    return arg.stdout.replace(/\n$/, '');
  }
  return `${String(arg)}`;
}

/**
 * @param {number} exitCode
 * @returns {string}
 */
function exitCodeInfo(exitCode) {
  return (
    {
      2: 'Misuse of shell builtins',
      126: 'Invoked command cannot execute',
      127: 'Command not found',
      128: 'Invalid exit argument',
      129: 'Hangup',
      130: 'Interrupt',
      131: 'Quit and dump core',
      132: 'Illegal instruction',
      133: 'Trace/breakpoint trap',
      134: 'Process aborted',
      135: 'Bus error: "access to undefined portion of memory object"',
      136: 'Floating point exception: "erroneous arithmetic operation"',
      137: 'Kill (terminate immediately)',
      138: 'User-defined 1',
      139: 'Segmentation violation',
      140: 'User-defined 2',
      141: 'Write to pipe with no one reading',
      142: 'Signal raised by alarm',
      143: 'Termination (request to terminate)',
      145: 'Child process terminated, stopped (or continued*)',
      146: 'Continue if stopped',
      147: 'Stop executing temporarily',
      148: 'Terminal stop signal',
      149: 'Background process attempting to read from tty ("in")',
      150: 'Background process attempting to write to tty ("out")',
      151: 'Urgent data available on socket',
      152: 'CPU time limit exceeded',
      153: 'File size limit exceeded',
      154: 'Signal raised by timer counting virtual time: "virtual timer expired"',
      155: 'Profiling timer expired',
      157: 'Pollable event',
      159: 'Bad syscall',
    }[exitCode] || `Unknown exit code: ${exitCode}`
  );
}