Move Cisco IOS connection under SSH transport

lib/train/transports/cisco_ios_connection.rb

@@ -0,0 +1,119 @@
+# encoding: utf-8
+
+class Train::Transports::SSH
+  class CiscoIOSConnection < BaseConnection
+    class BadEnablePassword < Train::TransportError; end
+
+    def initialize(options)
+      super(options)
+
+      # Extract 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)
+
+      # Use all options left that are not `nil` for `Net::SSH.start` later
+      @ssh_options = options.reject { |_key, value| value.nil? }
+
+      @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, @ssh_options)
+    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 = /#{Regexp.quote(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

lib/train/transports/ssh.rb

@@ -37,6 +37,7 @@ class SSH < Train.plugin(1) # rubocop:disable Metrics/ClassLength
     name 'ssh'
 
     require 'train/transports/ssh_connection'
+    require 'train/transports/cisco_ios_connection'
 
     # add options for submodules
     include_options Train::Extras::CommandWrapper
@@ -193,6 +194,19 @@ def create_new_connection(options, &block)
 
       @connection_options = options
       conn = Connection.new(options, &block)
+
+      # Cisco IOS requires a special implementation of `Net:SSH`. This uses the
+      # SSH transport to identify the platform, but then replaces SSHConnection
+      # with a CiscoIOSConnection in order to behave as expected for the user.
+      if defined?(conn.platform.cisco_ios?) && conn.platform.cisco_ios?
+        ios_options = {}
+        ios_options[:host] = @options[:host]
+        ios_options[:user] = @options[:user]
+        ios_options[:enable_password] = @options[:enable_password]
+        ios_options.merge!(@connection_options)
+        conn = CiscoIOSConnection.new(ios_options)
+      end
+
       @connection = conn unless conn.nil?
     end