Adding proper bastion support (#310)

* Adding proper bastion support
* Removing commented line
* Updating code to use proxy_command as per PR review
* Formatting fix

Signed-off-by: Noel Georgi <18496730+frezbo@users.noreply.github.com>

.gitignore

@@ -2,3 +2,4 @@ train-*.gem
 r-train-*.gem
 Gemfile.lock
 .kitchen/
+TAGS

lib/train/transports/ssh.rb

@@ -58,6 +58,10 @@ class SSH < Train.plugin(1) # rubocop:disable Metrics/ClassLength
     option :max_wait_until_ready, default: 600
     option :compression, default: false
     option :pty, default: false
+    option :proxy_command, default: nil
+    option :bastion_host, default: nil
+    option :bastion_user, default: 'root'
+    option :bastion_port, default: 22
 
     option :compression_level do |opts|
       # on nil or false: set compression level to 0
@@ -109,6 +113,10 @@ def validate_options(options)
         logger.warn('[SSH] PTY requested: stderr will be merged into stdout')
       end
 
+      if [options[:proxy_command], options[:bastion_host]].all? { |type| !type.nil? }
+        fail Train::ClientError, 'Only one of proxy_command or bastion_host needs to be specified'
+      end
+
       super
       self
     end
@@ -151,6 +159,9 @@ def connection_options(opts)
         password:               opts[:password],
         forward_agent:          opts[:forward_agent],
         proxy_command:          opts[:proxy_command],
+        bastion_host:           opts[:bastion_host],
+        bastion_user:           opts[:bastion_user],
+        bastion_port:           opts[:bastion_port],
         transport_options:      opts,
       }
 

lib/train/transports/ssh_connection.rb

@@ -43,6 +43,10 @@ def initialize(options)
       @session                = nil
       @transport_options      = @options.delete(:transport_options)
       @cmd_wrapper            = nil
+      @proxy_command          = @options.delete(:proxy_command)
+      @bastion_host           = @options.delete(:bastion_host)
+      @bastion_user           = @options.delete(:bastion_user)
+      @bastion_port           = @options.delete(:bastion_port)
       @cmd_wrapper            = CommandWrapper.load(self, @transport_options)
     end
 
@@ -55,8 +59,7 @@ def close
       @session = nil
     end
 
-    # (see Base::Connection#login_command)
-    def login_command
+    def ssh_opts
       level = logger.debug? ? 'VERBOSE' : 'ERROR'
       fwd_agent = options[:forward_agent] ? 'yes' : 'no'
 
