net-sftp 1.1.1 → 2.0.0

data/lib/net/sftp/operations/file.rb

@@ -0,0 +1,176 @@
+require 'net/ssh/loggable'
+require 'net/sftp/operations/file'
+
+module Net; module SFTP; module Operations
+
+  # A wrapper around an SFTP file handle, that exposes an IO-like interface
+  # for interacting with the remote file. All operations are synchronous
+  # (blocking), making this a very convenient way to deal with remote files.
+  #
+  # A wrapper is usually created via the Net::SFTP::Session#file factory:
+  #
+  #   file = sftp.file.open("/path/to/remote")
+  #   puts file.gets
+  #   file.close
+  class File
+    # A reference to the Net::SFTP::Session instance that drives this wrapper
+    attr_reader :sftp
+
+    # The SFTP file handle object that this object wraps
+    attr_reader :handle
+
+    # The current position within the remote file
+    attr_reader :pos
+
+    # Creates a new wrapper that encapsulates the given +handle+ (such as
+    # would be returned by Net::SFTP::Session#open!). The +sftp+ parameter
+    # must be the same Net::SFTP::Session instance that opened the file.
+    def initialize(sftp, handle)
+      @sftp     = sftp
+      @handle   = handle
+      @pos      = 0
+      @real_pos = 0
+      @real_eof = false
+      @buffer   = ""
+    end
+
+    # Repositions the file pointer to the given offset (relative to the
+    # start of the file). This will also reset the EOF flag.
+    def pos=(offset)
+      @real_pos = @pos = offset
+      @buffer = ""
+      @real_eof = false
+    end
+
+    # Closes the underlying file and sets the handle to +nil+. Subsequent
+    # operations on this object will fail.
+    def close
+      sftp.close!(handle)
+      @handle = nil
+    end
+
+    # Returns true if the end of the file has been encountered by a previous
+    # read. Setting the current file position via #pos= will reset this
+    # flag (useful if the file's contents have changed since the EOF was
+    # encountered).
+    def eof?
+      @real_eof && @buffer.empty?
+    end
+
+    # Reads up to +n+ bytes of data from the stream. Fewer bytes will be
+    # returned if EOF is encountered before the requested number of bytes
+    # could be read. Without an argument (or with a nil argument) all data
+    # to the end of the file will be read and returned.
+    #
+    # This will advance the file pointer (#pos).
+    def read(n=nil)
+      loop do
+        break if n && @buffer.length >= n
+        break unless fill
+      end
+
+      if n
+        result, @buffer = @buffer[0,n], (@buffer[n..-1] || "")
+      else
+        result, @buffer = @buffer, ""
+      end
+
+      @pos += result.length
+      return result
+    end
+
+    # Reads up to the next instance of +sep_string+ in the stream, and
+    # returns the bytes read (including +sep_string+). If +sep_string+ is
+    # omitted, it defaults to +$/+. If EOF is encountered before any data
+    # could be read, #gets will return +nil+.
+    def gets(sep_string=$/)
+      delim = if sep_string.length == 0
+        "#{$/}#{$/}"
+      else
+        sep_string
+      end
+
+      loop do
+        at = @buffer.index(delim)
+        if at
+          offset = at + delim.length
+          @pos += offset
+          line, @buffer = @buffer[0,offset], @buffer[offset..-1]
+          return line
+        elsif !fill
+          return nil if @buffer.empty?
+          @pos += @buffer.length
+          line, @buffer = @buffer, ""
+          return line
+        end
+      end
+    end
+
+    # Same as #gets, but raises EOFError if EOF is encountered before any
+    # data could be read.
+    def readline(sep_string=$/)
+      line = gets(sep_string)
+      raise EOFError if line.nil?
+      return line
+    end
+
+    # Writes the given data to the stream, incrementing the file position and
+    # returning the number of bytes written.
+    def write(data)
+      data = data.to_s
+      sftp.write!(handle, @real_pos, data)
+      @real_pos += data.length
+      @pos = @real_pos
+      data.length
+    end
+
+    # Writes each argument to the stream. If +$\+ is set, it will be written
+    # after all arguments have been written.
+    def print(*items)
+      items.each { |item| write(item) }
+      write($\) if $\
+      nil
+    end
+
+    # Writes each argument to the stream, appending a newline to any item
+    # that does not already end in a newline. Array arguments are flattened.
+    def puts(*items)
+      items.each do |item|
+        if Array === item
+          puts(*item)
+        else
+          write(item)
+          write("\n") unless item[-1] == ?\n
+        end
+      end
+      nil
+    end
+
+    # Performs an fstat operation on the handle and returns the attribute
+    # object (Net::SFTP::Protocol::V01::Attributes, Net::SFTP::Protool::V04::Attributes,
+    # or Net::SFTP::Protocol::V06::Attributes, depending on the SFTP protocol
+    # version in use).
+    def stat
+      sftp.fstat!(handle)
+    end
+
+    private
+
+      # Fills the buffer. Returns +true+ if it succeeded, and +false+ if
+      # EOF was encountered before any data was read.
+      def fill
+        data = sftp.read!(handle, @real_pos, 8192)
+
+        if data.nil?
+          @real_eof = true
+          return false
+        else
+          @real_pos += data.length
+          @buffer << data
+        end
+
+        !@real_eof
+      end
+  end
+
+end; end; end

