HOME

#!/usr/bin/env node

'use strict';

const EventEmitter = require('events');
const fs = require('fs');
const HttpsProxyAgent = require('https-proxy-agent');
const program = require('commander');
const read = require('read');
const readline = require('readline');
const tty = require('tty');
const WebSocket = require('ws');

/**
 * InputReader - processes console input.
 *
 * @extends EventEmitter
 */
class Console extends EventEmitter {
  constructor() {
    super();

    this.stdin = process.stdin;
    this.stdout = process.stdout;
    this.stderr = process.stderr;

    this.readlineInterface = readline.createInterface(this.stdin, this.stdout);

    this.readlineInterface
      .on('line', (data) => {
        this.emit('line', data);
      })
      .on('close', () => {
        this.emit('close');
      });

    this._resetInput = () => {
      this.clear();
    };
  }

  static get Colors() {
    return {
      Red: '\u001b[31m',
      Green: '\u001b[32m',
      Yellow: '\u001b[33m',
      Blue: '\u001b[34m',
      Default: '\u001b[39m'
    };
  }

  static get Types() {
    return {
      Incoming: '< ',
      Control: '',
      Error: 'error: '
    };
  }

  prompt() {
    this.readlineInterface.prompt(true);
  }

  print(type, msg, color) {
    if (tty.isatty(1)) {
      this.clear();

      if (programOptions.execute) color = type = '';
      else if (!programOptions.color) color = '';

      this.stdout.write(color + type + msg + Console.Colors.Default + '\n');
      this.prompt();
    } else if (type === Console.Types.Incoming) {
      this.stdout.write(msg + '\n');
    } else if (type === Console.Types.Error) {
      this.stderr.write(type + msg + '\n');
    } else {
      // is a control message and we're not in a tty... drop it.
    }
  }

  clear() {
    if (tty.isatty(1)) {
      this.stdout.write('\r\u001b[2K\u001b[3D');
    }
  }

  pause() {
    this.stdin.on('keypress', this._resetInput);
  }

  resume() {
    this.stdin.removeListener('keypress', this._resetInput);
  }
}

function collect(val, memo) {
  memo.push(val);
  return memo;
}

function noop() {}

/**
 * The actual application
 */
const version = require('../package.json').version;

program
  .version(version)
  .usage('[options] (--listen <port> | --connect <url>)')
  .option(
    '--auth <username:password>',
    'add basic HTTP authentication header (--connect only)'
  )
  .option('--ca <ca>', 'specify a Certificate Authority (--connect only)')
  .option('--cert <cert>', 'specify a Client SSL Certificate (--connect only)')
  .option('--host <host>', 'optional host')
  .option(
    '--key <key>',
    "specify a Client SSL Certificate's key (--connect only)"
  )
  .option(
    '--max-redirects [num]',
    'maximum number of redirects allowed (--connect only)',
    parseInt,
    10
  )
  .option('--no-color', 'run without color')
  .option(
    '--passphrase [passphrase]',
    "specify a Client SSL Certificate Key's passphrase (--connect only). " +
      "If you don't provide a value, it will be prompted for"
  )
  .option(
    '--proxy <[protocol://]host[:port]>',
    'connect via a proxy. Proxy must support CONNECT method'
  )
  .option(
    '--slash',
    'enable slash commands for control frames (/ping [data], /pong [data], ' +
      '/close [code [, reason]])'
  )
  .option('-c, --connect <url>', 'connect to a WebSocket server')
  .option(
    '-H, --header <header:value>',
    'set an HTTP header. Repeat to set multiple (--connect only)',
    collect,
    []
  )
  .option('-L, --location', 'follow redirects (--connect only)')
  .option('-l, --listen <port>', 'listen on port')
  .option('-n, --no-check', 'do not check for unauthorized certificates')
  .option('-o, --origin <origin>', 'optional origin')
  .option('-p, --protocol <version>', 'optional protocol version')
  .option(
    '-P, --show-ping-pong',
    'print a notification when a ping or pong is received'
  )
  .option('-s, --subprotocol <protocol>', 'optional subprotocol', collect, [])
  .option('-w, --wait <seconds>', 'wait given seconds after executing command')
  .option('-x, --execute <command>', 'execute command after connecting')
  .parse(process.argv);

const programOptions = program.opts();

if (programOptions.listen && programOptions.connect) {
  console.error('\u001b[33merror: Use either --listen or --connect\u001b[39m');
  process.exit(-1);
}

