cluster_bomb 0.2.6

data/lib/cluster_bomb/cli.rb

@@ -0,0 +1,82 @@
+require 'cluster_bomb/bomb'
+require 'cluster_bomb/logging'
+require 'getoptlong'
+class Cli
+  include ClusterBomb::Bomb
+  def initialize
+    @execution_environment={}
+    super
+    self.load File.join(File.dirname(__FILE__),'stdtasks.rb')
+    if File.exists? 'Bombfile'
+      self.load 'Bombfile'
+    else
+      puts "WARNING: no Bombfile in current directory"
+    end
+  end
+  def process_args
+    if ARGV.length < 1
+      puts "syntax: cb <task> [options]"
+      puts "available tasks: "
+      @tasks.each do |n,t|
+        puts "   #{t.name} - #{t.description}"
+      end
+      exit
+    end
+    @task=ARGV[0].to_sym  
+    
+    # Get program opts
+    opts = GetoptLong.new(
+      [ '--logfile', '-l', GetoptLong::REQUIRED_ARGUMENT ],
+      [ '--logmode', GetoptLong::REQUIRED_ARGUMENT ],
+      [ '--user', GetoptLong::REQUIRED_ARGUMENT ],
+      [ '--nolog', GetoptLong::NO_ARGUMENT ]      
+    )
+    logfile=nil
+    logmode=nil
+    nolog=false
+    @user=nil
+    opts.each do |opt, arg|
+      case opt
+      when '--logfile'
+        logfile=arg
+      when '--logmode'
+        raise "valid log modes are a or w" unless ['a','w'].include? arg
+        logmode=arg
+      when '--nolog'
+        nolog=true
+      when '--user'
+        @user=arg        
+      end      
+    end    
+    if logfile && !logmode
+      logmode='a'
+    end
+    unless nolog
+      ClusterBomb::Logging.log_init
+      ClusterBomb::Logging.log_enable(logfile,logmode)
+    end
+    # Set up environment, with special attention to roles
+    ARGV[1..-1].each do |arg|
+      pair = arg.split('=')
+      if pair.length == 2
+        k = pair[0].strip.to_sym
+        if k == :roles
+          list=pair[1].split(',').collect{|r|r.strip.to_sym}
+          @execution_environment[:roles]=list
+        elsif k == :hosts
+          list=pair[1].split(',').collect{|r|r.strip}
+          @execution_environment[:hosts]=list
+        else
+          self.set(k,pair[1].strip)
+        end
+      end
+    end  
+  end
+  def main
+    process_args
+    switch_user(@user) if @user   
+    exec @task, {:roles=>@execution_environment[:roles], :hosts=>@execution_environment[:hosts]}
+  end
+end
+
+Cli.new.main

data/lib/cluster_bomb/cluster.rb