data/lib/net/sftp/operations/file_factory.rb

@@ -0,0 +1,60 @@
+require 'net/ssh/loggable'
+require 'net/sftp/operations/file'
+
+module Net; module SFTP; module Operations
+
+  # A factory class for opening files and returning Operations::File instances
+  # that wrap the SFTP handles that represent them. This is a convenience
+  # class for use when working with files synchronously. Rather than relying
+  # on the programmer to provide callbacks that define a state machine that
+  # describes the behavior of the program, this class (and Operations::File)
+  # provide an interface where calls will block until they return, mimicking
+  # the IO class' interface.
+  class FileFactory
+    # The SFTP session object that drives this file factory.
+    attr_reader :sftp
+
+    # Create a new instance on top of the given SFTP session instance.
+    def initialize(sftp)
+      @sftp = sftp
+    end
+
+    # :call-seq:
+    #   open(name, flags="r", mode=nil) -> file
+    #   open(name, flags="r", mode=nil) { |file| ... }
+    #
+    # Attempt to open a file on the remote server. The +flags+ parameter
+    # accepts the same values as the standard Ruby ::File#open method. The
+    # +mode+ parameter must be an integer describing the permissions to use
+    # if a new file is being created.
+    #
+    # If a block is given, the new Operations::File instance will be yielded
+    # to it, and closed automatically when the block terminates. Otherwise
+    # the object will be returned, and it is the caller's responsibility to
+    # close the file.
+    #
+    #   sftp.file.open("/tmp/names.txt", "w") do |f|
+    #     # ...
+    #   end
+    def open(name, flags="r", mode=nil, &block)
+      handle = sftp.open!(name, flags, :permissions => mode)
+      file = Operations::File.new(sftp, handle)
+
+      if block_given?
+        begin
+          yield file
+        ensure
+          file.close
+        end
+      else
+        return file
+      end
+    end
+
+    # Returns +true+ if the argument refers to a directory on the remote host.
+    def directory?(path)
+      sftp.lstat!(path).directory?
+    end
+  end
+
+end; end; end

data/lib/net/sftp/operations/upload.rb