@@ -65,13 +68,32 @@ def login_command
       args += %w{ -o IdentitiesOnly=yes } if options[:keys]
       args += %W( -o LogLevel=#{level} )
       args += %W( -o ForwardAgent=#{fwd_agent} ) if options.key?(:forward_agent)
-      args += %W( -o ProxyCommand='#{options[:proxy_command]}' ) unless options[:proxy_command].nil?
       Array(options[:keys]).each do |ssh_key|
         args += %W( -i #{ssh_key} )
       end
+      args
+    end
+
+    def check_proxy
+      [@proxy_command, @bastion_host].any? { |type| !type.nil? }
+    end
+
+    def generate_proxy_command
+      return @proxy_command unless @proxy_command.nil?
+      args = %w{ ssh }
+      args += ssh_opts
+      args += %W( #{@bastion_user}@#{@bastion_host} )
+      args += %W( -p #{@bastion_port} )
+      args += %w{ -W %h:%p }
+      args.join(' ')
+    end
+
+    # (see Base::Connection#login_command)
+    def login_command
+      args = ssh_opts
+      args += %W( -o ProxyCommand='#{generate_proxy_command}' ) if check_proxy
       args += %W( -p #{@port} )
       args += %W( #{@username}@#{@hostname} )
-
       LoginCommand.new('ssh', args)
     end
 
@@ -145,10 +167,9 @@ def uri
     # @api private
     def establish_connection(opts)
       logger.debug("[SSH] opening connection to #{self}")
-      if @options[:proxy_command]
+      if check_proxy
         require 'net/ssh/proxy/command'
-        @options[:proxy] = Net::SSH::Proxy::Command.new(@options[:proxy_command])
-        @options.delete(:proxy_command)
+        @options[:proxy] = Net::SSH::Proxy::Command.new(generate_proxy_command)
       end
       Net::SSH.start(@hostname, @username, @options.clone.delete_if { |_key, value| value.nil? })
     rescue *RESCUE_EXCEPTIONS_ON_ESTABLISH => e

test/unit/transports/ssh_test.rb

@@ -96,8 +96,8 @@
         "-o", "IdentitiesOnly=yes",
         "-o", "LogLevel=VERBOSE",
         "-o", "ForwardAgent=no",
-        "-o", "ProxyCommand='ssh root@127.0.0.1 -W %h:%p'",
         "-i", conf[:key_files],
+        "-o", "ProxyCommand='ssh root@127.0.0.1 -W %h:%p'",
         "-p", "22",
         "root@#{conf[:host]}",
       ])
@@ -175,3 +175,144 @@
     end
   end
 end
+
+describe 'ssh transport with bastion' do
+  let(:cls) do
+    plat = Train::Platforms.name('mock').in_family('linux')
+    plat.add_platform_methods
+    Train::Platforms::Detect.stubs(:scan).returns(plat)
+    Train::Transports::SSH
+  end
+
+  let(:conf) {{
+    host: rand.to_s,
+    password: rand.to_s,
+    key_files: rand.to_s,
+    bastion_host: 'bastion_dummy',
+  }}
+  let(:cls_agent) { cls.new({ host: rand.to_s }) }
+
+  describe 'bastion' do
+    describe 'default options' do
+      let(:ssh) { cls.new({ bastion_host: 'bastion_dummy' }) }
+
+      it 'configures the host' do
+        ssh.options[:bastion_host].must_equal 'bastion_dummy'
+      end
+
+      it 'has default port' do
+        ssh.options[:bastion_port].must_equal 22
+      end
+
+      it 'has default user' do
+        ssh.options[:bastion_user].must_equal 'root'
+      end
+    end
+
+    describe 'opening a connection' do
+      let(:ssh) { cls.new(conf) }
+      let(:connection) { ssh.connection }
+
+      it 'provides a run_command_via_connection method' do
+        methods = connection.class.private_instance_methods(false)
+        methods.include?(:run_command_via_connection).must_equal true
+      end
+
+      it 'provides a file_via_connection method' do
+        methods = connection.class.private_instance_methods(false)
+        methods.include?(:file_via_connection).must_equal true
+      end
+
+      it 'gets the connection' do
+        connection.must_be_kind_of Train::Transports::SSH::Connection
+      end
+
+      it 'provides a uri' do
+        connection.uri.must_equal "ssh://root@#{conf[:host]}:22"
+      end
+
+      it 'must respond to wait_until_ready' do
+        connection.must_respond_to :wait_until_ready
+      end
+
+      it 'can be closed' do
+        connection.close.must_be_nil
+      end
+
+      it 'has a login command == ssh' do
+        connection.login_command.command.must_equal 'ssh'
+      end
+
+      it 'has login command arguments' do
+        connection.login_command.arguments.must_equal([
+          "-o", "UserKnownHostsFile=/dev/null",
+          "-o", "StrictHostKeyChecking=no",
+          "-o", "IdentitiesOnly=yes",
+          "-o", "LogLevel=VERBOSE",
+          "-o", "ForwardAgent=no",
+          "-i", conf[:key_files],
+          "-o", "ProxyCommand='ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -o IdentitiesOnly=yes -o LogLevel=VERBOSE -o ForwardAgent=no -i #{conf[:key_files]} root@bastion_dummy -p 22 -W %h:%p'",
+          "-p", "22",
+          "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
+
+      it 'sets up a proxy when ssh proxy command is specified' do
+        mock = MiniTest::Mock.new
+        mock.expect(:call, true) do |hostname, username, options|
+          options[:proxy].kind_of?(Net::SSH::Proxy::Command) &&
+            "ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -o IdentitiesOnly=yes -o LogLevel=VERBOSE -o ForwardAgent=no -i #{conf[:key_files]} root@bastion_dummy -p 22 -W %h:%p" == options[:proxy].command_line_template
+        end
+        connection.stubs(:run_command)
+        Net::SSH.stub(:start, mock) do
+          connection.wait_until_ready
+        end
+        mock.verify
+      end
+    end
+  end
+end
+
+describe 'ssh transport with bastion and proxy' do
+  let(:cls) do
+    plat = Train::Platforms.name('mock').in_family('linux')
+    plat.add_platform_methods
+    Train::Platforms::Detect.stubs(:scan).returns(plat)
+    Train::Transports::SSH
+  end
+
+  let(:conf) {{
+    host: rand.to_s,
+    password: rand.to_s,
+    key_files: rand.to_s,
+    bastion_host: 'bastion_dummy',
+    proxy_command: 'dummy'
+  }}
+  let(:cls_agent) { cls.new({ host: rand.to_s }) }
+
+  describe 'bastion and proxy' do
+    it 'will throw an exception when both proxy_command and bastion_host is specified' do
+      proc { cls.new(conf).connection }.must_raise Train::ClientError
+    end
+  end
+end