if (programOptions.listen) {
  const wsConsole = new Console();
  wsConsole.pause();

  let ws = null;
  const wss = new WebSocket.Server({ port: programOptions.listen }, () => {
    wsConsole.print(
      Console.Types.Control,
      `Listening on port ${programOptions.listen} (press CTRL+C to quit)`,
      Console.Colors.Green
    );
    wsConsole.clear();
  });

  wsConsole.on('close', () => {
    if (ws) ws.close();
    process.exit(0);
  });

  wsConsole.on('line', (data) => {
    if (ws) {
      ws.send(data);
      wsConsole.prompt();
    }
  });

  wss.on('connection', (newClient) => {
    if (ws) return newClient.terminate();

    ws = newClient;
    wsConsole.resume();
    wsConsole.prompt();
    wsConsole.print(
      Console.Types.Control,
      'Client connected',
      Console.Colors.Green
    );

    ws.on('close', (code, reason) => {
      wsConsole.print(
        Console.Types.Control,
        `Disconnected (code: ${code}, reason: "${reason}")`,
        Console.Colors.Green
      );
      wsConsole.clear();
      wsConsole.pause();
      ws = null;
    });

    ws.on('error', (err) => {
      wsConsole.print(Console.Types.Error, err.message, Console.Colors.Yellow);
    });

    ws.on('message', (data) => {
      wsConsole.print(Console.Types.Incoming, data, Console.Colors.Blue);
    });
  });

  wss.on('error', (err) => {
    wsConsole.print(Console.Types.Error, err.message, Console.Colors.Yellow);
    process.exit(-1);
  });
} else if (programOptions.connect) {
  const options = {};
  const cont = () => {
    const wsConsole = new Console();

    const headers = programOptions.header.reduce((acc, cur) => {
      const i = cur.indexOf(':');
      const key = cur.slice(0, i);
      const value = cur.slice(i + 1);
      acc[key] = value;
      return acc;
    }, {});

    if (programOptions.auth) {
      headers.Authorization =
        'Basic ' + Buffer.from(programOptions.auth).toString('base64');
    }
    if (programOptions.host) headers.Host = programOptions.host;
    if (programOptions.protocol)
      options.protocolVersion = +programOptions.protocol;
    if (programOptions.origin) options.origin = programOptions.origin;
    if (!programOptions.check)
      options.rejectUnauthorized = programOptions.check;
    if (programOptions.ca) options.ca = fs.readFileSync(programOptions.ca);
    if (programOptions.cert)
      options.cert = fs.readFileSync(programOptions.cert);
    if (programOptions.key) options.key = fs.readFileSync(programOptions.key);
    if (programOptions.proxy)
      options.agent = new HttpsProxyAgent(programOptions.proxy);
    if (programOptions.location) options.followRedirects = true;

    options.maxRedirects = programOptions.maxRedirects;
    options.headers = headers;

    let connectUrl = programOptions.connect;
    if (!connectUrl.match(/\w+:\/\/.*$/i)) {
      connectUrl = `ws://${connectUrl}`;
    }

    const ws = new WebSocket(connectUrl, programOptions.subprotocol, options);

    ws.on('open', () => {
      if (programOptions.execute) {
        ws.send(programOptions.execute);
        setTimeout(
          () => {
            ws.close();
          },
          programOptions.wait ? programOptions.wait * 1000 : 2000
        );
      } else {
        wsConsole.print(
          Console.Types.Control,
          'Connected (press CTRL+C to quit)',
          Console.Colors.Green
        );

        wsConsole.on('line', (data) => {
          if (programOptions.slash && data[0] === '/') {
            const toks = data.split(/\s+/);
            switch (toks[0].substr(1)) {
            case 'ping':
                if (toks.length >= 2) {
                    ws.ping(toks[1]);
                } else {
                    ws.ping(noop);
                }
                break;
            case 'pong':
                if (toks.length >= 2) {
                    ws.pong(toks[1]);
                } else {
                    ws.pong(noop);
                }
                break;
              case 'close': {
                let closeStatusCode = 1000;
                let closeReason = '';
                if (toks.length >= 2) {
                  closeStatusCode = parseInt(toks[1]);
                }
                if (toks.length >= 3) {
                  closeReason = toks.slice(2).join(' ');
                }
                if (closeReason.length > 0) {
                  ws.close(closeStatusCode, closeReason);
                } else {
                  ws.close(closeStatusCode);
                }
                break;
              }
              default:
                wsConsole.print(
                  Console.Types.Error,
                  'Unrecognized slash command.',
                  Console.Colors.Yellow
                );
            }
          } else {
            ws.send(data);
          }
          wsConsole.prompt();
        });
      }
    });

    ws.on('close', (code, reason) => {
      if (!programOptions.execute) {
        wsConsole.print(
          Console.Types.Control,
          `Disconnected (code: ${code}, reason: "${reason}")`,
          Console.Colors.Green
        );
      }
      wsConsole.clear();
      process.exit();
    });

    ws.on('error', (err) => {
      wsConsole.print(Console.Types.Error, err.message, Console.Colors.Yellow);
      process.exit(-1);
    });

    ws.on('message', (data) => {
      wsConsole.print(Console.Types.Incoming, data, Console.Colors.Blue);
    });

    ws.on('ping', (data) => {
      if (programOptions.showPingPong) {
        wsConsole.print(
          Console.Types.Incoming,
          `Received ping (data: "${data}"`,
          Console.Colors.Blue
        );
      }
    });

    ws.on('pong', (data) => {
      if (programOptions.showPingPong) {
        wsConsole.print(
          Console.Types.Incoming,
          `Received pong (data: "${data}")`,
          Console.Colors.Blue
        );
      }
    });

    wsConsole.on('close', () => {
      ws.close();
      process.exit();
    });
  };

  if (programOptions.passphrase === true) {
    read(
      {
        prompt: 'Passphrase: ',
        silent: true,
        replace: '*'
      },
      (err, passphrase) => {
        options.passphrase = passphrase;
        cont();
      }
    );
  } else if (typeof programOptions.passphrase === 'string') {
    options.passphrase = programOptions.passphrase;
    cont();
  } else {
    cont();
  }
} else {
  program.help();
}