robot-army 0.1.8

data/lib/robot-army/eval_command.rb

@@ -0,0 +1,17 @@
+class RobotArmy::EvalCommand
+  def initialize
+    yield self if block_given?
+  end
+
+  attr_accessor :user
+
+  attr_accessor :proc
+
+  attr_accessor :args
+
+  attr_accessor :context
+
+  attr_accessor :dependencies
+
+  attr_accessor :keychain
+end

data/lib/robot-army/gate_keeper.rb

@@ -0,0 +1,28 @@
+class RobotArmy::GateKeeper
+  attr_reader :connection_class
+  
+  def initialize(connection_class=RobotArmy::OfficerConnection)
+    @connection_class = connection_class
+  end
+  
+  def connect(host)
+    connections[host] ||= establish_connection(host)
+  end
+  
+  def establish_connection(host)
+    connection = connections[host] = connection_class.new(host)
+    connection.open
+  end
+  
+  def connections
+    @connections ||= {}
+  end
+  
+  def close
+    connections.each { |host,c| c.close unless c.closed? }
+  end
+  
+  def self.shared_instance
+    @shared_instance ||= new
+  end
+end

data/lib/robot-army/io.rb

@@ -0,0 +1,106 @@
+class RobotArmy::IO
+  attr_reader :name
+  
+  # Starts capturing output of the named stream.
+  # 
+  def start_capture
+    eval "$#{name} = self"
+  end
+  
+  # Stops capturing output of the named stream.
+  # 
+  def stop_capture
+    eval "$#{name} = #{name.upcase}"
+  end
+  
+  def puts(*args) #:nodoc:
+    post capture(:puts, *args)
+  end
+  
+  def print(*args) #:nodoc:
+    post capture(:print, *args)
+  end
+  
+  def write(*args) #:nodoc:
+    post capture(:write, *args)
+  end
+  
+private
+  
+  def initialize(name)
+    @name = name.to_s
+    start_capture
+  end
+  
+  def post(string)
+    RobotArmy.upstream.post(:status => 'output', :data => {:stream => name, :string => string})
+  end
+  
+  def capture(*call)
+    stream = StringIO.new
+    stream.send(*call)
+    stream.string
+  end
+  
+  class <<self
+    # Determines whether the given stream has any data to be read.
+    # 
+    # @param stream [IO]
+    #   The +IO+ stream to check.
+    # 
+    # @return [Boolean]
+    #   +true+ if stream has data to be read, +false+ otherwise.
+    # 
+    def has_data?(stream)
+      selected, _ = IO.select([stream], nil, nil, 0.5)
+      return selected && !selected.empty?
+    end
+
+    # Reads immediately available data from the given stream.
+    # 
+    #   # echo foo | ruby test.rb
+    #   RobotArmy::IO.read_data($stdin) # => "foo\n"
+    # 
+    # @param stream [IO]
+    #   The +IO+ stream to read from.
+    # 
+    # @return [String]
+    #   The data read from the stream.
+    # 
+    def read_data(stream)
+      data = []
+      data << stream.readpartial(1024) while has_data?(stream)
+      return data.join
+    end
+    
+    # Redirects the named stream to a +StringIO+.
+    # 
+    #   RobotArmy::IO.capture(:stdout) { puts "foo" } # => "foo\n"
+    # 
+    #   RobotArmy::IO.silence(:stderr) { system "rm non-existent-file" }
+    # 
+    # @param stream [Symbol]
+    #   The name of the stream to redirect (i.e. +:stderr+, +:stdout+).
+    # 
+    # @yield
+    #   The block whose output we should capture.
+    # 
+    # @return [String]
+    #   The string result of the output produced by the block.
+    # 
+    def capture(stream)
+      begin
+        stream = stream.to_s
+        eval "$#{stream} = StringIO.new"
+        yield
+        result = eval("$#{stream}").string
+      ensure 
+        eval("$#{stream} = #{stream.upcase}")
+      end
+
+      result
+    end
+    
+    alias_method :silence, :capture
+  end
+end

data/lib/robot-army/keychain.rb

@@ -0,0 +1,10 @@
+class RobotArmy::Keychain
+  def get_password_for_user_on_host(user, host)
+    read_with_prompt("[sudo] password for #{user}@#{host}: ")
+  end
+
+  def read_with_prompt(prompt)
+    require 'highline'
+    HighLine.new.ask(prompt) {|q| q.echo = false}
+  end
+end

data/lib/robot-army/loader.rb

