lib/shexy.rb in shexy-0.2 vs lib/shexy.rb in shexy-0.3

@@ -28,44 +28,80 @@
 #
 # Copying files (local -> remote):
 #
 #     Shexy.copy_to 'test@test-host', '/home/rubiojr/my-uber-file', '/tmp/'
 #
+
 module Shexy
+    
+  VERSION = '0.3'
 
-  VERSION = '0.2'
+  [:user, :password, :key, :cmd, :host].each do |n|
+    instance_eval %{
+      def #{n}; Thread.current[:shexy_#{n}]; end
+      def #{n}=(v); Thread.current[:shexy_#{n}] = v; end
+    }
+  end
 
-  @flags = {}
+  def self.flags=(f);Thread.current[:shexy_flags] = f;end
+  def self.flags; Thread.current[:shexy_flags] ||= {} ; end
+  def self.sudo?;Thread.current[:shexy_use_sudo]; end
+  def self.use_sudo(v = true); Thread.current[:shexy_use_sudo] = v; end
 
-  def self.password=(password); flags[:password] = password; end
-  def self.password; flags[:password]; end
-  def self.key=(key); flags[:keys] = [File.expand_path(key)]; end
-  def self.key; flags[:keys]; end
-  def self.use_sudo(v=true); @sudo = v; end
-  def self.sudo?; @sudo ||= false; end
+  def self.wait_for_ssh(timeout = 60)
+    Timeout.timeout(timeout) do
+      begin
+        sleep(1) until tcp_test_ssh(host) do
+        end
+      rescue Errno::ECONNRESET
+        # safe to ignore, we need to retry all the time.
+      end
+    end
+    true
+  rescue Exception => e
+    $stderr.puts e.message
+    false
+  end
 
-  class << self
-    attr_accessor :host, :user, :flags
-    attr_reader :cmd
+  def self.tcp_test_ssh
+    tcp_socket = TCPSocket.new(host, 22)
+    readable = IO.select([tcp_socket], nil, nil, 5)
+    if readable
+      yield
+      true
+    else
+      false
+    end
+  rescue Errno::ETIMEDOUT, Errno::EPERM
+    false
+  rescue Errno::ECONNREFUSED, Errno::EHOSTUNREACH, Errno::ENETUNREACH
+    sleep 2
+    false
+  ensure
+    tcp_socket && tcp_socket.close
   end
 
   def self.exe(*args)
     args.flatten!
     if args.size > 1
-      @host = args[0]
-      if @host =~ /@/
-        @user, @host = @host.split '@'
-      end
-      @cmd = args[1]
+      self.host = args[0]
+      self.user, self.host = self.host.split '@' if self.host =~ /@/
+      self.cmd = args[1]
     else
-      @cmd = args[0]
+      self.cmd = args[0]
     end
-    Net::SSH.start(host, user, flags) do |sh|
+    self.flags[:password] = self.password if self.password
+    self.flags[:keys] = [self.key] if self.key
+    Net::SSH.start(self.host, self.user, self.flags) do |sh|
       sh.open_channel do |ch| 
         if sudo?
           ch.request_pty 
-          @cmd = "sudo #{@cmd}"
+          if cmd =~ /(&&|\|\||&|\|)/
+            self.cmd = cmd.gsub(/(&&|\|\||&|\|)/, $1 + 'sudo')
+          end
+          self.cmd = "sudo #{cmd}"
+          puts "SHEXY: #{cmd}" if $DEBUG
         end
         ch.exec cmd do
           # FIXME: I don't think it's a good idea
           # to implement access to stdout,stderr this way
           stdout = ""
@@ -93,10 +129,29 @@
         end
       end
     end
   end
 
+  def self.batch(&block)
+    require 'tempfile'
+    out, err = nil,nil
+    def self.script(script)
+      f = Tempfile.new 'shexy'
+      begin
+        f.puts script 
+        f.flush
+        copy_to f.path, "#{f.path}.remote"
+        out, err = exe "/bin/bash #{f.path}.remote" 
+      ensure
+        f.close
+        f.unlink
+      end
+      return out, err
+    end
+    instance_eval &block
+  end
+
   #
   # Shexy.copy_to '[e-mail]', 'source_file', 'dest_file'
   #
   # or
   #
@@ -111,22 +166,23 @@
       args.delete :recursive
     end
     
     if args.size > 2
       # First arg assumed to be [e-mail]
-      @host = args[0]
-      if @host =~ /@/
-        @user, @host = @host.split '@'
+      self.host = args[0]
+      if self.host =~ /@/
+        self.user, self.host = self.host.split '@'
       end
       from = args[1]
       to = args[2]
     else
       # user, host already set via Shexy.host and
       # Shexy.user
       from = args[0]
       to = args[1]
     end
+    self.flags[:password] = self.password if self.password
     from = File.expand_path from
     Net::SCP.start(host, user, flags) do |scp|
       scp.upload! from, to, opts
     end
   end
@@ -161,8 +217,42 @@
   #
   def self.issue
     issue, err = ((exe 'cat /etc/issue')[0].strip.chomp || Shexy.exe('lsb_release -i')[0].strip.chomp).lines.first
     raise Exception.new "Error reading release info." unless (err.nil? or err.empty?)
     issue
+  end
+
+  def self.copy_ssh_pubkey(path, dest_dir = '~/.ssh')
+    path = File.expand_path path
+    raise ArgumentError.new("Invalid key file") unless File.exist?(path)
+    key = File.read path
+    batch do
+      script <<-EOH
+      mkdir -p #{dest_dir} && chmod 700 #{dest_dir}
+      echo '#{key}' >> #{dest_dir}/authorized_keys
+      EOH
+    end
+  end
+
+  #
+  # Set PermitRootLogin to value in /etc/ssh/sshd_config
+  # value accepted: yes,now, without-password
+  #
+  def self.permit_root_login(value)
+    value = value.to_s
+    unless value =~ /^(yes|no|without-password)$/
+      raise ArgumentError.new "Argument should be yes|no|without-password"
+    end
+    using_sudo = Shexy.sudo?
+    Shexy.use_sudo 
+    out, err = batch do
+      script <<-EOH
+      sed -i 's/^#\?PermitRootLogin.*$/PermitRootLogin\ #{value}/' /etc/ssh/sshd_config
+      test -f /etc/init.d/ssh && /etc/init.d/ssh restart
+      test -f /etc/init.d/sshd && /etc/init.d/sshd restart
+      EOH
+    end
+    Shexy.use_sudo(false) unless using_sudo
+    return out, err
   end
 
 end