@@ -0,0 +1,387 @@
+require 'net/ssh/loggable'
+
+module Net; module SFTP; module Operations
+
+  # A general purpose uploader module for Net::SFTP. It can upload IO objects,
+  # files, and even entire directory trees via SFTP, and provides a flexible
+  # progress reporting mechanism.
+  #
+  # To upload a single file to the remote server, simply specify both the
+  # local and remote paths:
+  #
+  #   uploader = sftp.upload("/path/to/local.txt", "/path/to/remote.txt")
+  #
+  # By default, this operates asynchronously, so if you want to block until
+  # the upload finishes, you can use the 'bang' variant:
+  #
+  #   sftp.upload!("/path/to/local.txt", "/path/to/remote.txt")
+  #
+  # Or, if you have multiple uploads that you want to run in parallel, you can
+  # employ the #wait method of the returned object:
+  #
+  #   uploads = %w(file1 file2 file3).map { |f| sftp.upload(f, "remote/#{f}") }
+  #   uploads.each { |u| u.wait }
+  #
+  # To upload an entire directory tree, recursively, simply pass the directory
+  # path as the first parameter:
+  #
+  #   sftp.upload!("/path/to/directory", "/path/to/remote")
+  #
+  # This will upload "/path/to/directory", it's contents, it's subdirectories,
+  # and their contents, recursively, to "/path/to/remote" on the remote server.
+  #
+  # If you want to send data to a file on the remote server, but the data is
+  # in memory, you can pass an IO object and upload it's contents:
+  #
+  #   require 'stringio'
+  #   io = StringIO.new(data)
+  #   sftp.upload!(io, "/path/to/remote")
+  #
+  # The following options are supported:
+  #
+  # * <tt>:progress</tt> - either a block or an object to act as a progress
+  #   callback. See the discussion of "progress monitoring" below.
+  # * <tt>:requests</tt> - the number of pending SFTP requests to allow at
+  #   any given time. When uploading an entire directory tree recursively,
+  #   this will default to 16, otherwise it will default to 2. Setting this
+  #   higher might improve throughput. Reducing it will reduce throughput.
+  # * <tt>:read_size</tt> - the maximum number of bytes to read at a time
+  #   from the source. Increasing this value might improve throughput. It
+  #   defaults to 32,000 bytes.
+  # * <tt>:name</tt> - the filename to report to the progress monitor when
+  #   an IO object is given as +local+. This defaults to "<memory>".
+  #
+  # == Progress Monitoring
+  #
+  # Sometimes it is desirable to track the progress of an upload. There are
+  # two ways to do this: either using a callback block, or a special custom
+  # object.
+  #
+  # Using a block it's pretty straightforward:
+  #
+  #   sftp.upload!("local", "remote") do |event, uploader, *args|
+  #     case event
+  #     when :open then
+  #       # args[0] : file metadata
+  #       puts "starting upload: #{args[0].local} -> #{args[0].remote} (#{args[0].size} bytes}"
+  #     when :put then
+  #       # args[0] : file metadata
+  #       # args[1] : byte offset in remote file
+  #       # args[2] : data being written (as string)
+  #       puts "writing #{args[2].length} bytes to #{args[0].remote} starting at #{args[1]}"
+  #     when :close then
+  #       # args[0] : file metadata
+  #       puts "finished with #{args[0].remote}"
+  #     when :mkdir then
+  #       # args[0] : remote path name
+  #       puts "creating directory #{args[0]}"
+  #     when :finish then
+  #       puts "all done!"
+  #   end
+  #
+  # However, for more complex implementations (e.g., GUI interfaces and such)
+  # a block can become cumbersome. In those cases, you can create custom
+  # handler objects that respond to certain methods, and then pass your handler
+  # to the uploader:
+  #
+  #   class CustomHandler
+  #     def on_open(uploader, file)
+  #       puts "starting upload: #{file.local} -> #{file.remote} (#{file.size} bytes)"
+  #     end
+  #
+  #     def on_put(uploader, file, offset, data)
+  #       puts "writing #{data.length} bytes to #{file.remote} starting at #{offset}"
+  #     end
+  #
+  #     def on_close(uploader, file)
+  #       puts "finished with #{file.remote}"
+  #     end
+  #
+  #     def on_mkdir(uploader, path)
+  #       puts "creating directory #{path}"
+  #     end
+  #
+  #     def on_finish(uploader)
+  #       puts "all done!"
+  #     end
+  #   end
+  #
+  #   sftp.upload!("local", "remote", :progress => CustomHandler.new)
+  #
+  # If you omit any of those methods, the progress updates for those missing
+  # events will be ignored. You can create a catchall method named "call" for
+  # those, instead.
+  class Upload
+    include Net::SSH::Loggable
+
+    # The source of the upload (on the local server)
+    attr_reader :local
+
+    # The destination of the upload (on the remote server)
+    attr_reader :remote
+
+    # The hash of options that were given when the object was instantiated
+    attr_reader :options
+
+    # The SFTP session object used by this upload instance
+    attr_reader :sftp
+
+    # The properties hash for this object
+    attr_reader :properties
+
+    # Instantiates a new uploader process on top of the given SFTP session.
+    # +local+ is either an IO object containing data to upload, or a string
+    # identifying a file or directory on the local host. +remote+ is a string
+    # identifying the location on the remote host that the upload should
+    # target.
+    #
+    # This will return immediately, and requires that the SSH event loop be
+    # run in order to effect the upload. (See #wait.)
+    def initialize(sftp, local, remote, options={}, &progress) #:nodoc:
+      @sftp = sftp
+      @local = local
+      @remote = remote
+      @progress = progress || options[:progress]
+      @options = options
+      @properties = options[:properties] || {}
+      @active = 0
+
+      self.logger = sftp.logger
+
+      @uploads = []
+      @recursive = local.respond_to?(:read) ? false : ::File.directory?(local)
+
+      if recursive?
+        @stack = [entries_for(local)]
+        @local_cwd = local
+        @remote_cwd = remote
+
+        @active += 1
+        sftp.mkdir(remote) do |response|
+          @active -= 1
+          raise StatusException.new(response, "mkdir `#{remote}'") unless response.ok?
+          (options[:requests] || RECURSIVE_READERS).to_i.times do
+            break unless process_next_entry
+          end
+        end
+      else
+        raise ArgumentError, "expected a file to upload" unless local.respond_to?(:read) || ::File.exists?(local)
+        @stack = [[local]]
+        process_next_entry
+      end
+    end
+
+    # Returns true if a directory tree is being uploaded, and false if only a
+    # single file is being uploaded.
+    def recursive?
+      @recursive
+    end
+
+    # Returns true if the uploader is currently running. When this is false,
+    # the uploader has finished processing.
+    def active?
+      @active > 0 || @stack.any?
+    end
+
+    # Forces the transfer to stop.
+    def abort!
+      @active = 0
+      @stack.clear
+      @uploads.clear
+    end
+
+    # Blocks until the upload has completed.
+    def wait
+      sftp.loop { active? }
+      self
+    end
+
+    # Returns the property with the given name. This allows Upload instances
+    # to store their own state when used as part of a state machine.
+    def [](name)
+      @properties[name.to_sym]
+    end
+
+    # Sets the given property to the given name. This allows Upload instances
+    # to store their own state when used as part of a state machine.
+    def []=(name, value)
+      @properties[name.to_sym] = value
+    end
+
+    private
+
+      #--
+      # "ruby -w" hates private attributes, so we have to do this longhand.
+      #++
+
+      # The progress handler for this instance. Possibly nil.
+      def progress; @progress; end
+
+      # A simple struct for recording metadata about the file currently being
+      # uploaded.
+      LiveFile = Struct.new(:local, :remote, :io, :size, :handle)
+
+      # The default # of bytes to read from disk at a time.
+      DEFAULT_READ_SIZE   = 32_000
+
+      # The number of readers to use when uploading a single file.
+      SINGLE_FILE_READERS = 2
+
+      # The number of readers to use when uploading a directory.
+      RECURSIVE_READERS   = 16
+
+      # Examines the stack and determines what action to take. This is the
+      # starting point of the state machine.
+      def process_next_entry
+        if @stack.empty?
+          if @uploads.any?
+            write_next_chunk(@uploads.first)
+          elsif !active?
+            update_progress(:finish)
+          end
+          return false
+        elsif @stack.last.empty?
+          @stack.pop
+          @local_cwd = ::File.dirname(@local_cwd)
+          @remote_cwd = ::File.dirname(@remote_cwd)
+          process_next_entry
+        elsif recursive?
+          entry = @stack.last.shift
+          lpath = ::File.join(@local_cwd, entry)
+          rpath = ::File.join(@remote_cwd, entry)
+
+          if ::File.directory?(lpath)
+            @stack.push(entries_for(lpath))
+            @local_cwd = lpath
+            @remote_cwd = rpath
+
+            @active += 1
+            update_progress(:mkdir, rpath)
+            request = sftp.mkdir(rpath, &method(:on_mkdir))
+            request[:dir] = rpath
+          else
+            open_file(lpath, rpath)
+          end
+        else
+          open_file(@stack.pop.first, remote)
+        end
+        return true
+      end
+
+      # Prepares to send +local+ to +remote+.
+      def open_file(local, remote)
+        @active += 1
+
+        if local.respond_to?(:read)
+          file = local
+          name = options[:name] || "<memory>"
+        else
+          file = ::File.open(local)
+          name = local
+        end
+
+        if file.respond_to?(:stat)
+          size = file.stat.size
+        else
+          size = file.size
+        end
+
+        metafile = LiveFile.new(name, remote, file, size)
+        update_progress(:open, metafile)
+
+        request = sftp.open(remote, "w", &method(:on_open))
+        request[:file] = metafile
+      end
+
+      # Called when a +mkdir+ request finishes, successfully or otherwise.
+      # If the request failed, this will raise a StatusException, otherwise
+      # it will call #process_next_entry to continue the state machine.
+      def on_mkdir(response)
+        @active -= 1
+        dir = response.request[:dir]
+        raise StatusException.new(response, "mkdir #{dir}") unless response.ok?
+
+        process_next_entry
+      end
+
+      # Called when an +open+ request finishes. Raises StatusException if the
+      # open failed, otherwise it calls #write_next_chunk to begin sending
+      # data to the remote server.
+      def on_open(response)
+        @active -= 1
+        file = response.request[:file]
+        raise StatusException.new(response, "open #{file.remote}") unless response.ok?
+
+        file.handle = response[:handle]
+
+        @uploads << file
+        write_next_chunk(file)
+
+        if !recursive?
+          (options[:requests] || SINGLE_FILE_READERS).to_i.times { write_next_chunk(file) }
+        end
+      end
+
+      # Called when a +write+ request finishes. Raises StatusException if the
+      # write failed, otherwise it calls #write_next_chunk to continue the
+      # write.
+      def on_write(response)
+        @active -= 1
+        file = response.request[:file]
+        raise StatusException.new(response, "write #{file.remote}") unless response.ok?
+        write_next_chunk(file)
+      end
+
+      # Called when a +close+ request finishes. Raises a StatusException if the
+      # close failed, otherwise it calls #process_next_entry to continue the
+      # state machine.
+      def on_close(response)
+        @active -= 1
+        file = response.request[:file]
+        raise StatusException.new(response, "close #{file.remote}") unless response.ok?
+        process_next_entry
+      end
+
+      # Attempts to send the next chunk from the given file (where +file+ is
+      # a LiveFile instance).
+      def write_next_chunk(file)
+        if file.io.nil?
+          process_next_entry
+        else
+          @active += 1
+          offset = file.io.pos
+          data = file.io.read(options[:read_size] || DEFAULT_READ_SIZE)
+          if data.nil?
+            update_progress(:close, file)
+            request = sftp.close(file.handle, &method(:on_close))
+            request[:file] = file
+            file.io.close
+            file.io = nil
+            @uploads.delete(file)
+          else
+            update_progress(:put, file, offset, data)
+            request = sftp.write(file.handle, offset, data, &method(:on_write))
+            request[:file] = file
+          end
+        end
+      end
+
+      # Returns all directory entries for the given path, removing the '.'
+      # and '..' relative paths.
+      def entries_for(local)
+        ::Dir.entries(local).reject { |v| %w(. ..).include?(v) }
+      end
+
+      # Attempts to notify the progress monitor (if one was given) about
+      # progress made for the given event.
+      def update_progress(event, *args)
+        on = "on_#{event}"
+        if progress.respond_to?(on)
+          progress.send(on, self, *args)
+        elsif progress.respond_to?(:call)
+          progress.call(event, self, *args)
+        end
+      end
+  end
+
+end; end; end