@@ -0,0 +1,298 @@
+require 'net/ssh'
+require 'net/scp'
+require 'net/sftp'
+require 'cluster_bomb/logging'
+
+# Represents a cluster of hosts over which to operate
+module ClusterBomb
+  class Cluster
+    include Logging
+    class Host
+      attr_accessor :name
+      attr_accessor :buffer_stderr, :buffer_stdout, :buffer_console
+      attr_accessor :exception
+      attr_accessor :data
+      attr_accessor :connected
+      attr_accessor :connect_failed
+      
+      def initialize(name, cluster)
+        self.name=name
+        self.clear!
+        self.data={}
+        self.connected=false
+        self.connect_failed=false
+        @cluster = cluster
+      end
+      def nickname
+        @cluster.nicknames[self.name]
+      end
+      def stdout
+        buffer_stdout.join('')
+      end
+      def stderr
+        buffer_stderr.join('')
+      end    
+      def console
+        buffer_console.join('')
+      end
+      def clear!
+        self.buffer_stderr=[]
+        self.buffer_stdout=[]
+        self.buffer_console=[]
+        self.exception=nil
+      end
+      # match stdout, returning arrays
+      def match(rex, default=nil)
+        m = rex.match self.stdout
+        if m
+          if m.length > 2
+            m[1..-1]
+          else
+            m[1]
+          end
+        else
+          default
+        end
+      end
+    end
+      
+    attr_accessor :hosts, :nicknames
+    def initialize(user, options={})
+      @user_name ||= user
+      @connections=[]
+      @hosts=[]
+      @ssh_options=options
+      @connection_mutex = Mutex.new
+      @connected=false
+      @connection_cache={}
+      @nicknames={}
+      @start_time =nil
+      @max_time=nil
+    end
+    
+    def connect!(host_list)
+      return if host_list.empty?
+      @hosts=[]
+      # Build up results rray
+      host_list.each {|hostname| @hosts << Host.new(hostname, self)}
+    
+      # Connect. Build up connections array
+      # Seems like there would be an async call to do this -- but it looks like
+      # Not -- so we resort to threads
+      puts "Connecting to #{hosts.length} hosts"
+      ensure_connected!
+      @connected=true    
+      puts "Connected to #{hosts.length} hosts"
+    end
+    
+    # Credentials to be used for next connection attempt. 
+    def credentials(user, ssh_opts)
+       @ssh_options=ssh_opts
+       @user_name = user
+    end
+    
+    # Be sure all hosts are connected that have not previously failed to connect
+    def ensure_connected!      
+      if @ssh_options[:timeout]
+        total_timeout = @ssh_options[:timeout] * 2
+      else
+        total_timeout = 30
+      end
+      # puts "Total timeout: #{total_timeout}"
+      @connections=[]
+      hosts_to_connect = @hosts.inject(0) {|sum,h| sum += (h.connect_failed ? 0:1)}
+      # puts "#{hosts_to_connect} to connect"
+      @hosts.each do |host|
+        if @connection_cache[host.name] || host.connected
+          @connection_mutex.synchronize { @connections << {:connection=>@connection_cache[host.name], :host=>host} }
+          host.connected=true
+        elsif !host.connect_failed
+          Thread.new {
+            begin
+              #puts "Connecting #{host.name}"          
+              c = Net::SSH.start(host.name, @user_name, @ssh_options)
+              @connection_cache[host.name] = c
+              @connection_mutex.synchronize { @connections << {:connection=>c, :host=>host} }
+              host.connected=true
+              #puts "Connected #{host.name}"          
+            rescue Exception => e
+              host.connect_failed = true
+              host.connected=false
+              error "Unable to connect to #{host.name}\n#{e.message}"
+              @connection_mutex.synchronize {@connections << {:connection=>nil, :host=>host} }
+              host.exception=e
+            end        
+          }
+        end
+      end
+      s = Time.now
+      loop do
+        l=0
+        @connection_mutex.synchronize { l = @connections.length }
+        break if l == hosts_to_connect
+        sleep(0.1)
+        if Time.now - s > total_timeout
+          puts "Warning -- total connection time expired"
+          puts "Failed to connect:"
+          hosts.each do |h|
+            unless h.connected
+              puts "    #{h.name}" 
+              h.connect_failed=true
+              # TODO: Need to handle this situations much better. Attempt to kill thread and/or mark connection in cache as unreachable
+            end
+          end
+          break
+        end
+      end      
+    end
+    
+    def set_run_timer(options)
+      if options[:max_run_time]
+        @max_time = options[:max_run_time].to_i
+        @start_time = Time.now
+      else
+        @max_time = nil
+        @start_time = nil
+      end
+    end
+    
+    def download(remote, local, options={}, &task)
+      ensure_connected!
+      set_run_timer(options)
+      @connections.each do |c|
+        next if c[:connection].nil?
+        c[:completed]=false      
+        c[:connection].scp.download(remote,local)
+      end
+      event_loop(task)
+      @hosts      
+    end
+    
+    def upload(local, remote, options={}, &task)
+      opts={:chunk_size=>16384}.merge(options)
+      ensure_connected!
+      set_run_timer(options)      
+      @connections.each do |c|
+        next if c[:connection].nil?
+        c[:completed]=false      
+        c[:connection].scp.upload(local,remote, opts)
+      end
+      event_loop(task)
+      @hosts      
+    end
+       
+    # Build sudo-fied command. Really only works for bash afaik   
+    def mksudo(command)
+      "sudo sh -c '(#{command})'"
+    end 
+    
+    def run(command, options={}, &task)
+      # Execute
+      ensure_connected!
+      if options[:sudo]
+        command=mksudo(command)
+      end      
+      # puts "Command: #{command}"
+      set_run_timer(options)      
+      @connections.each do |c|
+        next if c[:connection].nil?
+        c[:completed]=false
+        c[:host].clear! if options[:echo] # Clear out before starting
+        c[:connection].exec command do |ch, stream, data|
+          c[:host].buffer_console << data
+          if stream == :stderr
+            c[:host].buffer_stderr << data
+          else
+            "#{c[:host].name}=> #{data}"
+            c[:host].buffer_stdout << data
+          end
+          puts "#{c[:host].name}::#{data}" if options[:debug]
+          print "." if options[:dotty]
+        end
+      end
+      event_loop(task, options)
+      @hosts
+    end
+    
+    def event_loop(task, options={})
+      # Event loop
+      condition = Proc.new { |s| s.busy?(true) }
+      # Count up non-nil connections
+      count = 0
+      @connections.each {|c| count +=1 if c[:connection]}
+      loop do
+        @connections.each do |conn|
+          next if conn[:connection].nil? || conn[:completed]
+          ex=nil
+          busy=true
+          begin
+            busy = conn[:connection].process(0.1, &condition)
+            if @start_time && Time.now - @start_time > @max_time
+              # Soft exception here -- stay connected
+              conn[:host].exception = Exception.new("Execution time exceeded: #{@max_time}")
+              puts "Execution time exceeded: #{@max_time}"
+              busy=false
+            end
+          rescue Exception => e
+            # As far as I can tell, if we ever get here, the session is fucked. 
+            # Close out the connection and indicate that we want to be reconnected later
+            # In general, its upload/download exceptions that get us here. Even bad filenames can do the trick
+            puts "#{e.message}"
+            host = conn[:host]
+            @connection_cache[host.name].close
+            @connection_cache[host.name]=nil
+            host.connected=false # disconnect
+            busy=false
+          end
+          if !busy
+            conn[:completed] = true
+            count -=1 
+            h = conn[:host]
+            if task
+              task.call(h)
+            elsif options[:echo]
+              puts "#{h.name}\n#{h.console}\n"
+            end
+          end
+        end
+        break if count <=0
+      end
+      # Reset these
+      @start_time=nil
+      @max_time = nil      
+    end
+  
+    def connected?
+      @connected
+    end
+  
+    def disconnect!
+      if @connections
+        @connection_cache.each do |k, conn|
+          begin
+            conn.close
+          rescue Exception=>e
+            puts "Non-fatal EXCEPTION closing connection: #{e.message}"
+          end
+        end
+      end
+      @hosts.each {|h| h.connected=false}
+      @connections=[]
+      @connection_cache={}
+      @connected=false
+      @hosts=[]
+    end
+    
+    def reset!
+      @connections=[]
+      @connected=false  
+    end
+  
+    def clear!
+      @hosts.each {|h| h.clear!}
+    end
+    def error(msg)
+      puts msg
+    end
+  end
+end

