<commit_before>/* @flow */

import Path from 'path'
import SSH2 from 'ssh2'
import invariant from 'assert'
import scanDirectory from 'sb-scandir'
import shellEscape from 'shell-escape'
import * as Helpers from './helpers'
import type { ConfigGiven } from './types'

class SSH {
  connection: ?SSH2;
  constructor() {
    this.connection = null
  }
  connect(givenConfig: ConfigGiven): Promise<this> {
    const connection = this.connection = new SSH2()
    return new Promise(function(resolve) {
      resolve(Helpers.normalizeConfig(givenConfig))
    }).then(config =>
      new Promise((resolve, reject) => {
        connection.on('error', reject)
        connection.on('ready', () => {
          connection.removeListener('error', reject)
          resolve(this)
        })
        connection.on('end', () => {
          if (this.connection === connection) {
            this.connection = null
          }
        })
        connection.connect(config)
      })
    )
  }
  async requestShell(): Promise<Object> {
    const connection = this.connection
    invariant(connection, 'Not connected to server')
    return await new Promise(function(resolve, reject) {
      connection.shell(Helpers.generateCallback(resolve, reject))
    })
  }
  async requestSFTP(): Promise<Object> {
    const connection = this.connection
    invariant(connection, 'Not connected to server')
    return await new Promise(function(resolve, reject) {
      connection.sftp(Helpers.generateCallback(resolve, reject))
    })
  }
  async mkdir(path: string, type: 'exec' | 'sftp' = 'sftp', givenSftp: ?Object = null): Promise<void> {
    invariant(this.connection, 'Not connected to server')
    invariant(type === 'exec' || type === 'sftp', 'Type should either be sftp or exec')
    if (type === 'exec') {
      const output = await this.exec('mkdir', ['-p', path])
      if (output.stdout) {
        throw new Error(output.stdout)
      }
    } else {
      invariant(!givenSftp || typeof givenSftp === 'object', 'sftp must be an object')
      const sftp = givenSftp || await this.requestSFTP()

      const makeSftpDirectory = retry =>
        Helpers.mkdirSftp(path, sftp).catch((error) => {
          if (retry && error && error.message === 'No such file') {
            return this.mkdir(Path.dirname(path), 'sftp', sftp).then(() => makeSftpDirectory(false))
          }
          throw error
        })
      try {
        await makeSftpDirectory(true)
      } finally {
        if (!givenSftp) {
          sftp.end()
        }
      }
    }
  }
  async exec(command: string, parameters: Array<string> = [], options: { cwd?: string, stdin?: string, stream?: string } = {}): Promise<string | Object> {
    invariant(this.connection, 'Not connected to server')
    invariant(typeof options === 'object' && options, 'options must be an Object')
    invariant(!options.cwd || typeof options.cwd === 'string', 'options.cwd must be a string')
    invariant(!options.stdin || typeof options.stdin === 'string', 'options.stdin must be a string')
    invariant(!options.stream || ['stdout', 'stderr', 'both'].indexOf(options.stream) !== -1, 'options.stream must be among "stdout", "stderr" and "both"')
    const output = await this.execCommand([command].concat(shellEscape(parameters)).join(' '), options)
    if (!options.stream || options.stream === 'stdout') {
      if (output.stderr) {
        throw new Error(output.stderr)
      }
      return output.stdout
    }
    if (options.stream === 'stderr') {
      return output.stderr
    }
    return output
  }
  async execCommand(givenCommand: string, options: { cwd?: string, stdin?: string } = {}): Promise<{ stdout: string, stderr: string, code: number, signal: ?string }> {
    let command = givenCommand
    const connection = this.connection
    invariant(connection, 'Not connected to server')
    invariant(typeof options === 'object' && options, 'options must be an Object')
    invariant(!options.cwd || typeof options.cwd === 'string', 'options.cwd must be a string')
    invariant(!options.stdin || typeof options.stdin === 'string', 'options.stdin must be a string')

    if (options.cwd) {
      // NOTE: Output piping cd command to hide directory non-existent errors
      command = `cd ${shellEscape([options.cwd])} 1> /dev/null 2> /dev/null; ${command}`
    }
    const output = { stdout: [], stderr: [] }
    return await new Promise(function(resolve, reject) {
      connection.exec(command, Helpers.generateCallback(function(stream) {
        stream.on('data', function(chunk) {
          output.stdout.push(chunk)
        })
        stream.stderr.on('data', function(chunk) {
          output.stderr.push(chunk)
        })
        if (options.stdin) {
          stream.write(options.stdin)
          stream.end()
        }
        stream.on('close', function(code, signal) {
          resolve({ code, signal, stdout: output.stdout.join('').trim(), stderr: output.stderr.join('').trim() })
        })
      }, reject))
    })
  }
  async getFile(localFile: string, remoteFile: string, givenSftp: ?Object = null, givenOpts: ?Object = null): Promise<void> {
    invariant(this.connection, 'Not connected to server')
    invariant(typeof localFile === 'string' && localFile, 'localFile must be a string')
    invariant(typeof remoteFile === 'string' && remoteFile, 'remoteFile must be a string')
    invariant(!givenSftp || typeof givenSftp === 'object', 'sftp must be an object')
    invariant(!givenOpts || typeof givenOpts === 'object', 'opts must be an object')

    const opts = givenOpts || {}
    const sftp = givenSftp || await this.requestSFTP()

    try {
      await new Promise(function(resolve, reject) {
        sftp.fastGet(remoteFile, localFile, opts, Helpers.generateCallback(resolve, reject))
      })
    } finally {
      if (!givenSftp) {
        sftp.end()
      }
    }
  }
  async putFile(localFile: string, remoteFile: string, givenSftp: ?Object = null, givenOpts: ?Object = null): Promise<void> {<commit_msg>:new: Add support for passing direct options to ssh2.exec<commit_after>/* @flow */

import Path from 'path'
import SSH2 from 'ssh2'
import invariant from 'assert'
import scanDirectory from 'sb-scandir'
import shellEscape from 'shell-escape'
import * as Helpers from './helpers'
import type { ConfigGiven } from './types'

class SSH {
  connection: ?SSH2;
  constructor() {
    this.connection = null
  }
  connect(givenConfig: ConfigGiven): Promise<this> {
    const connection = this.connection = new SSH2()
    return new Promise(function(resolve) {
      resolve(Helpers.normalizeConfig(givenConfig))
    }).then(config =>
      new Promise((resolve, reject) => {
        connection.on('error', reject)
        connection.on('ready', () => {
          connection.removeListener('error', reject)
          resolve(this)
        })
        connection.on('end', () => {
          if (this.connection === connection) {
            this.connection = null
          }
        })
        connection.connect(config)
      })
    )
  }
  async requestShell(): Promise<Object> {
    const connection = this.connection
    invariant(connection, 'Not connected to server')
    return await new Promise(function(resolve, reject) {
      connection.shell(Helpers.generateCallback(resolve, reject))
    })
  }
  async requestSFTP(): Promise<Object> {
    const connection = this.connection
    invariant(connection, 'Not connected to server')
    return await new Promise(function(resolve, reject) {
      connection.sftp(Helpers.generateCallback(resolve, reject))
    })
  }
  async mkdir(path: string, type: 'exec' | 'sftp' = 'sftp', givenSftp: ?Object = null): Promise<void> {
    invariant(this.connection, 'Not connected to server')
    invariant(type === 'exec' || type === 'sftp', 'Type should either be sftp or exec')
    if (type === 'exec') {
      const output = await this.exec('mkdir', ['-p', path])
      if (output.stdout) {
        throw new Error(output.stdout)
      }
    } else {
      invariant(!givenSftp || typeof givenSftp === 'object', 'sftp must be an object')
      const sftp = givenSftp || await this.requestSFTP()

      const makeSftpDirectory = retry =>
        Helpers.mkdirSftp(path, sftp).catch((error) => {
          if (retry && error && error.message === 'No such file') {
            return this.mkdir(Path.dirname(path), 'sftp', sftp).then(() => makeSftpDirectory(false))
          }
          throw error
        })
      try {
        await makeSftpDirectory(true)
      } finally {
        if (!givenSftp) {
          sftp.end()
        }
      }
    }
  }
  async exec(command: string, parameters: Array<string> = [], options: { cwd?: string, stdin?: string, stream?: string, options?: Object } = {}): Promise<string | Object> {
    invariant(this.connection, 'Not connected to server')
    invariant(typeof options === 'object' && options, 'options must be an Object')
    invariant(!options.cwd || typeof options.cwd === 'string', 'options.cwd must be a string')
    invariant(!options.stdin || typeof options.stdin === 'string', 'options.stdin must be a string')
    invariant(!options.stream || ['stdout', 'stderr', 'both'].indexOf(options.stream) !== -1, 'options.stream must be among "stdout", "stderr" and "both"')
    invariant(!options.options || typeof options.options === 'object', 'options.options must be an object')

    const output = await this.execCommand([command].concat(shellEscape(parameters)).join(' '), options)
    if (!options.stream || options.stream === 'stdout') {
      if (output.stderr) {
        throw new Error(output.stderr)
      }
      return output.stdout
    }
    if (options.stream === 'stderr') {
      return output.stderr
    }
    return output
  }
  async execCommand(givenCommand: string, options: { cwd?: string, stdin?: string, options?: Object } = {}): Promise<{ stdout: string, stderr: string, code: number, signal: ?string }> {
    let command = givenCommand
    const connection = this.connection
    invariant(connection, 'Not connected to server')
    invariant(typeof options === 'object' && options, 'options must be an Object')
    invariant(!options.cwd || typeof options.cwd === 'string', 'options.cwd must be a string')
    invariant(!options.stdin || typeof options.stdin === 'string', 'options.stdin must be a string')
    invariant(!options.options || typeof options.options === 'object', 'options.options must be an object')

    if (options.cwd) {
      // NOTE: Output piping cd command to hide directory non-existent errors
      command = `cd ${shellEscape([options.cwd])} 1> /dev/null 2> /dev/null; ${command}`
    }
    const output = { stdout: [], stderr: [] }
    return await new Promise(function(resolve, reject) {
      connection.exec(command, Helpers.generateCallback(function(stream) {
        stream.on('data', function(chunk) {
          output.stdout.push(chunk)
        })
        stream.stderr.on('data', function(chunk) {
          output.stderr.push(chunk)
        })
        if (options.stdin) {
          stream.write(options.stdin)
          stream.end()
        }
        stream.on('close', function(code, signal) {
          resolve({ code, signal, stdout: output.stdout.join('').trim(), stderr: output.stderr.join('').trim() })
        })
      }, reject), options.options || {})
    })
  }
  async getFile(localFile: string, remoteFile: string, givenSftp: ?Object = null, givenOpts: ?Object = null): Promise<void> {
    invariant(this.connection, 'Not connected to server')
    invariant(typeof localFile === 'string' && localFile, 'localFile must be a string')
    invariant(typeof remoteFile === 'string' && remoteFile, 'remoteFile must be a string')
    invariant(!givenSftp || typeof givenSftp === 'object', 'sftp must be an object')
    invariant(!givenOpts || typeof givenOpts === 'object', 'opts must be an object')

    const opts = givenOpts || {}
    const sftp = givenSftp || await this.requestSFTP()

    try {
      await new Promise(function(resolve, reject) {
        sftp.fastGet(remoteFile, localFile, opts, Helpers.generateCallback(resolve, reject))
      })
    } finally {
      if (!givenSftp) {
        sftp.end()
      }
    }
  }
  async putFile(localFile: string, remoteFile: string, givenSftp: ?Object = null, givenOpts: ?Object = null): Promise<void> {
{'fraction_non_alphanumeric': 0.144287109375, 'fraction_numerical': 0.001220703125, 'mean_word_length': 3.1086593112671346, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4572751', 'n_tokens_mistral': 3568, 'n_tokens_neox': 3490, 'n_words': 955}