@@ -0,0 +1,85 @@
+class RobotArmy::Loader
+  attr_accessor :messenger
+  
+  def libraries
+    @libraries ||= []
+  end
+  
+  def render
+    %{
+      begin
+        ##
+        ## setup
+        ##
+        
+        $TESTING = #{$TESTING.inspect}
+        $stdout.sync = $stdin.sync = true
+        #{libraries.map{|l| "require #{l.inspect}"}.join("\n")}
+        
+        
+        ##
+        ## local Robot Army objects to communicate with the parent
+        ##
+        
+        loader = #{self.class.name}.new
+        RobotArmy.upstream = RobotArmy::Messenger.new($stdin, $stdout)
+        loader.messenger = RobotArmy.upstream
+        loader.messenger.post(:status => 'ok')
+        
+        ##
+        ## event loop
+        ##
+        
+        loader.load
+      rescue Object => e
+        ##
+        ## exception handler of last resort
+        ##
+        
+        if defined?(RobotArmy::Exit) && e.is_a?(RobotArmy::Exit)
+          # don't stomp on our own "let me out" exception
+          exit(e.status)
+        else
+          # if we got here that means something up to and including loader.load
+          # went unexpectedly wrong. this could be a missing library, or it 
+          # could be a bug in Robot Army. either way we should report the error
+          # back to the place we came from so that they may re-raise the exception
+          
+          # a little bit of un-DRY
+          print Base64.encode64(Marshal.dump(:status => 'error', :data => e))+'|'
+          exit(1)
+        end
+      end
+    }
+  end
+  
+  def safely
+    begin
+      return yield, true
+    rescue RobotArmy::Exit
+      # let RobotArmy::Exit through
+      raise
+    rescue Object => e
+      messenger.post(:status => 'error', :data => e)
+      return nil, false
+    end
+  end
+  
+  def safely_or_die(&block)
+    retval, success = safely(&block)
+    exit(1) unless success
+    return retval
+  end
+  
+  def load
+    # create a soldier
+    soldier = safely_or_die{ RobotArmy::Soldier.new(messenger) }
+    
+    # use the soldier to start listening to incoming commands
+    # at this point everything has been loaded successfully, so we
+    # don't have to exit if an exception is thrown
+    loop do
+      safely{ soldier.listen }
+    end
+  end
+end

data/lib/robot-army/marshal_ext.rb

@@ -0,0 +1,52 @@
+module Marshal
+  class <<self
+    # Determines whether a given object can be dumped.
+    #
+    # @param object Object
+    #   The object to check.
+    #
+    # @return [Boolean]
+    #   +true+ if dumping the object does not raise an error,
+    #   +false+ if a +TypeError+ is raised.
+    #
+    # @raise Exception
+    #   Whatever +Marshal.dump+ might raise that isn't a +TypeError+.
+    #
+    def can_dump?(object)
+      begin
+        Marshal.dump(object)
+        return true
+      rescue TypeError
+        return false
+      end
+    end
+  end
+end
+
+class Object
+  # Syntactic sugar for +Marshal.can_dump?+.
+  #
+  # @see Marshal.can_dump?
+  def marshalable?
+    Marshal.can_dump?(self) &&
+    instance_variables.all? do |ivar|
+      instance_variable_get(ivar).marshalable?
+    end
+  end
+end
+
+class Array
+  def marshalable?
+    super &&
+    all? {|item| item.marshalable?}
+  end
+end
+
+class Hash
+  def marshalable?
+    super &&
+    keys.marshalable? &&
+    values.marshalable? &&
+    default.marshalable?
+  end
+end

data/lib/robot-army/messenger.rb

@@ -0,0 +1,31 @@
+module RobotArmy
+  class Messenger
+    attr_reader :input, :output
+    
+    def initialize(input, output)
+      @input, @output = input, output
+    end
+    
+    def post(response)
+      dump = Marshal.dump(response)
+      dump = Base64.encode64(dump) + '|'
+      output << dump
+    end
+    
+    def get
+      data = nil
+      loop do
+        case data = input.gets('|')
+        when nil, ''
+          return nil
+        when /^\s*$/
+          # again!
+        else
+          break
+        end
+      end
+      data = Base64.decode64(data.chop)
+      Marshal.load(data)
+    end
+  end
+end

data/lib/robot-army/officer.rb

