Manage SSH tunnels to remote hosts.

tunnels.rb

require 'net/ssh/gateway'

module Net
  module SSH

    # A class that manages URLs, tunnelling them if necessary.
    class Tunnels

      HP = Struct.new(:host, :port) do
        def to_s
          '%s:%d' % [self.host, self.port]
        end
      end
      GW = Struct.new(:user, :host, :port, :gateway)

      def initialize(loggers=[])
        @gateway_mappings = {}
        @gateways = []
        @loggers = Array(loggers)
      end

      # Add a gateway and an array of destinations to which it provides tunnels.
      # - gateway is of the form host[:port] (default port is 22)
      # - destination is a URI or host[:port] combination (default port is 22)
      def add_gateway(gateway, user, pass, destns)
        # The gateway may need to be tunnelled.
        gw_hp = host_port(tunnel(gateway))
        gw = Net::SSH::Gateway.new(gw_hp.host, user, password: pass, port: gw_hp.port)
        @gateways << GW.new(user, gw_hp.host, gw_hp.port, gw)
        log_info "Connected to gateway at #{user}@#{gw_hp.host}:#{gw_hp.port}"
        destns.each do |destn|
          gw_hp = host_port(destn)
          raise "Invalid destination '#{destn}'" unless gw_hp
          local_port = gw.open(gw_hp.host, gw_hp.port)
          log_info "=> Created tunnel via port #{local_port} to #{gw_hp}"
          @gateway_mappings[gw_hp] = local_port
        end
      end

      HOST_PORT_RE = /^(?<host>[^:]+):(?<port>\d+)/
      JDBC_RE = /[jJ]dbc:[^@]+@(?<host>[^:]+):(?<port>\d+)/
      JDBC_RE_2 = /[jJ]dbc:.+HOST=\s*(?<host>[^\s\)]+).*PORT=\s*(?<port>\d+)/
      # Extract the host and port from a destination.
      def host_port(destn)
        # Match host and port some way.
        match = HOST_PORT_RE.match(destn) || JDBC_RE.match(destn) || JDBC_RE_2.match(destn)
        if match
          HP.new(match[:host], match[:port].to_i)
        else
          uri = URI(destn)
          if uri.host && uri.port
            # Destination is a URI.
            HP.new(uri.host, uri.port)
          else
            # Assume hostname only, on port 22 (for an SSH gateway).
            HP.new(destn, 22)
          end
        end
      end

      # Provide the tunnelled version of a URL or host:port if necessary.
      def tunnel(destn)
        hp = host_port(destn)
        if @gateway_mappings[hp]
          sub_tunnel(destn, hp, @gateway_mappings[hp])
        else
          destn
        end
      end

      # Substitute the tunnelled host and port into a destination.
      def sub_tunnel(destn, hp,  local_port)
        uri = URI(destn)
        if uri.host && uri.port
          uri.host = 'localhost'
          uri.port = local_port
          uri.to_s
        elsif HOST_PORT_RE.match(destn) || JDBC_RE.match(destn) || JDBC_RE_2.match(destn)
          destn.sub(hp.host, 'localhost').sub(hp.port.to_s, local_port.to_s)
        else
          "#{destn.sub(hp.host, 'localhost')}:#{local_port}"
        end
      end

      # Shut down the gateways in reverse order of creation.
      def tidy_up
        @gateways.reverse_each do |gw|
          gw.gateway.shutdown!
          log_info "Disconnected from gateway #{gw.user}@#{gw.host}:#{gw.port}"
        end
      end

      def log_debug(msg)
        @loggers.each { |l| l.debug msg }
      end

      def log_info(msg)
        @loggers.each { |l| l.info msg }
      end

    end

  end
end