data/lib/cluster_bomb/configuration.rb

@@ -0,0 +1,83 @@
+require 'yaml'
+module ClusterBomb
+  module Configuration
+    HOME=File.join(ENV["HOME"], ".cluster_bomb")
+    CFG_FILE_NAME = File.join(HOME,"config.yaml")
+    KEYS={
+      :screen_width=>{:default=>120},
+      :logging=>{:default=>true},
+      :logfile=>{:default=>"logs/cb-#{Time.now().strftime('%m-%d-%Y')}.log}"},
+      :max_run_time=>{:default=>nil},
+      :max_history=>{:default=>1000},
+      :username=>{:default=>ENV["USER"]}
+    }
+    def initialize
+      @configuration={}
+      @ssh_options_for_user={} # By user
+      Configuration.check_dir
+    end
+    # Check for home dir, and create it if it doesn't exist
+    def self.check_dir
+      `mkdir #{HOME}` unless File.exists? HOME
+    end
+    # Save it out
+    def save
+      Configuration.check_dir
+      File.open(CFG_FILE_NAME,"w") {|f| f.write(@configuration.to_yaml)}
+    end
+    def configuration
+      @configuration
+    end  
+    def set(key, val)
+      k = key.class==Symbol ? key : key.to_sym
+      # TODO: VALIDATION
+      raise "Invalid setting [#{key}]" unless KEYS[k]
+      @configuration[k] = val
+    end
+    def get(key)
+      k = key.class==Symbol ? key : key.to_sym
+      raise "Config.get ==> Uknown key #{k}" unless (@configuration.has_key? k or KEYS.has_key? k)
+      ret = @configuration[k] || KEYS[k][:default]
+      ret
+    end
+    def keys
+      KEYS.keys
+    end
+    def valid_key?(k)
+      KEYS.has_key? k.to_sym
+    end
+    def load!
+      Configuration.check_dir
+      begin
+        buf = File.read(CFG_FILE_NAME)
+      rescue
+        `touch #{CFG_FILE_NAME}`
+        buf = File.read(CFG_FILE_NAME)
+      end
+      unless buf.empty?
+        @configuration = YAML.load(buf)
+      end
+    end
+    def ssh_options(username)
+      if has_ssh_options? username
+        if !@ssh_options_for_user[username]
+          @ssh_options_for_user[username] = YAML.load(File.read(File.join(HOME,"ssh_#{username}.yml")))
+        end
+      end
+      @ssh_options_for_user[username] || {}
+    end
+    def has_ssh_options?(username)
+      return true if  @ssh_options_for_user[username] 
+      File.exists? File.join(HOME,"ssh_#{username}.yml")
+    end    
+    def method_missing(symbol, *args)
+      str = symbol.to_s
+      if str.match(/=$/)
+        self.set(str.sub('=', ''),args[0])
+      else
+        self.get(str)
+      end
+    end
+  end
+end
+

