Add transport for Cisco IOS (#271)

* Add transport for Cisco IOS
* Add more unit tests and minor style changes
* Clean up code after chatting with @TrevorBramble
* Modify `#connection` to only validate options once
* Modify `#format_result` to return `CommandResult`
* Remove useless `nil` assignments

Signed-off-by: Jerry Aldrich <jerryaldrichiii@gmail.com>

lib/train/transports/cisco_ios.rb

@@ -0,0 +1,140 @@
+# encoding: utf-8
+
+require 'train/plugins'
+require 'train/transports/ssh'
+
+module Train::Transports
+  class BadEnablePassword < Train::TransportError; end
+
+  class CiscoIOS < SSH
+    name 'cisco_ios'
+
+    option :host, required: true
+    option :user, required: true
+    option :port, default: 22, required: true
+
+    option :password, required: true
+
+    # Used to elevate to enable mode (similar to `sudo su` in Linux)
+    option :enable_password
+
+    def connection
+      @connection ||= Connection.new(validate_options(@options).options)
+    end
+
+    class Connection < BaseConnection
+      def initialize(options)
+        super(options)
+
+        # Delete options to avoid passing them in to `Net::SSH.start` later
+        @host = @options.delete(:host)
+        @user = @options.delete(:user)
+        @port = @options.delete(:port)
+        @enable_password = @options.delete(:enable_password)
+
+        @prompt = /^\S+[>#]\r\n.*$/
+      end
+
+      def uri
+        "ssh://#{@user}@#{@host}:#{@port}"
+      end
+
+      private
+
+      def establish_connection
+        logger.debug("[SSH] opening connection to #{self}")
+
+        Net::SSH.start(
+          @host,
+          @user,
+          @options.reject { |_key, value| value.nil? },
+        )
+      end
+
+      def session
+        return @session unless @session.nil?
+
+        @session = open_channel(establish_connection)
+
+        # Escalate privilege to enable mode if password is given
+        if @enable_password
+          run_command_via_connection("enable\r\n#{@enable_password}")
+        end
+
+        # Prevent `--MORE--` by removing terminal length limit
+        run_command_via_connection('terminal length 0')
+
+        @session
+      end
+
+      def run_command_via_connection(cmd)
+        # Ensure buffer is empty before sending data
+        @buf = ''
+
+        logger.debug("[SSH] Running `#{cmd}` on #{self}")
+        session.send_data(cmd + "\r\n")
+
+        logger.debug('[SSH] waiting for prompt')
+        until @buf =~ @prompt
+          raise BadEnablePassword if @buf =~ /Bad secrets/
+          session.connection.process(0)
+        end
+
+        # Save the buffer and clear it for the next command
+        output = @buf.dup
+        @buf = ''
+
+        format_result(format_output(output, cmd))
+      end
+
+      ERROR_MATCHERS = [
+        'Bad IP address',
+        'Incomplete command',
+        'Invalid input detected',
+        'Unrecognized host',
+      ].freeze
+
+      # IOS commands do not have an exit code so we must compare the command
+      # output with partial segments of known errors. Then, we return a
+      # `CommandResult` with arguments in the correct position based on the
+      # result.
+      def format_result(result)
+        if ERROR_MATCHERS.none? { |e| result.include?(e) }
+          CommandResult.new(result, '', 0)
+        else
+          CommandResult.new('', result, 1)
+        end
+      end
+
+      # The buffer (@buf) contains all data sent/received on the SSH channel so
+      # we need to format the data to match what we would expect from Train
+      def format_output(output, cmd)
+        leading_prompt = /(\r\n|^)\S+[>#]/
+        command_string = /#{cmd}\r\n/
+        trailing_prompt = /\S+[>#](\r\n|$)/
+        trailing_line_endings = /(\r\n)+$/
+
+        output
+          .sub(leading_prompt, '')
+          .sub(command_string, '')
+          .gsub(trailing_prompt, '')
+          .gsub(trailing_line_endings, '')
+      end
+
+      # Create an SSH channel that writes to @buf when data is received
+      def open_channel(ssh)
+        logger.debug("[SSH] opening SSH channel to #{self}")
+        ssh.open_channel do |ch|
+          ch.on_data do |_, data|
+            @buf += data
+          end
+
+          ch.send_channel_request('shell') do |_, success|
+            raise 'Failed to open SSH shell' unless success
+            logger.debug('[SSH] shell opened')
+          end
+        end
+      end
+    end
+  end
+end

test/unit/transports/cisco_ios.rb

@@ -0,0 +1,94 @@
+# encoding: utf-8
+
+require 'helper'
+require 'train/transports/cisco_ios'
+
+describe 'Train::Transports::CiscoIOS' do
+  let(:cls) do
+    plat = Train::Platforms.name('mock').in_family('cisco_ios')
+    plat.add_platform_methods
+    Train::Platforms::Detect.stubs(:scan).returns(plat)
+    Train::Transports::CiscoIOS
+  end
+
+  let(:opts) do
+    {
+      host: 'fakehost',
+      user: 'fakeuser',
+      password: 'fakepassword',
+    }
+  end
+
+  let(:cisco_ios) do
+    cls.new(opts)
+  end
+
+  describe 'CiscoIOS::Connection' do
+    let(:connection) { cls.new(opts).connection }
+
+    describe '#initialize' do
+      it 'raises an error when user is missing' do
+        opts.delete(:user)
+        err = proc { cls.new(opts).connection }.must_raise(Train::ClientError)
+        err.message.must_match(/must provide.*user/)
+      end
+
+      it 'raises an error when host is missing' do
+        opts.delete(:host)
+        err = proc { cls.new(opts).connection }.must_raise(Train::ClientError)
+        err.message.must_match(/must provide.*host/)
+      end
+
+      it 'raises an error when password is missing' do
+        opts.delete(:password)
+        err = proc { cls.new(opts).connection }.must_raise(Train::ClientError)
+        err.message.must_match(/must provide.*password/)
+      end
+
+      it 'provides a uri' do
+        connection.uri.must_equal 'ssh://fakeuser@fakehost:22'
+      end
+    end
+
+    describe '#format_result' do
+      it 'returns correctly when result is `good`' do
+        output = 'good'
+        Train::Extras::CommandResult.expects(:new).with(output, '', 0)
+        connection.send(:format_result, 'good')
+      end
+
+      it 'returns correctly when result matches /Bad IP address/' do
+        output = "Translating \"nope\"\r\n\r\nTranslating \"nope\"\r\n\r\n% Bad IP address or host name\r\n% Unknown command or computer name, or unable to find computer address\r\n"
+        Train::Extras::CommandResult.expects(:new).with('', output, 1)
+        connection.send(:format_result, output)
+      end
+
+      it 'returns correctly when result matches /Incomplete command/' do
+        output = "% Incomplete command.\r\n\r\n"
+        Train::Extras::CommandResult.expects(:new).with('', output, 1)
+        connection.send(:format_result, output)
+      end
+
+      it 'returns correctly when result matches /Invalid input detected/' do
+        output = "             ^\r\n% Invalid input detected at '^' marker.\r\n\r\n"
+        Train::Extras::CommandResult.expects(:new).with('', output, 1)
+        connection.send(:format_result, output)
+      end
+
+      it 'returns correctly when result matches /Unrecognized host/' do
+        output = "Translating \"nope\"\r\n% Unrecognized host or address, or protocol not running.\r\n\r\n"
+        Train::Extras::CommandResult.expects(:new).with('', output, 1)
+        connection.send(:format_result, output)
+      end
+    end
+
+    describe '#format_output' do
+      it 'returns output containing only the output of the command executed' do
+        cmd = 'show calendar'
+        output = "show calendar\r\n10:35:50 UTC Fri Mar 23 2018\r\n7200_ios_12#\r\n7200_ios_12#"
+        result = connection.send(:format_output, output, cmd)
+        result.must_equal '10:35:50 UTC Fri Mar 23 2018'
+      end
+    end
+  end
+end