require 'chef/knife'

# отличается от ssh:
# 1. host_key_verify => false
# 2. concurrency => 50
# 3. коннектиться по дефолту к private_ip
# 4. выдает exit status не нулевой если есть сфеленные по коннекту
# 5. выводит список всех сфейленных машин в самом конце

class Chef
  class Knife
    class CloudSsh < Knife

      deps do
        require 'net/ssh'
        require 'net/ssh/multi'
        require 'readline'
        require 'chef/exceptions'
        require 'chef/search/query'
        require 'chef/mixin/shell_out'
        require 'mixlib/shellout'
      end

      include Chef::Mixin::ShellOut

      attr_writer :password

      banner "knife cloud_ssh QUERY COMMAND (options)"

      option :concurrency,
        :short => "-C NUM",
        :long => "--concurrency NUM",
        :description => "The number of concurrent connections",
        :default => 50,
        :proc => lambda { |o| o.to_i }

      option :attribute,
        :short => "-a ATTR",
        :long => "--attribute ATTR",
        :description => "The attribute to use for opening the connection - default depends on the context",
        :proc => Proc.new { |key| Chef::Config[:knife][:ssh_attribute] = key.strip }

      option :manual,
        :short => "-m",
        :long => "--manual-list",
        :boolean => true,
        :description => "QUERY is a space separated list of servers",
        :default => false

      option :ssh_user,
        :short => "-x USERNAME",
        :long => "--ssh-user USERNAME",
        :description => "The ssh username"

      option :ssh_password,
        :short => "-P PASSWORD",
        :long => "--ssh-password PASSWORD",
        :description => "The ssh password"

      option :ssh_port,
        :short => "-p PORT",
        :long => "--ssh-port PORT",
        :description => "The ssh port",
        :proc => Proc.new { |key| Chef::Config[:knife][:ssh_port] = key.strip }

      option :ssh_gateway,
        :short => "-G GATEWAY",
        :long => "--ssh-gateway GATEWAY",
        :description => "The ssh gateway",
        :proc => Proc.new { |key| Chef::Config[:knife][:ssh_gateway] = key.strip }

      option :forward_agent,
        :short => "-A",
        :long => "--forward-agent",
        :description => "Enable SSH agent forwarding",
        :boolean => true

      option :identity_file,
        :short => "-i IDENTITY_FILE",
        :long => "--identity-file IDENTITY_FILE",
        :description => "The SSH identity file used for authentication"

      option :host_key_verify,
        :long => "--[no-]host-key-verify",
        :description => "Verify host key, disabled by default.",
        :boolean => true,
        :default => false

      def session
        config[:on_error] ||= :skip
        ssh_error_handler = Proc.new do |server|

          @failed_connect_nodes << { "node" => server.host , "message" => "#{$!.class.name}: #{$!.message}" }
          if config[:manual]
            node_name = server.host
          else
            @action_nodes.each do |n|
              node_name = n if format_for_display(n)[config[:attribute]] == server.host
            end
          end
          case config[:on_error]
          when :skip
            ui.warn "Failed to connect to #{node_name} -- #{$!.class.name}: #{$!.message}"
            $!.backtrace.each { |l| Chef::Log.error(l) }
          when :raise
            #Net::SSH::Multi magic to force exception to be re-raised.
            throw :go, :raise
          end
        end

        @session ||= Net::SSH::Multi.start(:concurrent_connections => config[:concurrency], :on_error => ssh_error_handler)
      end

      def configure_gateway
        config[:ssh_gateway] ||= Chef::Config[:knife][:ssh_gateway]
        if config[:ssh_gateway]
          gw_host, gw_user = config[:ssh_gateway].split('@').reverse
          gw_host, gw_port = gw_host.split(':')
          gw_opts = gw_port ? { :port => gw_port } : {}

          session.via(gw_host, gw_user || config[:ssh_user], gw_opts)
        end
      rescue Net::SSH::AuthenticationFailed
        user = gw_user || config[:ssh_user]
        prompt = "Enter the password for #{user}@#{gw_host}: "
        gw_opts.merge!(:password => prompt_for_password(prompt))
        session.via(gw_host, user, gw_opts)
      end

      def configure_session
        list = case config[:manual]
               when true
                 @all_nodes += @name_args[0].split(" ")
                 @name_args[0].split(" ")
               when false
                 r = Array.new
                 q = Chef::Search::Query.new
                 @action_nodes = q.search(:node, @name_args[0])[0]
                 @action_nodes.each do |item|
                   # we should skip the loop to next iteration if the item returned by the search is nil
                   next if item.nil?
                   # if a command line attribute was not passed, and we have a cloud public_hostname, use that.
                   # see #configure_attribute for the source of config[:attribute] and config[:override_attribute]
                   if !config[:override_attribute] && item[:cloud] and item[:cloud][:public_hostname]
                     i = item[:cloud][:public_hostname]
                   elsif config[:override_attribute]
                     i = extract_nested_value(item, config[:override_attribute])
                   else
                     i = extract_nested_value(item, config[:attribute])
                   end
                   # next if we couldn't find the specified attribute in the returned node object
                   next if i.nil?
                   r.push(i)
                   @all_nodes << item[:fqdn]
                 end
                 r
               end
        if list.length == 0
          if @action_nodes.length == 0
            ui.fatal("No nodes returned from search!")
          else
            ui.fatal("#{@action_nodes.length} #{@action_nodes.length > 1 ? "nodes":"node"} found, " +
                     "but does not have the required attribute to establish the connection. " +
                     "Try setting another attribute to open the connection using --attribute.")
          end
          exit 10
        end
        session_from_list(list)
      end

      def session_from_list(list)
        list.each do |item|
          Chef::Log.debug("Adding #{item}")
          session_opts = {}

          ssh_config = Net::SSH.configuration_for(item)

          # Chef::Config[:knife][:ssh_user] is parsed in #configure_user and written to config[:ssh_user]
          user = config[:ssh_user] || ssh_config[:user]
          hostspec = user ? "#{user}@#{item}" : item
          session_opts[:keys] = File.expand_path(config[:identity_file]) if config[:identity_file]
          session_opts[:keys_only] = true if config[:identity_file]
          session_opts[:password] = config[:ssh_password] if config[:ssh_password]
          session_opts[:forward_agent] = config[:forward_agent]
          session_opts[:port] = config[:ssh_port] || Chef::Config[:knife][:ssh_port] || ssh_config[:port]
          session_opts[:logger] = Chef::Log.logger if Chef::Log.level == :debug

          if !config[:host_key_verify]
            session_opts[:paranoid] = false
            session_opts[:user_known_hosts_file] = "/dev/null"
          end

          session.use(hostspec, session_opts)

          @longest = item.length if item.length > @longest
        end

        session
      end

      def fixup_sudo(command)
        command.sub(/^sudo/, 'sudo -p \'knife sudo password: \'')
      end

      def print_data(host, data)
        @buffers ||= {}
        if leftover = @buffers[host]
          @buffers[host] = nil
          print_data(host, leftover + data)
        else
          if newline_index = data.index("\n")
            line = data.slice!(0...newline_index)
            data.slice!(0)
            print_line(host, line)
            print_data(host, data)
          else
            @buffers[host] = data
          end
        end
      end

      def print_line(host, data)
        padding = @longest - host.length
        str = ui.color(host, :cyan) + (" " * (padding + 1)) + data
        ui.msg(str)
      end

      def ssh_command(command, subsession=nil)
        exit_status = 0
        subsession ||= session
        command = fixup_sudo(command)
        command.force_encoding('binary') if command.respond_to?(:force_encoding)
        subsession.open_channel do |ch|
          ch.request_pty
          ch.exec command do |ch, success|
            raise ArgumentError, "Cannot execute #{command}" unless success
            ch.on_data do |ichannel, data|
              print_data(ichannel[:host], data)
              if data =~ /^knife sudo password: /
                print_data(ichannel[:host], "\n")
                ichannel.send_data("#{get_password}\n")
              end
            end
            ch.on_request "exit-status" do |ichannel, data|
              exit_status = [exit_status, data.read_long].max
              if exit_status != 0
                @not_zerro_nodes << { 'node' => ichannel[:host], "message" => "exit status: #{exit_status}" }
              else
                @zerro_nodes << {'node' => ichannel[:host]}
              end
            end
          end
        end
        session.loop
        exit_status
      end

      def get_password
        @password ||= prompt_for_password
      end

      def prompt_for_password(prompt = "Enter your password: ")
        ui.ask(prompt) { |q| q.echo = false }
      end

      def configure_attribute
        config[:override_attribute] = config[:attribute] || Chef::Config[:knife][:ssh_attribute]
        config[:attribute] = (Chef::Config[:knife][:ssh_attribute] ||
                              config[:attribute] ||
                              "private_ipaddress").strip
      end

      def read_line
        loop do
          command = reader.readline("#{ui.color('knife-ssh>', :bold)} ", true)

          if command.nil?
            command = "exit"
            puts(command)
          else
            command.strip!
          end

          unless command.empty?
            return command
          end
        end
      end

      def reader
        Readline
      end

      def interactive
        puts "Connected to #{ui.list(session.servers_for.collect { |s| ui.color(s.host, :cyan) }, :inline, " and ")}"
        puts
        puts "To run a command on a list of servers, do:"
        puts "  on SERVER1 SERVER2 SERVER3; COMMAND"
        puts "  Example: on latte foamy; echo foobar"
        puts
        puts "To exit interactive mode, use 'quit!'"
        puts
        while 1
          command = read_line
          case command
          when 'quit!'
            puts 'Bye!'
            break
          when /^on (.+?); (.+)$/
            raw_list = $1.split(" ")
            server_list = Array.new
            session.servers.each do |session_server|
              server_list << session_server if raw_list.include?(session_server.host)
            end
            command = $2
            ssh_command(command, session.on(*server_list))
          else
            ssh_command(command)
          end
        end
      end

      def get_stripped_unfrozen_value(value)
        return nil if value.nil?
        value.strip
      end

      def configure_user
        config[:ssh_user] = get_stripped_unfrozen_value(config[:ssh_user] ||
                             Chef::Config[:knife][:ssh_user])
      end

      def configure_identity_file
        config[:identity_file] = get_stripped_unfrozen_value(config[:identity_file] ||
                             Chef::Config[:knife][:ssh_identity_file])
      end

      def extract_nested_value(data_structure, path_spec)
        ui.presenter.extract_nested_value(data_structure, path_spec)
      end

      def print_all_status
        succ = @all_nodes.count - @failed_connect_nodes.count - @not_zerro_nodes.count
        failed = @failed_connect_nodes.count + @not_zerro_nodes.count
        @failed_connect_nodes.each { |msg| ui.error("#{msg['node']}  #{msg['message']}") }
        @not_zerro_nodes.each { |msg| ui.error("#{msg['node']}  #{msg['message']}") }
        if failed > 0 || @zerro_nodes.count != succ
          ui.warn "All: #{@all_nodes.count}, successful exit: #{@zerro_nodes.count}, connection problem: #{@failed_connect_nodes.count}, failed: #{@not_zerro_nodes.count}"
        else
          ui.info "All: #{@all_nodes.count}, successful: #{succ}"
        end
      end

      def run
        extend Chef::Mixin::Command

        @longest = 0

        @all_nodes = []
        @failed_connect_nodes = []
        @not_zerro_nodes = []
        @zerro_nodes = []

        configure_attribute
        configure_user
        configure_identity_file
        configure_gateway
        configure_session

        exit_status =
        case @name_args[1]
        when "interactive"
          interactive
        else
          ssh_command(@name_args[1..-1].join(" "))
        end


        session.close

        print_all_status

        if (@not_zerro_nodes + @failed_connect_nodes).compact.empty?
          exit_status
        else
          my_status = exit_status == 0 ? 1 : exit_status
          exit my_status
        end
      end

    end
  end
end