robot-army 0.1.8

data/lib/robot-army/ruby2ruby_ext.rb

@@ -0,0 +1,19 @@
+class Proc
+  def arguments
+    (to_ruby[/\Aproc \{ \|([^\|]+)\|/, 1] || '').split(/\s*,\s*/)
+  end
+
+  def to_ruby_without_proc_wrapper
+    to_ruby[/\Aproc\s*\{\s*(\|[^\|]+\|)?\s*(.*?)\s*\}\Z/m, 2] || raise("Unable to parse proc's Ruby: #{to_ruby}")
+  end
+end
+
+class Method
+  def arguments
+    (to_ruby[/\A(def [^\s\(]+)(?:\(([^\)]*)\))?/, 2] || '').split(/\s*,\s*/)
+  end
+
+  def to_ruby_without_method_declaration
+    to_ruby[/\Adef [^\s\(]+(?:\([^\)]*\))?\s*(.*?)\s*end\Z/m, 1] || raise("Unable to parse method's Ruby: #{to_ruby}")
+  end
+end

data/lib/robot-army/soldier.rb

@@ -0,0 +1,37 @@
+class RobotArmy::Soldier
+  attr_reader :messenger
+  
+  def initialize(messenger)
+    @messenger = messenger
+  end
+  
+  def listen
+    request  = messenger.get
+    result   = run(request[:command], request[:data])
+    if result.marshalable?
+      response = {:status => 'ok', :data => result}
+    else
+      response = {
+        :status => 'warning', 
+        :data => "ignoring invalid remote return value #{result.inspect}"}
+    end
+    debug "#{self.class} post(#{response.inspect})"
+    messenger.post response
+  end
+  
+  def run(command, data)
+    debug "#{self.class} running command=#{command.inspect}"
+    case command
+    when :info
+      {:pid => Process.pid, :type => self.class.name}
+    when :eval
+      instance_eval(data[:code], data[:file], data[:line])
+    when :exit
+      # tell the parent we're okay before we exit
+      messenger.post(:status => 'ok')
+      raise RobotArmy::Exit
+    else
+      raise ArgumentError, "Unrecognized command #{command.inspect}"
+    end
+  end
+end

data/lib/robot-army/task_master.rb

@@ -0,0 +1,317 @@
+module RobotArmy
+  # The place where the magic happens
+  #
+  # ==== Types (shortcuts for use in this file)
+  # HostList:: <Array[String], String, nil>
+  class TaskMaster < Thor
+
+  no_tasks do
+
+    def initialize(*args)
+      super
+      @dep_loader = DependencyLoader.new
+    end
+
+    # Gets or sets a single host that instances of +RobotArmy::TaskMaster+ subclasses will use.
+    #
+    # @param host [String, :localhost]
+    #   The fully-qualified domain name to connect to.
+    #
+    # @return [String, :localhost]
+    #   The current value for the host.
+    #
+    # @raise RobotArmy::HostArityError
+    #   If you're using the getter form of this method and you've already
+    #   set multiple hosts, an error will be raised.
+    #
+    def self.host(host=nil)
+      if host
+        @hosts = nil
+        @host = host
+      elsif @hosts
+        raise RobotArmy::HostArityError,
+          "There are #{@hosts.size} hosts, so calling host doesn't make sense"
+      else
+        @host
+      end
+    end
+
+    # Gets or sets the hosts that instances of +RobotArmy::TaskMaster+ subclasses will use.
+    #
+    # @param hosts [Array[String]]
+    #   A list of fully-qualified domain names to connect to.
+    #
+    # @return [Array[String]]
+    #   The current list of hosts.
+    #
+    def self.hosts(hosts=nil)
+      if hosts
+        @host = nil
+        @hosts = hosts
+      elsif @host
+        [@host]
+      else
+        @hosts || []
+      end
+    end
+
+    # Gets the first host for this instance of +RobotArmy::TaskMaster+.
+    #
+    # @return [String, :localhost]
+    #   The host value to use.
+    #
+    # @raise RobotArmy::HostArityError
+    #   If you're using the getter form of this method and you've already
+    #   set multiple hosts, an error will be raised.
+    #
+    def host
+      if @host
+        @host
+      elsif @hosts
+        raise RobotArmy::HostArityError,
+          "There are #{@hosts.size} hosts, so calling host doesn't make sense"
+      else
+        self.class.host
+      end
+    end
+
+    # Sets a single host for this instance of +RobotArmy::TaskMaster+.
+    #
+    # @param host [String, :localhost]
+    #   The host value to use.
+    #
+    def host=(host)
+      @hosts = nil
+      @host = host
+    end
+
+    # Gets the hosts for the instance of +RobotArmy::TaskMaster+.
+    #
+    # @return [Array[String]]
+    #   A list of hosts.
+    #
+    def hosts
+      if @hosts
+        @hosts
+      elsif @host
+        [@host]
+      else
+        self.class.hosts
+      end
+    end
+
+    # Sets the hosts for this instance of +RobotArmy::TaskMaster+.
+    #
+    # @param hosts [Array[String]]
+    #   A list of hosts.
+    #
+    def hosts=(hosts)
+      @host = nil
+      @hosts = hosts
+    end
+
+    # Gets an open connection for the host this instance is configured to use.
+    #
+    # @return RobotArmy::Connection
+    #   An open connection with an active Ruby process.
+    #
+    def connection(host)
+      RobotArmy::GateKeeper.shared_instance.connect(host)
+    end
+
+    # Runs a block of Ruby on the machine specified by a host string as root
+    # and returns the return value of the block. Example:
+    #
+    #   sudo { `shutdown -r now` }
+    #
+    # You may also specify a user other than root. In this case +sudo+ is the
+    # same as +remote+:
+    #
+    #   sudo(:user => 'www-data') { `/etc/init.d/apache2 restart` }
+    #
+    # @param host [String, :localhost]
+    #   The fully-qualified domain name of the machine to connect to, or
+    #   +:localhost+ if you want to use the same machine.
+    #
+    # @options options
+    #   :user -> String => shell user
+    #
+    # @raise Exception
+    #   Whatever is raised by the block.
+    #
+    # @return [Object]
+    #   Whatever is returned by the block.
+    #
+    # @see remote
+    def sudo(hosts=self.hosts, options={}, &proc)
+      options, hosts = hosts, self.hosts if hosts.is_a?(Hash)
+      remote hosts, {:user => 'root'}.merge(options), &proc
+    end
+
+    # Runs a block of Ruby on the machine specified by a host string and
+    # returns the return value of the block. Example:
+    #
+    #   remote { "foo" } # => "foo"
+    #
+    # Local variables accessible from the block are also passed along to the
+    # remote process:
+    #
+    #   foo = "bar"
+    #   remote { foo } # => "bar"
+    #
+    # Objects which can't be marshalled, such as IO streams, will be proxied
+    # instead:
+    #
+    #   file = File.open("README.markdown", "r")
+    #   remote { file.gets } # => "Robot Army\n"
+    #
+    # @param hosts [HostList]
+    #   Which hosts to run the block on.
+    #
+    # @options options
+    #   :user -> String => shell user
+    #
+    # @raise Exception
+    #   Whatever is raised by the block.
+    #
+    # @return [Object]
+    #   Whatever is returned by the block.
+    #
+    def remote(hosts=self.hosts, options={}, &proc)
+      options, hosts = hosts, self.hosts if hosts.is_a?(Hash)
+      results = Array(hosts).map {|host| remote_eval({:host => host}.merge(options), &proc) }
+      results.size == 1 ? results.first : results
+    end
+
+    # Copies src to dest on each host.
+    #
+    # @param src [String]
+    #   A local file to copy.
+    #
+    # @param dest [String]
+    #   The path of a remote file to copy to.
+    #
+    # @raise Errno::EACCES
+    #   If the destination path cannot be written to.
+    #
+    # @raise Errno::ENOENT
+    #   If the source path cannot be read.
+    #
+    def scp(src, dest, hosts=self.hosts)
+      Array(hosts).each do |host|
+        output = `scp -q #{src} #{"#{host}:" unless host == :localhost}#{dest} 2>&1`
+        case output
+        when /Permission denied/i
+          raise Errno::EACCES, output.chomp
+        when /No such file or directory/i
+          raise Errno::ENOENT, output.chomp
+        end unless $?.exitstatus == 0
+      end
+
+      return nil
+    end
+
+    # Copies path to a temporary directory on each host.
+    #
+    # @param path [String]
+    #   A local file to copy.
+    #
+    # @param hosts [HostList]
+    #   Which hosts to connect to.
+    #
+    # @yield [path]
+    #   Yields the path of the newly copied file on each remote host.
+    #
+    # @yieldparam [String] path
+    #   The path of the file under in a new directory under a
+    #   temporary directory on the remote host.
+    #
+    # @return [Array<String>]
+    #   An array of destination paths.
+    #
+    def cptemp(path, hosts=self.hosts, options={}, &block)
+      hosts, options = self.hosts, hosts if hosts.is_a?(Hash)
+
+      results = remote(hosts, options) do
+        File.join(%x{mktemp -d -t robot-army.XXXX}.chomp, File.basename(path))
+      end
+
+      me = ENV['USER']
+      host_and_path = Array(hosts).zip(Array(results))
+      # copy them over
+      host_and_path.each do |host, tmp|
+        sudo(host) { FileUtils.chown(me, nil, File.dirname(tmp)) } if options[:user]
+        scp path, tmp, host
+        sudo(host) { FileUtils.chown(options[:user], nil, File.dirname(tmp)) } if options[:user]
+      end
+      # call the block on each host
+      results = host_and_path.map do |host, tmp|
+        remote(host, options.merge(:args => [tmp]), &block)
+      end if block
+
+      # delete it when we're done
+      RobotArmy::AtExit.shared_instance.at_exit do
+        host_and_path.each do |host, tmp|
+          remote(host, options) do
+            FileUtils.rm_rf(tmp)
+          end
+        end
+      end
+
+      results.size == 1 ? results.first : results
+    end
+
+    # Add a gem dependency this TaskMaster checks for on each remote host.
+    #
+    # @param dep [String]
+    #   The name of the gem to check for.
+    #
+    # @param ver [String]
+    #   The version string of the gem to check for.
+    #
+    def dependency(dep, ver = nil)
+      @dep_loader.add_dependency dep, ver
+    end
+
+  private
+
+    def say(something)
+      something = HighLine.new.color(something, :bold) if defined?(HighLine)
+      puts "** #{something}"
+    end
+
+    # Handles remotely eval'ing a Ruby Proc.
+    #
+    # @options options
+    #   :host -> [String, :localhost] => remote host
+    #   :user -> String => shell user
+    #   :password -> [String, nil] => sudo password
+    #
+    # @return Object
+    #   Whatever the block returns.
+    #
+    # @raise Exception
+    #   Whatever the block raises.
+    #
+    def remote_eval(options, &proc)
+      evaler = RemoteEvaler.new(connection(options[:host]), EvalCommand.new do |command|
+        command.user = options[:user]
+        command.proc = proc
+        command.args = options[:args] || []
+        command.context = self
+        command.dependencies = @dep_loader
+        command.keychain = keychain
+      end)
+
+      return evaler.execute_command
+    end
+
+    # Returns the default password manager. This should be
+    # overridden if you wish to have different sudo behavior.
+    def keychain
+      @keychain ||= Keychain.new
+    end
+  end
+
+  end
+end

data/spec/at_exit_spec.rb

@@ -0,0 +1,25 @@
+require File.dirname(__FILE__) + '/spec_helper'
+
+describe RobotArmy::AtExit do
+  before do
+    @at_exit = RobotArmy::AtExit.shared_instance
+  end
+
+  it "runs the provided block when directed" do
+    foo = 'foo'
+    @at_exit.at_exit { foo = 'bar' }
+    foo.should == 'foo'
+    @at_exit.do_exit
+    foo.should == 'bar'
+  end
+
+  it "does not run the same block twice" do
+    foo = 0
+    @at_exit.at_exit { foo += 1 }
+    foo.should == 0
+    @at_exit.do_exit
+    foo.should == 1
+    @at_exit.do_exit
+    foo.should == 1
+  end
+end

data/spec/connection_spec.rb

@@ -0,0 +1,126 @@
+require File.dirname(__FILE__) + '/spec_helper'
+
+describe RobotArmy::Connection do
+  before do
+    # given
+    @host = 'example.com'
+    @connection = RobotArmy::Connection.new(@host)
+    @messenger  = mock(:messenger)
+    @messenger.stub!(:post)
+    @connection.stub!(:messenger).and_return(@messenger)
+    @connection.stub!(:start_child)
+  end
+
+  it "is not closed after opening" do
+    # when
+    @connection.open
+
+    # then
+    @connection.should_not be_closed
+  end
+
+  it "returns itself from open" do
+    @connection.open.should == @connection
+  end
+
+  it "returns the result of the block passed to open" do
+    # when
+    @connection.stub!(:close)
+
+    # then
+    @connection.open{ 3 }.should == 3
+  end
+
+  it "closes the connection if a block is passed to open" do
+    # then
+    proc{ @connection.open{ 3 } }.
+      should_not change(@connection, :closed?).from(true)
+  end
+
+  it "closes the connection even if an exception is raised in the block passed to open" do
+    # then
+    proc do
+      proc{ @connection.open{ raise 'BOO!' } }.
+        should_not change(@connection, :closed?).from(true)
+    end.
+      should raise_error('BOO!')
+  end
+
+  it "does not start another child process if we're already open" do
+    # then
+    @connection.should_not_receive(:start_child)
+
+    # when
+    @connection.stub!(:closed?).and_return(false)
+    @connection.open
+  end
+
+  it "raises an exception when calling close if a connection is already closed" do
+    # when
+    @connection.stub!(:closed?).and_return(true)
+
+    # then
+    proc{ @connection.close }.should raise_error(RobotArmy::ConnectionNotOpen)
+  end
+
+  it "sends an exit command to its child upon closing" do
+    # then
+    @messenger.should_receive(:post).with(:command => :exit)
+
+    # when
+    @connection.stub!(:closed?).and_return(false)
+    @connection.close
+  end
+
+  it "starts a closed local connection when calling localhost" do
+    RobotArmy::Connection.localhost.should be_closed
+  end
+
+  it "raises a Warning when handling a message with status=warning" do
+    proc{ @connection.handle_response(:status => 'warning', :data => 'foobar') }.
+      should raise_error(RobotArmy::Warning)
+  end
+end
+
+describe RobotArmy::Connection, 'answer_sudo_prompt' do
+  before do
+    @connection = RobotArmy::Connection.new(:localhost, 'root')
+    @password_proc = proc { 'password' }
+    @stdin = StringIO.new
+    @stderr = stub(:stderr, :readpartial => nil)
+  end
+
+  it "calls back using the password proc if it is a proc" do
+    # when
+    @connection.stub!(:password).and_return(@password_proc)
+    @connection.stub!(:asking_for_password?).and_return(true, false)
+    @connection.answer_sudo_prompt(@stdin, @stderr)
+
+    # then
+    @stdin.string.should == "password\n"
+  end
+
+  it "raises if password is a string and is rejected" do
+    # when
+    @connection.stub!(:password).and_return('password')
+    @connection.stub!(:asking_for_password?).and_return(true)
+
+    # then
+    proc { @connection.answer_sudo_prompt(@stdin, @stderr) }.
+      should raise_error(RobotArmy::InvalidPassword)
+  end
+
+  it "calls back three times before raising if password is a proc" do
+    calls = 0
+
+    # when
+    @connection.stub!(:password).and_return(proc{ calls += 1 })
+
+    # then
+    @connection.should_receive(:asking_for_password?).
+      exactly(4).times.and_return(true)
+    proc { @connection.answer_sudo_prompt(@stdin, @stderr) }.
+      should raise_error(RobotArmy::InvalidPassword)
+    calls.should == 3
+  end
+end