data/lib/cluster_bomb/dispatcher.rb

@@ -0,0 +1,110 @@
+require 'readline'
+module ClusterBomb
+  module Dispatcher
+    COMMANDS = [
+      {:rex=> /^qu.*/, :method=>:quit, :name=>':quit', :description=>'Quit. Ctrl-d also works'},
+      {:rex=> /^wi.*/, :method=>:with, :name=>':with', :description=>'Set roles in use. This will determine the remote host set. ex: :with apache,database'},
+      {:rex=> /^use.*/, :method=>:use, :name=>':use', :description=>'Run a command against a host or hosts. Ex. use foo.bar.com,x.y.com ls -la'},
+      {:rex=> /^ex.*/, :method=>:exec, :name=>':exec', :description=>'Execute a configured task task'},
+      {:rex=> /^li.*/, :method=>:list, :name=>':list', :description=>'List available tasks to run'},
+      {:rex=> /^his.*/, :method=>:history, :name=>':history', :description=>'Command history. Can be followed by a filter regexp'},      
+      {:rex=> /^disc.*/, :method=>:disconnect, :name=>':disconnect', :description=>'Disconnect all cached connections.'},
+      {:rex=> /^up.*/, :method=>:upload, :name=>':upload', :description=>'Upload one or more files to servers. Supports wildcards and autocomplete for source filename. ex: upload sourcepath destpath'},          
+      {:rex=> /^he.*/, :method=>:help, :name=>':help', :description=>'Quick help.'},    
+      {:rex=> /^se.*/, :method=>:set, :name=>':set', :description=>'Set a variable'},          
+      {:rex=> /^ho.*/, :method=>:host_list, :name=>':hosts', :description=>'List current hosts'},
+      {:rex=> /^switch/, :method=>:switch, :name=>':switch', :description=>'Switch user'},
+      {:rex=> /^sudo/, :method=>:sudo, :name=>':sudo', :description=>'Sudo mode on/off'}      
+      ]
+    COMMAND_AUTOCOMPLETE = COMMANDS.collect{|c|c[:name]} + ['with','use']
+    def init_autocomplete
+      Readline.completion_proc=proc {|s| self.dispatcher_completion_proc(s)}
+      Readline.completer_word_break_characters = 7.chr  
+      Readline.completion_case_fold = true
+      Readline.completion_append_character = ''      
+    end
+    def process_cmd(line)
+       # cmd = cmd.strip.downcase
+       m=line.match(/([^ ]+) *?(.*)/)
+       return true unless m
+       cmd = m[1]
+       params = m[2] || ""
+       params.strip!
+       return false if cmd =~ /^q.*/
+       found=false
+       COMMANDS.each do |cr|
+         if cmd =~ cr[:rex]
+           # p = cmd_param(cmd)
+           self.send cr[:method], params
+           found=true
+           break
+         end
+       end
+       if !found
+         puts "Available commands:"
+         COMMANDS.each do |cr|
+           puts "    #{cr[:name]} - #{cr[:description]}"
+         end         
+       end
+       return true
+    end
+    def dispatcher_completion_proc(s)      
+      # No line buffer for mac, so no way to get command context
+      # Very lame with libedit. Will not return candidates, and 
+      # We cannot get the current line so we can do context-based 
+      # edits
+      ret=[]
+      tokens = s.split(' ')
+      if s =~ /^:?\w.* $/
+        tokens << ''
+      elsif s =~ /^\\\w*/
+        tokens << ''
+      end
+      # Initial command only
+      if tokens.length <= 1 && tokens[0] != '\\'
+        ret = COMMAND_AUTOCOMPLETE.grep(/^#{s}/)
+      else
+        if tokens[0]=~/:?with.*/
+          ret = secondary_completion_proc(tokens, @bomb.role_list.collect{|r| r[:name].to_s})
+        elsif tokens[0]=~/:?use.*/ && @bomb.valid_role?(:all)
+          ret = secondary_completion_proc(tokens, @bomb.servers([:all]))
+        elsif tokens[0]=~/:exec.*/
+          ret = secondary_completion_proc(tokens, @bomb.task_list.collect{|t| t.name.to_s})
+        elsif tokens[0]=~/:upload.*/
+          ret = dir_completion_proc(tokens)
+        elsif tokens[0]=~/:set.*/
+          ret = secondary_completion_proc(tokens, @bomb.configuration.keys)
+        elsif tokens[0]=~/^\\/
+          ret = shawtie_names.grep(/#{tokens[0][1..-1]}/)
+          ret = ret.collect {|r|"\\#{r}"}
+        end
+        
+      end      
+      ret
+    end
+    
+    def dir_completion_proc(tokens)
+      choices=dir_list(tokens[1])
+      secondary_completion_proc(tokens,choices)
+    end
+    
+    def dir_list(token)
+      m = token.match(/(.*\/).*$/)
+      if m && m[1]
+        ret=Dir.glob("#{m[1]}*")
+      else
+        ret=Dir.glob("*")
+      end      
+      ret.collect {|p| (File.directory? p) ? "#{p}/" : "#{p}"}
+    end
+    
+    def secondary_completion_proc(tokens, choices)
+      if tokens[1]==''
+        choices
+      else
+        tokens[1] = tokens[1].gsub(/\./,'\.')
+        choices.grep(/^#{tokens[1]}/).collect {|c| "#{tokens[0]} #{c}"}
+      end
+    end
+  end # dispatcher
+end # clusterbomb

data/lib/cluster_bomb/history.rb

@@ -0,0 +1,42 @@
+require 'readline'
+module ClusterBomb
+  module History
+    HISTORY_FILE=File.join(Configuration::HOME,'history')
+    
+    def save_history
+      save_history =  Readline::HISTORY.to_a.dup
+      start=0
+      start = save_history.length - @max_history if save_history.length > @max_history
+      
+      File.open(HISTORY_FILE, "w") do |f|
+        save_history[start..-1].each {|l| f.puts(l)}
+      end
+    end
+    
+    def load_history(max_history=500)
+      @max_history = max_history      
+      return unless File.exists? HISTORY_FILE
+      File.open(HISTORY_FILE, "r") do |f|
+        c=0
+        while line = f.gets
+          Readline::HISTORY.push(line.strip)
+          if c==0 && libedit? # libedit work-around 
+            Readline::HISTORY.push(line.strip)
+          end
+          c+=1
+        end
+      end
+    end
+    # Cheesy check to see if libedit is in use -- will affect history
+    def libedit?
+      libedit = false
+      # If NotImplemented then this might be libedit
+      begin
+        Readline.emacs_editing_mode
+      rescue NotImplementedError
+        libedit = true
+      end
+      libedit
+    end    
+  end # History
+end # ClusterBomb

data/lib/cluster_bomb/logging.rb

@@ -0,0 +1,35 @@
+module ClusterBomb
+  module Logging
+    DEFAULT_LOGFILENAME="logs/cb-#{Time.now().strftime('%m-%d-%Y')}.log"
+    def self.log_init
+      @logging_enabled=false
+      @io=nil            
+    end
+    def self.log_enable(filename=DEFAULT_LOGFILENAME, filemode='a')
+      filename ||= DEFAULT_LOGFILENAME
+      filemode ||= 'a'
+      log_dir = File.dirname(filename)
+      `mkdir -p #{log_dir}` unless File.exists? log_dir
+      @io = File.open(filename,filemode)
+      @logging_enabled=true
+    end
+    def self.log_disable
+      @io.close if @io
+      @io=nil
+      @logging_enabled=false
+    end
+    def self.log(msg)
+      if @logging_enabled && @io
+        @io.puts "#{Time.now().strftime('%m-%d-%Y %H:%M:%S')} - #{msg}"
+        @io.flush
+      end      
+    end
+    def self.puts(msg)      
+      Kernel.puts msg
+      self.log(msg)
+    end
+    def puts(msg)
+      Logging.puts(msg)
+    end
+  end
+end # ClusterBomb

data/lib/cluster_bomb/roles.rb

@@ -0,0 +1,56 @@
+module Roles
+  def role(*args, &task)
+    if args.length==2
+      name=args[0]
+      servers=args[1]
+      ra=servers.split(',')       
+    elsif args.length==1 && task != nil
+      name=args[0]
+      ra = yield(task)
+    else
+      raise "role command takes a role name and a list of hosts OR a block yielding a list of hosts"
+    end    
+    set_role(name, ra)
+  end
+
+  # Clear out role
+  def clear_role(name)
+    rl = @roles[name.to_sym]
+    @roles[name.to_sym] =[] if rl
+  end
+  
+  # produce a list of servers from role list
+  # empty or nil role list implies all roles
+  def servers(role_list)
+    @roles ||={}
+    ra=[]
+    if role_list && !role_list.empty?
+      if role_list.class==String
+        ra = role_list.split(',').collect{|r| r.to_sym}
+      else
+        ra=role_list
+      end
+    end
+    server_list=[]
+    ra.each do |role|
+      raise "Role #{role} not found" unless @roles[role]
+      server_list += @roles[role]
+    end
+    server_list
+  end
+  
+  def valid_role?(name)
+    @roles[name.to_sym]
+  end
+  
+  private
+  def set_role(name, hosts)
+    @roles ||={}
+    @roles[name] ||=[]
+    # puts "ROLE: #{name} #{hosts.inspect}"
+    hosts.each do |host|
+      host.strip!
+      @roles[name] <<  host unless @roles[name].include? host
+    end
+  end
+end