pssh.rb in pec2-0.6.0

- old
+ new
@@ -1,40 +1,70 @@
-require 'shellwords'
+require 'net/ssh'
+require 'parallel'
+require 'colorize'
 
 module Pec2
   class Pssh
 
-    PSSH_PATH = File.expand_path('../../../exe/bin/pssh', __FILE__)
-
-    def initialize(options, hosts_file, parallel = 1)
-      @pssh_command = "#{PSSH_PATH} -t 0 -x '-tt' -h #{hosts_file} -O StrictHostKeyChecking=no"
-      if options[:print]
-        @pssh_command = "#{@pssh_command} -P"
-      end
-
-      if options[:user]
-        @pssh_command = "#{@pssh_command} -l #{options[:user]}"
-      end
-
-      if options[:log]
-        @pssh_command = "#{@pssh_command} -o #{options[:log]}"
-      end
-
-      @pssh_command = "#{@pssh_command} -p #{options[:parallel] || parallel}"
+    def initialize(options, servers, parallel = 1)
+      @parallel = parallel
+      color_index = 0
+      colors = String.colors.select{ |color|
+        !color.to_s.start_with?('light_') && !color.to_s.include?('red') && !color.to_s.include?('yellow')
+      }
+      @servers = servers.map { |server|
+        result = {}
+        result[:host] = server
+        result[:color] = colors[color_index]
+        if colors.size == color_index + 1
+          color_index = 0
+        else
+          color_index = color_index + 1
+        end
+        result
+      }
+      @user = options[:user]
+      @print = options[:print]
       @sudo_password = options[:sudo_password]
+      @ssh_options = {
+        verify_host_key: false,
+        user_known_hosts_file: '/dev/null',
+      }
+      @logger = Logger.new(STDOUT)
     end
 
-    def build_pssh_command(command)
-      if @sudo_password
-        %Q{(echo #{@sudo_password}) | #{@pssh_command} -I #{Shellwords.escape(command)}}
-      else
-        %Q{#{@pssh_command} -i #{Shellwords.escape(command)}}
-      end
-    end
-
     def exec_pssh_command(command)
       return false if command.nil? || command.empty?
-      build_command = build_pssh_command(command)
-      system(build_command)
+      error_servers = []
+      Parallel.each(@servers, in_threads: @parallel) do |server|
+        begin
+          Net::SSH.start(server[:host], @user, @ssh_options) do |ssh|
+            channel = ssh.open_channel do |channel, success|
+              channel.on_data do |channel, data|
+                if data =~ /^\[sudo\] password for /
+                  channel.send_data "#{@sudo_password}\n"
+                else
+                  data.to_s.lines.each do |line|
+                    if @print
+                      print %Q{#{server[:host]}:#{line}}.colorize(server[:color])
+                    end
+                  end
+                end
+              end
+              channel.request_pty
+              channel.exec(command)
+              channel.wait
+            end
+            channel.wait
+          end
+        rescue => e
+          error_servers << server[:host]
+          puts "\n#{e.message}\n#{e.backtrace.join("\n")}"
+        end
+      end
+      if error_servers.size > 0
+        @logger.error "error servers => #{error_servers.join(', ')}".colorize(:red)
+      end
+      return true
     end
   end
 end