From 5c97832a6ae72a946fc2bb50aa636c7543854615 Mon Sep 17 00:00:00 2001
From: Michal Bicz <mbicz@twilio.com>
Date: Sun, 15 Feb 2015 16:20:45 -0800
Subject: [PATCH] add KITCHEN_SYNC_IDENTITY to specify which key-pair to use
 with more loggin

---
 README.md                   | 12 +++++++----
 lib/kitchen-sync/rsync.rb   | 41 ++++++++++++++++++++++++-------------
 lib/kitchen-sync/version.rb |  2 +-
 3 files changed, 36 insertions(+), 19 deletions(-)

diff --git a/README.md b/README.md
index 6ce1189..fe5c86f 100644
--- a/README.md
+++ b/README.md
@@ -37,10 +37,14 @@ used directly.
 ### Rsync
 
 The rsync mode is based on the work done by [Mikhail Bautin](https://github.com/test-kitchen/test-kitchen/pull/359).
-This is the fastest mode, but it does have a few downsides. The biggest is that
-you must be using `ssh-agent` and have an identity loaded for it to use. It also
-requires that rsync be available on the remote side. Consider this implementation
-more experimental than the others at this time.
+This is the fastest mode, but it does have a few downsides.
+There are 2 options to provide authenticatoin for Rsync.
+If no KITCHEN_SYNC_IDENTITY env variable is set. `.ssh/config` is consulted
+for the host and it's `IdentityFile` is used for authenticaton.
+`KITCHEN_SYNC_IDENTITY=0` will inform kitchen-sync to use 1st identity
+from ssh-agent. You can see all identities configure by issuing `ssh-add -l`.
+It also requires that rsync be available on the remote side.
+Consider this implementation more experimental than the others at this time.
 
 License
 -------
diff --git a/lib/kitchen-sync/rsync.rb b/lib/kitchen-sync/rsync.rb
index 3df3b08..bea004e 100644
--- a/lib/kitchen-sync/rsync.rb
+++ b/lib/kitchen-sync/rsync.rb
@@ -25,7 +25,14 @@ def upload(local, remote, recursive=true)
       upload_done = false
       if !@rsync_failed && recursive && File.exists?('/usr/bin/rsync')
         ssh_command = "ssh #{ssh_args.join(' ')}"
-        copy_identity
+        ssh_identity ||= ENV["KITCHEN_SYNC_IDENTITY"].to_i
+        if ssh_identity
+          ssh_identity_key = `ssh-add -l`.split("\n")[ssh_identity]
+          @logger.info("[sync:rsync] Using SSH identity ##{ssh_identity} key => #{ssh_identity_key}.")
+          copy_identity(ssh_identity)
+        else
+          @logger.info("[sync:rsync] Using default ssh_config IdentityFile configuration.")
+        end
         rsync_cmd = "/usr/bin/rsync -e '#{ssh_command}' -az #{local} #{@session.options[:user]}@#{@session.host}:#{remote}"
         @logger.info("[sync:rsync] Running rsync command: #{rsync_cmd}")
         if system(rsync_cmd)
@@ -41,20 +48,26 @@ def upload(local, remote, recursive=true)
     end
 
     # Copy your SSH identity, creating a new one if needed
-    def copy_identity
+    def copy_identity(ssh_identity)
       return if @copied_identity
-      key = Net::SSH::Authentication::Agent.connect.identities.first
-      enc_key = Base64.encode64(key.to_blob).gsub("\n", '')
-      identitiy = "ssh-rsa #{enc_key} #{key.comment}"
-      @session.exec! <<-EOT
-        test -e ~/.ssh || mkdir ~/.ssh
-        test -e ~/.ssh/authorized_keys || touch ~/.ssh/authorized_keys
-        if ! grep -q "#{identitiy}" ~/.ssh/authorized_keys ; then
-          chmod go-w ~ ~/.ssh ~/.ssh/authorized_keys ; \
-          echo "#{identitiy}" >> ~/.ssh/authorized_keys
-        fi
-      EOT
-      @copied_identity = true
+      key = Net::SSH::Authentication::Agent.connect.identities[ssh_identity]
+      if key
+        enc_key = Base64.encode64(key.to_blob).gsub("\n", '')
+        identitiy = "ssh-rsa #{enc_key} #{key.comment}"
+        @session.exec! <<-EOT
+          test -e ~/.ssh || mkdir ~/.ssh
+          test -e ~/.ssh/authorized_keys || touch ~/.ssh/authorized_keys
+          if ! grep -q "#{identitiy}" ~/.ssh/authorized_keys ; then
+            chmod go-w ~ ~/.ssh ~/.ssh/authorized_keys ; \
+            echo "#{identitiy}" >> ~/.ssh/authorized_keys
+          fi
+        EOT
+        @copied_identity = true
+        @logger.info("[sync:rsync] Successfully copied SSH identity ##{ssh_identity}.")
+      else
+        @logger.warn("[sync:rsync] Failed to copy SSH identity ##{ssh_identity}.\
+          Falling back to default ssh_config IdentityFile configuration")
+      end
     end
 
     def ssh_args
diff --git a/lib/kitchen-sync/version.rb b/lib/kitchen-sync/version.rb
index 1aa1ddb..e408e37 100644
--- a/lib/kitchen-sync/version.rb
+++ b/lib/kitchen-sync/version.rb
@@ -17,5 +17,5 @@
 #
 
 class KitchenSync
-  VERSION = '1.0.1'
+  VERSION = '1.0.2'
 end