<commit_before># encoding: utf-8
#
# Author:: <NAME> (<<EMAIL>>)
# Author:: <NAME> (<<EMAIL>>)
# Author:: <NAME> (<<EMAIL>>)
#
# Copyright (C) 2014, <NAME>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

require 'net/ssh'
require 'net/scp'
require 'timeout'

class Train::Transports::SSH
  # A Connection instance can be generated and re-generated, given new
  # connection details such as connection port, hostname, credentials, etc.
  # This object is responsible for carrying out the actions on the remote
  # host such as executing commands, transferring files, etc.
  #
  # @author <NAME> <<EMAIL>>
  class Connection < BaseConnection # rubocop:disable Metrics/ClassLength
    def initialize(options)
      super(options)
      @username               = @options.delete(:username)
      @hostname               = @options.delete(:hostname)
      @port                   = @options[:port] # don't delete from options
      @connection_retries     = @options.delete(:connection_retries)
      @connection_retry_sleep = @options.delete(:connection_retry_sleep)
      @max_wait_until_ready   = @options.delete(:max_wait_until_ready)
      @session                = nil
      @transport_options      = @options.delete(:transport_options)
      @cmd_wrapper            = nil
      @cmd_wrapper            = CommandWrapper.load(self, @transport_options)
    end

    # (see Base::Connection#close)
    def close
      return if @session.nil?
      logger.debug("[SSH] closing connection to #{self}")
      session.close
    ensure
      @session = nil
    end

    def os
      @os ||= OS.new(self)
    end

    def file(path)
      @files[path] ||= \
        if os.aix?
          AixFile.new(self, path)
        elsif os.solaris?
          UnixFile.new(self, path)
        else
          LinuxFile.new(self, path)
        end
    end

    # (see Base::Connection#run_command)
    def run_command(cmd)
      stdout = stderr = ''
      exit_status = nil
      cmd.dup.force_encoding('binary') if cmd.respond_to?(:force_encoding)
      logger.debug("[SSH] #{self} (#{cmd})")

      session.open_channel do |channel|
        # wrap commands if that is configured
        cmd = @cmd_wrapper.run(cmd) unless @cmd_wrapper.nil?

        if @transport_options[:pty]
          channel.request_pty do |_ch, success|
            fail Train::Transports::SSHPTYFailed, 'Requesting PTY failed' unless success
          end
        end

        channel.exec(cmd) do |_, success|
          abort 'Couldn\'t execute command on SSH.' unless success

          channel.on_data do |_, data|
            stdout += data
          end

          channel.on_extended_data do |_, _type, data|
            stderr += data
          end

          channel.on_request('exit-status') do |_, data|
            exit_status = data.read_long
          end

          channel.on_request('exit-signal') do |_, data|
            exit_status = data.read_long
          end
        end
      end
      @session.loop

      CommandResult.new(stdout, stderr, exit_status)
    rescue Net::SSH::Exception => ex
      raise Train::Transports::SSHFailed, "SSH command failed (#{ex.message})"
    end

    # (see Base::Connection#login_command)
    def login_command
      level = logger.debug? ? 'VERBOSE' : 'ERROR'
      fwd_agent = options[:forward_agent] ? 'yes' : 'no'

      args  = %w{ -o UserKnownHostsFile=/dev/null }
      args += %w{ -o StrictHostKeyChecking=no }
      args += %w{ -o IdentitiesOnly=yes } if options[:keys]
      args += %W( -o LogLevel=#{level} )
      args += %W( -o ForwardAgent=#{fwd_agent} ) if options.key?(:forward_agent)
      Array(options[:keys]).each do |ssh_key|
        args += %W( -i #{ssh_key} )
      end
      args += %W( -p #{@port} )
      args += %W( #{@username}@#{@hostname} )

      LoginCommand.new('ssh', args)
    end

    # (see Base::Connection#upload)
    def upload(locals, remote)
      Array(locals).each do |local|
        opts = File.directory?(local) ? { recursive: true } : {}

        session.scp.upload!(local, remote, opts) do |_ch, name, sent, total|
          logger.debug("Uploaded #{name} (#{total} bytes)") if sent == total
        end
      end
    rescue Net::SSH::Exception => ex
      raise Train::Transports::SSHFailed, "SCP upload failed (#{ex.message})"
    end

    # (see Base::Connection#wait_until_ready)
    def wait_until_ready
      delay = 3
      session(
        retries: @max_wait_until_ready / delay,
        delay:   delay,
        message: "Waiting for SSH service on #{@hostname}:#{@port}, " \
                 "retrying in #{delay} seconds",
      )
      execute(PING_COMMAND.dup)
    end

    def uri
      "ssh://#{@username}@#{@hostname}:#{@port}"
    end

    private

    PING_COMMAND = "echo '[SSH] Established'".freeze

    RESCUE_EXCEPTIONS_ON_ESTABLISH = [
      Errno::EACCES, Errno::EADDRINUSE, Errno::ECONNREFUSED, Errno::ETIMEDOUT,
      Errno::ECONNRESET, Errno::ENETUNREACH, Errno::EHOSTUNREACH,
      Net::SSH::Disconnect, Net::SSH::AuthenticationFailed, Net::SSH::ConnectionTimeout,
      Timeout::Error
    ].freeze

    # Establish an SSH session on the remote host.
    #
    # @param opts [Hash] retry options
    # @option opts [Integer] :retries the number of times to retry before
    #   failing
    # @option opts [Float] :delay the number of seconds to wait until
    #   attempting a retry
    # @option opts [String] :message an optional message to be logged on
    #   debug (overriding the default) when a rescuable exception is raised
    # @return [Net::SSH::Connection::Session] the SSH connection session
    # @api private
    def establish_connection(opts)
      logger.debug("[SSH] opening connection to #{self}")
      Net::SSH.start(@hostname, @username, @options)
    rescue *RESCUE_EXCEPTIONS_ON_ESTABLISH => e
      if (opts[:retries] -= 1) <= 0
        logger.warn("[SSH] connection failed, terminating (#{e.inspect})")<commit_msg>Rescue EPIPE on connect in ssh transport<commit_after># encoding: utf-8
#
# Author:: <NAME> (<<EMAIL>>)
# Author:: <NAME> (<<EMAIL>>)
# Author:: <NAME> (<<EMAIL>>)
#
# Copyright (C) 2014, <NAME>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

require 'net/ssh'
require 'net/scp'
require 'timeout'

class Train::Transports::SSH
  # A Connection instance can be generated and re-generated, given new
  # connection details such as connection port, hostname, credentials, etc.
  # This object is responsible for carrying out the actions on the remote
  # host such as executing commands, transferring files, etc.
  #
  # @author <NAME> <<EMAIL>>
  class Connection < BaseConnection # rubocop:disable Metrics/ClassLength
    def initialize(options)
      super(options)
      @username               = @options.delete(:username)
      @hostname               = @options.delete(:hostname)
      @port                   = @options[:port] # don't delete from options
      @connection_retries     = @options.delete(:connection_retries)
      @connection_retry_sleep = @options.delete(:connection_retry_sleep)
      @max_wait_until_ready   = @options.delete(:max_wait_until_ready)
      @session                = nil
      @transport_options      = @options.delete(:transport_options)
      @cmd_wrapper            = nil
      @cmd_wrapper            = CommandWrapper.load(self, @transport_options)
    end

    # (see Base::Connection#close)
    def close
      return if @session.nil?
      logger.debug("[SSH] closing connection to #{self}")
      session.close
    ensure
      @session = nil
    end

    def os
      @os ||= OS.new(self)
    end

    def file(path)
      @files[path] ||= \
        if os.aix?
          AixFile.new(self, path)
        elsif os.solaris?
          UnixFile.new(self, path)
        else
          LinuxFile.new(self, path)
        end
    end

    # (see Base::Connection#run_command)
    def run_command(cmd)
      stdout = stderr = ''
      exit_status = nil
      cmd.dup.force_encoding('binary') if cmd.respond_to?(:force_encoding)
      logger.debug("[SSH] #{self} (#{cmd})")

      session.open_channel do |channel|
        # wrap commands if that is configured
        cmd = @cmd_wrapper.run(cmd) unless @cmd_wrapper.nil?

        if @transport_options[:pty]
          channel.request_pty do |_ch, success|
            fail Train::Transports::SSHPTYFailed, 'Requesting PTY failed' unless success
          end
        end

        channel.exec(cmd) do |_, success|
          abort 'Couldn\'t execute command on SSH.' unless success

          channel.on_data do |_, data|
            stdout += data
          end

          channel.on_extended_data do |_, _type, data|
            stderr += data
          end

          channel.on_request('exit-status') do |_, data|
            exit_status = data.read_long
          end

          channel.on_request('exit-signal') do |_, data|
            exit_status = data.read_long
          end
        end
      end
      @session.loop

      CommandResult.new(stdout, stderr, exit_status)
    rescue Net::SSH::Exception => ex
      raise Train::Transports::SSHFailed, "SSH command failed (#{ex.message})"
    end

    # (see Base::Connection#login_command)
    def login_command
      level = logger.debug? ? 'VERBOSE' : 'ERROR'
      fwd_agent = options[:forward_agent] ? 'yes' : 'no'

      args  = %w{ -o UserKnownHostsFile=/dev/null }
      args += %w{ -o StrictHostKeyChecking=no }
      args += %w{ -o IdentitiesOnly=yes } if options[:keys]
      args += %W( -o LogLevel=#{level} )
      args += %W( -o ForwardAgent=#{fwd_agent} ) if options.key?(:forward_agent)
      Array(options[:keys]).each do |ssh_key|
        args += %W( -i #{ssh_key} )
      end
      args += %W( -p #{@port} )
      args += %W( #{@username}@#{@hostname} )

      LoginCommand.new('ssh', args)
    end

    # (see Base::Connection#upload)
    def upload(locals, remote)
      Array(locals).each do |local|
        opts = File.directory?(local) ? { recursive: true } : {}

        session.scp.upload!(local, remote, opts) do |_ch, name, sent, total|
          logger.debug("Uploaded #{name} (#{total} bytes)") if sent == total
        end
      end
    rescue Net::SSH::Exception => ex
      raise Train::Transports::SSHFailed, "SCP upload failed (#{ex.message})"
    end

    # (see Base::Connection#wait_until_ready)
    def wait_until_ready
      delay = 3
      session(
        retries: @max_wait_until_ready / delay,
        delay:   delay,
        message: "Waiting for SSH service on #{@hostname}:#{@port}, " \
                 "retrying in #{delay} seconds",
      )
      execute(PING_COMMAND.dup)
    end

    def uri
      "ssh://#{@username}@#{@hostname}:#{@port}"
    end

    private

    PING_COMMAND = "echo '[SSH] Established'".freeze

    RESCUE_EXCEPTIONS_ON_ESTABLISH = [
      Errno::EACCES, Errno::EADDRINUSE, Errno::ECONNREFUSED, Errno::ETIMEDOUT,
      Errno::ECONNRESET, Errno::ENETUNREACH, Errno::EHOSTUNREACH, Errno::EPIPE,
      Net::SSH::Disconnect, Net::SSH::AuthenticationFailed, Net::SSH::ConnectionTimeout,
      Timeout::Error
    ].freeze

    # Establish an SSH session on the remote host.
    #
    # @param opts [Hash] retry options
    # @option opts [Integer] :retries the number of times to retry before
    #   failing
    # @option opts [Float] :delay the number of seconds to wait until
    #   attempting a retry
    # @option opts [String] :message an optional message to be logged on
    #   debug (overriding the default) when a rescuable exception is raised
    # @return [Net::SSH::Connection::Session] the SSH connection session
    # @api private
    def establish_connection(opts)
      logger.debug("[SSH] opening connection to #{self}")
      Net::SSH.start(@hostname, @username, @options)
    rescue *RESCUE_EXCEPTIONS_ON_ESTABLISH => e
      if (opts[:retries] -= 1) <= 0
        logger.warn("[SSH] connection failed, terminating (#{e.inspect})")
{'fraction_non_alphanumeric': 0.11914926647519987, 'fraction_numerical': 0.0018629201272995421, 'mean_word_length': 2.674843126069595, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 35, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4996711', 'n_tokens_mistral': 3957, 'n_tokens_neox': 3886, 'n_words': 1086}