Allow SFTP to be used for upload!/download! instead of SCP

EXAMPLES.md

@@ -162,6 +162,23 @@ end
 In this case the `recursive: true` option mirrors the same options which are
 available to [`Net::{SCP,SFTP}`](http://net-ssh.github.io/net-scp/).
 
+## Set the upload/download method (SCP or SFTP).
+
+SSHKit can use SCP or SFTP for file transfers. The default is SCP, but this can be changed to SFTP per host:
+
+```ruby
+host = SSHKit::Host.new('user@example.com')
+host.transfer_method = :sftp
+```
+
+or globally:
+
+```ruby
+SSHKit::Backend::Netssh.configure do |ssh|
+  ssh.transfer_method = :sftp
+end
+```
+
 ## Setting global SSH options
 
 Setting global SSH options, these will be overwritten by options set on the

README.md

@@ -68,7 +68,8 @@ you can pass the `strip: false` option: `capture(:ls, '-l', strip: false)`
 ### Transferring files
 
 All backends also support the `upload!` and `download!` methods for transferring files.
-For the remote backend, the file is transferred with scp.
+For the remote backend, the file is transferred with scp by default, but sftp is also
+supported. See [EXAMPLES.md](EXAMPLES.md) for details.
 
 ```ruby
 on '1.example.com' do

lib/sshkit/backends/netssh.rb

@@ -2,7 +2,6 @@
 require 'strscan'
 require 'mutex_m'
 require 'net/ssh'
-require 'net/scp'
 
 module Net
   module SSH
@@ -23,10 +22,27 @@ module SSHKit
   module Backend
 
     class Netssh < Abstract
+      def self.assert_valid_transfer_method!(method)
+        return if [:scp, :sftp].include?(method)
+
+        raise ArgumentError, "#{method.inspect} is not a valid transfer method. Supported methods are :scp, :sftp."
+      end
+
       class Configuration
         attr_accessor :connection_timeout, :pty
+        attr_reader :transfer_method
         attr_writer :ssh_options
 
+        def initialize
+          self.transfer_method = :scp
+        end
+
+        def transfer_method=(method)
+          Netssh.assert_valid_transfer_method!(method)
+
+          @transfer_method = method
+        end
+
         def ssh_options
           default_options.merge(@ssh_options ||= {})
         end
@@ -64,16 +80,16 @@ def assign_defaults
       def upload!(local, remote, options = {})
         summarizer = transfer_summarizer('Uploading', options)
         remote = File.join(pwd_path, remote) unless remote.to_s.start_with?("/") || pwd_path.nil?
-        with_ssh do |ssh|
-          ssh.scp.upload!(local, remote, options, &summarizer)
+        with_transfer(summarizer) do |transfer|
+          transfer.upload!(local, remote, options)
         end
       end
 
       def download!(remote, local=nil, options = {})
         summarizer = transfer_summarizer('Downloading', options)
         remote = File.join(pwd_path, remote) unless remote.to_s.start_with?("/") || pwd_path.nil?
-        with_ssh do |ssh|
-          ssh.scp.download!(remote, local, options, &summarizer)
+        with_transfer(summarizer) do |transfer|
+          transfer.download!(remote, local, options)
         end
       end
 
@@ -105,7 +121,7 @@ def transfer_summarizer(action, options = {})
         last_percentage = nil
         proc do |_ch, name, transferred, total|
           percentage = (transferred.to_f * 100 / total.to_f)
-          unless percentage.nan?
+          unless percentage.nan? || percentage.infinite?
             message = "#{action} #{name} #{percentage.round(2)}%"
             percentage_r = (percentage / log_percent).truncate * log_percent
             if percentage_r > 0 && (last_name != name || last_percentage != percentage_r)
@@ -183,6 +199,20 @@ def with_ssh(&block)
         )
       end
 
+      def with_transfer(summarizer)
+        transfer_method = host.transfer_method || self.class.config.transfer_method
+        transfer_class = if transfer_method == :sftp
+                           require_relative "netssh/sftp_transfer"
+                           SftpTransfer
+                         else
+                           require_relative "netssh/scp_transfer"
+                           ScpTransfer
+                         end
+
+        with_ssh do |ssh|
+          yield(transfer_class.new(ssh, summarizer))
+        end
+      end
     end
   end
 

lib/sshkit/backends/netssh/scp_transfer.rb

@@ -0,0 +1,26 @@
+require "net/scp"
+
+module SSHKit
+  module Backend
+    class Netssh < Abstract
+      class ScpTransfer
+        def initialize(ssh, summarizer)
+          @ssh = ssh
+          @summarizer = summarizer
+        end
+
+        def upload!(local, remote, options)
+          ssh.scp.upload!(local, remote, options, &summarizer)
+        end
+
+        def download!(remote, local, options)
+          ssh.scp.download!(remote, local, options, &summarizer)
+        end
+
+        private
+
+        attr_reader :ssh, :summarizer
+      end
+    end
+  end
+end

lib/sshkit/backends/netssh/sftp_transfer.rb

@@ -0,0 +1,46 @@
+require "net/sftp"
+
+module SSHKit
+  module Backend
+    class Netssh < Abstract
+      class SftpTransfer
+        def initialize(ssh, summarizer)
+          @ssh = ssh
+          @summarizer = summarizer
+        end
+
+        def upload!(local, remote, options)
+          options = { progress: self }.merge(options || {})
+          ssh.sftp.connect!
+          ssh.sftp.upload!(local, remote, options)
+        ensure
+          ssh.sftp.close_channel
+        end
+
+        def download!(remote, local, options)
+          options = { progress: self }.merge(options || {})
+          destination = local ? local : StringIO.new.tap { |io| io.set_encoding('BINARY') }
+
+          ssh.sftp.connect!
+          ssh.sftp.download!(remote, destination, options)
+          local ? true : destination.string
+        ensure
+          ssh.sftp.close_channel
+        end
+
+        def on_get(download, entry, offset, data)
+          entry.size ||= download.sftp.file.open(entry.remote, &:size)
+          summarizer.call(nil, entry.remote, offset + data.bytesize, entry.size)
+        end
+
+        def on_put(_upload, file, offset, data)
+          summarizer.call(nil, file.local, offset + data.bytesize, file.size)
+        end
+
+        private
+
+        attr_reader :ssh, :summarizer
+      end
+    end
+  end
+end

lib/sshkit/host.rb

@@ -7,6 +7,7 @@ module SSHKit
   class Host
 
     attr_accessor :password, :hostname, :port, :user, :ssh_options
+    attr_reader :transfer_method
 
     def key=(new_key)
       @keys = [new_key]
@@ -41,6 +42,12 @@ def initialize(host_string_or_options_hash)
       end
     end
 
+    def transfer_method=(method)
+      Backend::Netssh.assert_valid_transfer_method!(method) unless method.nil?
+
+      @transfer_method = method
+    end
+
     def local?
       @local
     end

sshkit.gemspec

@@ -23,6 +23,7 @@ Gem::Specification.new do |gem|
   gem.add_runtime_dependency('mutex_m')
   gem.add_runtime_dependency('net-ssh',  '>= 2.8.0')
   gem.add_runtime_dependency('net-scp',  '>= 1.1.2')
+  gem.add_runtime_dependency('net-sftp', '>= 2.1.2')
 
   gem.add_development_dependency('danger')
   gem.add_development_dependency('minitest', '>= 5.0.0')

test/functional/backends/netssh_transfer_tests.rb

@@ -0,0 +1,83 @@
+require 'securerandom'
+
+module SSHKit
+  module Backend
+    module NetsshTransferTests
+      def setup
+        super
+        @output = String.new
+        SSHKit.config.output_verbosity = :debug
+        SSHKit.config.output = SSHKit::Formatter::SimpleText.new(@output)
+      end
+
+      def a_host
+        VagrantWrapper.hosts['one']
+      end
+
+      def test_upload_and_then_capture_file_contents
+        actual_file_contents = ""
+        file_name = File.join("/tmp", SecureRandom.uuid)
+        File.open file_name, 'w+' do |f|
+          f.write "Some Content\nWith a newline and trailing spaces    \n "
+        end
+        Netssh.new(a_host) do
+          upload!(file_name, file_name)
+          actual_file_contents = capture(:cat, file_name, strip: false)
+        end.run
+        assert_equal "Some Content\nWith a newline and trailing spaces    \n ", actual_file_contents
+      end
+
+      def test_upload_within
+        file_name = SecureRandom.uuid
+        file_contents = "Some Content"
+        dir_name = SecureRandom.uuid
+        actual_file_contents = ""
+        Netssh.new(a_host) do |_host|
+          within("/tmp") do
+            execute :mkdir, "-p", dir_name
+            within(dir_name) do
+              upload!(StringIO.new(file_contents), file_name)
+            end
+          end
+          actual_file_contents = capture(:cat, "/tmp/#{dir_name}/#{file_name}", strip: false)
+        end.run
+        assert_equal file_contents, actual_file_contents
+      end
+
+      def test_upload_string_io
+        file_contents = ""
+        Netssh.new(a_host) do |_host|
+          file_name = File.join("/tmp", SecureRandom.uuid)
+          upload!(StringIO.new('example_io'), file_name)
+          file_contents = download!(file_name)
+        end.run
+        assert_equal "example_io", file_contents
+      end
+
+      def test_upload_large_file
+        size      = 25
+        fills     = SecureRandom.random_bytes(1024*1024)
+        file_name = "/tmp/file-#{size}.txt"
+        File.open(file_name, 'wb') do |f|
+          (size).times {f.write(fills) }
+        end
+        file_contents = ""
+        Netssh.new(a_host) do
+          upload!(file_name, file_name)
+          file_contents = download!(file_name)
+        end.run
+        assert_equal File.open(file_name, 'rb').read, file_contents
+      end
+
+      def test_upload_via_pathname
+        file_contents = ""
+        Netssh.new(a_host) do |_host|
+          file_name = Pathname.new(File.join("/tmp", SecureRandom.uuid))
+          upload!(StringIO.new('example_io'), file_name)
+          file_contents = download!(file_name)
+        end.run
+        assert_equal "example_io", file_contents
+      end
+    end
+  end
+end

test/functional/backends/test_netssh.rb

@@ -1,5 +1,4 @@
 require 'helper'
-require 'securerandom'
 require 'benchmark'
 
 module SSHKit
@@ -136,71 +135,6 @@ def test_test_does_not_raise_on_non_zero_exit_status
         end.run
       end
 
-      def test_upload_and_then_capture_file_contents
-        actual_file_contents = ""
-        file_name = File.join("/tmp", SecureRandom.uuid)
-        File.open file_name, 'w+' do |f|
-          f.write "Some Content\nWith a newline and trailing spaces    \n "
-        end
-        Netssh.new(a_host) do
-          upload!(file_name, file_name)
-          actual_file_contents = capture(:cat, file_name, strip: false)
-        end.run
-        assert_equal "Some Content\nWith a newline and trailing spaces    \n ", actual_file_contents
-      end
-
-      def test_upload_within
-        file_name = SecureRandom.uuid
-        file_contents = "Some Content"
-        dir_name = SecureRandom.uuid
-        actual_file_contents = ""
-        Netssh.new(a_host) do |_host|
-          within("/tmp") do
-            execute :mkdir, "-p", dir_name
-            within(dir_name) do
-              upload!(StringIO.new(file_contents), file_name)
-            end
-          end
-          actual_file_contents = capture(:cat, "/tmp/#{dir_name}/#{file_name}", strip: false)
-        end.run
-        assert_equal file_contents, actual_file_contents
-      end
-
-      def test_upload_string_io
-        file_contents = ""
-        Netssh.new(a_host) do |_host|
-          file_name = File.join("/tmp", SecureRandom.uuid)
-          upload!(StringIO.new('example_io'), file_name)
-          file_contents = download!(file_name)
-        end.run
-        assert_equal "example_io", file_contents
-      end
-
-      def test_upload_large_file
-        size      = 25
-        fills     = SecureRandom.random_bytes(1024*1024)
-        file_name = "/tmp/file-#{size}.txt"
-        File.open(file_name, 'wb') do |f|
-          (size).times {f.write(fills) }
-        end
-        file_contents = ""
-        Netssh.new(a_host) do
-          upload!(file_name, file_name)
-          file_contents = download!(file_name)
-        end.run
-        assert_equal File.open(file_name, 'rb').read, file_contents
-      end
-
-      def test_upload_via_pathname
-        file_contents = ""
-        Netssh.new(a_host) do |_host|
-          file_name = Pathname.new(File.join("/tmp", SecureRandom.uuid))
-          upload!(StringIO.new('example_io'), file_name)
-          file_contents = download!(file_name)
-        end.run
-        assert_equal "example_io", file_contents
-      end
-
       def test_interaction_handler
         captured_command_result = nil
         Netssh.new(a_host) do