From 0d3702d49e6f1466334bee0b009f7d4280493163 Mon Sep 17 00:00:00 2001
From: Alex Pop <apop@chef.io>
Date: Mon, 21 Nov 2016 13:59:21 +0000
Subject: [PATCH] try to use ssh agent if no password or key files have been
 specified Signed-off-by: Alex Pop <apop@chef.io>

---
 Gemfile                          |  1 +
 README.md                        |  7 +++++++
 lib/train/transports/ssh.rb      | 28 ++++++++++++++++++++++------
 test/unit/transports/ssh_test.rb | 26 +++++++++++++++++++++++---
 4 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/Gemfile b/Gemfile
index c6074613..ed43bede 100644
--- a/Gemfile
+++ b/Gemfile
@@ -30,6 +30,7 @@ end
 
 group :tools do
   gem 'pry', '~> 0.10'
+  gem 'rb-readline'
   gem 'license_finder'
   gem 'github_changelog_generator', '~> 1'
 end
diff --git a/README.md b/README.md
index 84ff8cc8..8cda7716 100644
--- a/README.md
+++ b/README.md
@@ -39,6 +39,13 @@ train = Train.create('ssh',
   host: '1.2.3.4', port: 22, user: 'root', key_files: '/vagrant')
 ```
 
+If you don't specify the `key_files` and `password` options, SSH agent authentication will be attempted. For example:
+
+```ruby
+require 'train'
+train = Train.create('ssh', host: '1.2.3.4', port: 22, user: 'root')
+```
+
 **WinRM**
 
 ```ruby
diff --git a/lib/train/transports/ssh.rb b/lib/train/transports/ssh.rb
index 96debf98..054839b1 100644
--- a/lib/train/transports/ssh.rb
+++ b/lib/train/transports/ssh.rb
@@ -82,12 +82,6 @@ def validate_options(options)
       super(options)
 
       key_files = Array(options[:key_files])
-      if key_files.empty? and options[:password].nil?
-        fail Train::ClientError,
-             'You must configure at least one authentication method for SSH:'\
-             ' Password or key.'
-      end
-
       options[:auth_methods] ||= ['none']
 
       unless key_files.empty?
@@ -100,6 +94,17 @@ def validate_options(options)
         options[:auth_methods].push('password', 'keyboard-interactive')
       end
 
+      if options[:auth_methods] == ['none']
+        if ssh_known_identities.empty?
+          fail Train::ClientError,
+               'You must configure at least one authentication method for SSH:'\
+               ' Agent, Key or Password.'
+        else
+          logger.debug('[SSH] Using Agent keys as no password or key file have been specified')
+          options[:auth_methods].push('publickey')
+        end
+      end
+
       if options[:pty]
         logger.warn('[SSH] PTY requested: stderr will be merged into stdout')
       end
@@ -108,6 +113,17 @@ def validate_options(options)
       self
     end
 
+    # Creates an SSH Authentication KeyManager instance and saves it for
+    # potential future reuse.
+    #
+    # @return [Hash] hash of SSH Known Identities
+    # @api private
+    def ssh_known_identities
+      # Force KeyManager to load the key(s)
+      @manager ||= Net::SSH::Authentication::KeyManager.new(nil).each_identity {}
+      @manager.known_identities
+    end
+
     # Builds the hash of options needed by the Connection object on
     # construction.
     #
diff --git a/test/unit/transports/ssh_test.rb b/test/unit/transports/ssh_test.rb
index f037ebf2..86c7413c 100644
--- a/test/unit/transports/ssh_test.rb
+++ b/test/unit/transports/ssh_test.rb
@@ -23,6 +23,7 @@ def detect_family
     password: rand.to_s,
     key_files: rand.to_s,
   }}
+  let(:cls_agent) { cls.new({ host: rand.to_s }) }
 
   describe 'default options' do
     let(:ssh) { cls.new({ host: 'dummy' }) }
@@ -84,6 +85,26 @@ def detect_family
         "root@#{conf[:host]}",
       ])
     end
+
+    it 'sets the right auth_methods when password is specified' do
+      conf[:key_files] = nil
+      cls.new(conf).connection.method(:options).call[:auth_methods].must_equal ["none", "password", "keyboard-interactive"]
+    end
+
+    it 'sets the right auth_methods when keys are specified' do
+      conf[:password] = nil
+      cls.new(conf).connection.method(:options).call[:auth_methods].must_equal ["none", "publickey"]
+    end
+
+    it 'sets the right auth_methods for agent auth' do
+      cls_agent.stubs(:ssh_known_identities).returns({:some => 'rsa_key'})
+      cls_agent.connection.method(:options).call[:auth_methods].must_equal ['none', 'publickey']
+    end
+
+    it 'works with ssh agent auth' do
+      cls_agent.stubs(:ssh_known_identities).returns({:some => 'rsa_key'})
+      cls_agent.connection
+    end
   end
 
   describe 'converting connection to string for logging' do
@@ -111,9 +132,8 @@ def detect_family
     end
 
     it 'does not like key and password == nil' do
-      conf.delete(:password)
-      conf.delete(:key_files)
-      proc { cls.new(conf).connection }.must_raise Train::ClientError
+      cls_agent.stubs(:ssh_known_identities).returns({})
+      proc { cls_agent.connection }.must_raise Train::ClientError
     end
 
     it 'wont connect if it is not possible' do