@@ -0,0 +1,35 @@
+class RobotArmy::Officer < RobotArmy::Soldier
+  def run(command, data)
+    case command
+    when :eval
+      debug "officer delegating eval command for user=#{data[:user].inspect}"
+      RobotArmy::Connection.localhost(data[:user], proc{ ask_for_password(data[:user]) }) do |local|
+        local.post(:command => command, :data => data)
+        
+        loop do
+          # we want to stay in this loop as long as we 
+          # have proxy requests coming back from our child
+          response = local.get
+          case response[:status]
+          when 'proxy'
+            # forward proxy requests on to our parent
+            messenger.post(response)
+            # and send the response back to our child
+            local.post(messenger.get)
+          else
+            return RobotArmy::Connection.handle_response(response)
+          end
+        end
+      end
+    when :exit
+      super
+    else
+      super
+    end
+  end
+  
+  def ask_for_password(user)
+    messenger.post(:status => 'password', :data => {:as => user, :user => ENV['USER']})
+    RobotArmy::Connection.handle_response messenger.get
+  end
+end

data/lib/robot-army/officer_connection.rb

@@ -0,0 +1,5 @@
+class RobotArmy::OfficerConnection < RobotArmy::Connection
+  def loader
+    @loader ||= RobotArmy::OfficerLoader.new
+  end
+end

data/lib/robot-army/officer_loader.rb

@@ -0,0 +1,13 @@
+class RobotArmy::OfficerLoader < RobotArmy::Loader
+  def load
+    # create a soldier
+    soldier = safely_or_die{ RobotArmy::Officer.new(messenger) }
+    
+    # use the soldier to start listening to incoming commands
+    # at this point everything has been loaded successfully, so we
+    # don't have to exit if an exception is thrown
+    loop do
+      safely{ soldier.listen }
+    end
+  end
+end

data/lib/robot-army/proxy.rb

@@ -0,0 +1,35 @@
+class RobotArmy::Proxy
+  alias_method :__proxy_instance_eval, :instance_eval
+  instance_methods.each { |m| undef_method m unless m =~ /^__/ }
+
+  def initialize(messenger, hash)
+    @messenger = messenger
+    @hash = hash
+  end
+
+  def sh(binary, *args)
+    command = [binary, *args].join(' ')
+    output = %x{#{command} 2>&1}
+
+    if not $?.success?
+      raise RobotArmy::ShellCommandError.new(command, $?.exitstatus, output)
+    end
+  end
+
+  def self.generator_for(object)
+    "RobotArmy::Proxy.new(RobotArmy.upstream, #{object.hash.inspect})"
+  end
+
+private
+
+  def method_missing(*args, &block)
+    @messenger.post(:status => 'proxy', :data => {:hash => @hash, :call => args})
+    response = @messenger.get
+    case response[:status]
+    when 'proxy'
+      return RobotArmy::Proxy.new(@messenger, response[:data])
+    else
+      return RobotArmy::Connection.handle_response(response)
+    end
+  end
+end

data/lib/robot-army/remote_evaler.rb

@@ -0,0 +1,59 @@
+class RobotArmy::RemoteEvaler
+  attr_reader :connection, :command, :options, :proxies
+
+  def initialize(connection, command)
+    @connection = connection
+    @command = command
+  end
+
+  def execute_command
+    @options, @proxies = RobotArmy::EvalBuilder.build(command)
+    send_eval_command
+    return loop_until_done
+  end
+
+  private
+
+  def send_eval_command
+    debug("Evaling code remotely:\n#{options[:code]}")
+    connection.post(:command => :eval, :data => options)
+  end
+
+  def loop_until_done
+    catch :done do
+      loop { process_response(connection.messenger.get) }
+    end
+  end
+
+  def process_response(response)
+    case response[:status]
+    when 'proxy'
+      handle_proxy_response(response)
+    when 'password'
+      debug("Got password request: #{response[:data].inspect}")
+      connection.post :status => 'ok', :data => command.keychain.get_password_for_user_on_host(response[:data][:user], connection.host)
+    else
+      begin
+        throw :done, connection.handle_response(response)
+      rescue RobotArmy::Warning => e
+        $stderr.puts "WARNING: #{e.message}"
+        throw :done, nil
+      end
+    end
+  end
+
+  def handle_proxy_response(response)
+    begin
+      proxy = proxies[response[:data][:hash]]
+      data = proxy.send(*response[:data][:call])
+      if data.marshalable?
+        connection.post :status => 'ok', :data => data
+      else
+        proxies[data.hash] = data
+        connection.post :status => 'proxy', :data => data.hash
+      end
+    rescue Object => e
+      connection.post :status => 'error', :data => e
+    end
+  end
+end