From 984b3747a8ba8dbfdf8c26554d329b96e6543898 Mon Sep 17 00:00:00 2001
From: Matt Brictson <matt@mattbrictson.com>
Date: Fri, 22 Dec 2023 15:02:37 -0600
Subject: [PATCH] Implement SFTP transfer strategy

---
 lib/sshkit/backends/netssh.rb               | 11 ++++-
 lib/sshkit/backends/netssh/sftp_transfer.rb | 46 +++++++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)
 create mode 100644 lib/sshkit/backends/netssh/sftp_transfer.rb

diff --git a/lib/sshkit/backends/netssh.rb b/lib/sshkit/backends/netssh.rb
index 4656101c..e1c8f71c 100644
--- a/lib/sshkit/backends/netssh.rb
+++ b/lib/sshkit/backends/netssh.rb
@@ -200,10 +200,17 @@ def with_ssh(&block)
       end
 
       def with_transfer(summarizer)
-        require_relative "netssh/scp_transfer"
+        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(ScpTransfer.new(ssh, summarizer))
+          yield(transfer_class.new(ssh, summarizer))
         end
       end
     end
diff --git a/lib/sshkit/backends/netssh/sftp_transfer.rb b/lib/sshkit/backends/netssh/sftp_transfer.rb
new file mode 100644
index 00000000..20c230ba
--- /dev/null
+++ b/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