Uploading files via SFTP

sftp.rb

require 'pty'
require 'expect'
require 'semantic_logger'

# Sftp that works on both Ruby and JRuby.
#
# Example Usage:
#   Sftp.upload(
#     user:            'friend',
#     password:        'secure',
#     host:            'localhost',
#     input_file_name: 'sample.txt'
#   )
#
# PTY.spawn is used so that the password can be passed via stdin into the sftp command line process.
#
# The sftp ruby gem was extremely slow and could not upload files larger than 2GB.
# This approach is much faster and does not suffer from the CPU and memory utilization issues
# associated with the sftp ruby gem.
#
# Note: The retry mechanism is only required on JRuby.
class Sftp
  include SemanticLogger::Loggable

  class SftpError < RuntimeError
  end

  class ConnectionFailure < SftpError
  end

  class SpawnFailure < SftpError
  end

  class ReadTimeout < SftpError
  end

  class SendFailure < SftpError
  end

  attr_accessor :user, :password, :host, :port,
    :input_file_name, :output_file_name, :input_file_name, :path,
    :sftp_bin, :connect_timeout_s, :read_timeout_s
  attr_reader :reader, :writer, :pid

  def self.upload(args = {})
    retry_count = 0
    sftp        = new(args)
    begin
      sftp.spawn do
        begin
          sftp.connect
          sftp.mkdir
          sftp.upload
        ensure
          sftp.disconnect
        end
      end
      true
    rescue SpawnFailure
      retry_count += 1
      if retry_count <= 5
        logger.info "Retry attempt: #{retry_count}"
        retry
      end
    end
  end

  def initialize(user:, password:, host:, port: 22, input_file_name:, output_file_name: input_file_name, path: nil, sftp_bin: 'sftp', connect_timeout_s: 30, read_timeout_s: 5)
    @user              = user
    @password          = password
    @host              = host
    @port              = port
    @input_file_name   = input_file_name
    @output_file_name  = output_file_name
    @input_file_name   = input_file_name
    @path              = path
    @sftp_bin          = sftp_bin
    @connect_timeout_s = connect_timeout_s
    @read_timeout_s    = read_timeout_s
  end

  def spawn(&block)
    PTY.spawn(sftp_bin, "-P#{port}", "#{user}@#{host}") do |reader, writer, pid|
      @reader = reader
      @writer = writer
      @pid    = pid
      SemanticLogger.tagged("sftp_pid:#{pid}", "#{user}@#{host}:#{pid}", &block)
    end
  end

  def connect
    logger.measure_info('Connect') do
      logger.trace 'Waiting for password'
      raise(SpawnFailure, "Failed to spawn #{sftp_bin}") unless reader.expect(/Password:.*/, read_timeout_s)
      writer.puts(password)
      logger.trace 'Waiting for Connected'
      raise(ConnectionFailure, 'Failed to connect') unless reader.expect(/Connected to #{host}.*/, connect_timeout_s)
    end
  end

  def mkdir
    writer.puts "mkdir -p #{path}" if path
  end

  def upload
    logger.measure_info("Upload #{input_file_name} as #{output_file_name}") do
      writer.puts 'progress'
      logger.trace 'Waiting for Progress disabled'
      raise(ReadTimeout, 'Failed to disable progress') unless reader.expect(/Progress meter disabled.*/, read_timeout_s)

      writer.puts "put #{input_file_name} #{output_file_name}"
      logger.trace 'Waiting for put echo'
      raise(ReadTimeout, 'Failed to put file') unless reader.expect(/put.*/, read_timeout_s)

      logger.trace 'Waiting for blank line'
      logger.trace "Skipping: #{reader.gets}" # Skip remainder of put line above
      while (str = reader.gets).nil?
        sleep 0.1
        # Add limit
        logger.trace 'Waiting for more blank lines'
      end
      if str && str.include?('Uploading')
        writer.puts 'progress'
        logger.trace 'Waiting for Progress enabled (download complete)'
        # No timeout on the time it takes to send the file
        reader.expect(/Progress meter enabled.*/)
      else
        raise(SendFailure, "Failed to upload #{input_file_name} as #{output_file_name}: #{str}")
      end
    end
  end

  def disconnect
    logger.measure_info('Disconnect') do
      begin
        logger.trace 'puts bye'
        writer.puts 'bye'
        logger.trace 'writer close'
        writer.close
        logger.trace 'reader close'
        reader.close
      rescue SystemCallError
      end

      logger.trace 'process wait'
      Process.wait(pid, Process::WNOHANG)
      logger.trace "Finished. Status: #{$?.inspect}"
    end
  end

end