com.proofpoint.galaxy.cli 0.6 → 0.7

data/lib/galaxy.rb

@@ -2,7 +2,23 @@
 require 'optparse'
 require 'httpclient'
 require 'json'
+require 'uri'
+require 'base64'
+require 'digest/md5'
+require 'ssh/key/signer'
 require 'galaxy/version'
+require 'galaxy/colorize'
+require 'galaxy/shell'
+require 'galaxy/table'
+
+# stop http client from overriding the Date header
+class HTTPClient
+  class Session
+    def set_header(req)
+      # do nothing
+    end
+  end
+end
 
 module Galaxy
   GALAXY_VERSION = "0.1"
@@ -11,28 +27,89 @@ module Galaxy
       :success => 0,
       :no_slots => 1,
       :unsupported => 3,
-      :invalid_usage => 64
+      :invalid_usage => 64,
+      :no_signing_identities => 80,
+      :general_error => 99
   }
 
   #
   # Slot Information
   #
   class Slot
-    attr_reader :uuid, :host, :ip, :url, :binary, :config, :status
+    attr_reader :uuid, :short_id, :host, :ip, :url, :binary, :config, :status, :status_message, :path
 
-    def initialize(uuid, url, binary, config, status)
+    def initialize(uuid, short_id, url, binary, config, status, status_message, path)
       @uuid = uuid
+      @short_id = short_id
       @url = url
       @binary = binary
       @config = config
       @status = status
-      uri = URI.parse(url)
-      @host = uri.host
-      @ip = IPSocket::getaddress(host)
+      @status_message = status_message
+      @path = path
+
+      unless url.nil?
+        @host = URI.parse(url).host unless url.nil?
+        @ip = IPSocket::getaddress(host)
+      end
+
+      @host ||= "unknown"
+      @ip ||= "unknown"
+    end
+
+    def columns(colors = false)
+      status = @status
+      status_message = @status_message
+
+      if colors
+        status = case status
+          when "RUNNING" then Colorize::colorize(status, :bright, :green)
+          when "STOPPED" then status
+          else status
+        end
+        status_message = Colorize::colorize(status_message, :red)
+      end
+
+      return [@short_id, @host, status, @binary, @config, status_message].map { |value| value || '' }
+    end
+  end
+  #
+  # Agent Information
+  #
+  class Agent
+    attr_reader :agent_id, :host, :ip, :url, :status, :location, :path, :instance_type
+
+    def initialize(agent_id, status, url, location, instance_type)
+      @agent_id = agent_id
+      @url = url
+      @status = status
+      @path = path
+
+      unless url.nil?
+        @host = URI.parse(url).host unless url.nil?
+        @ip = IPSocket::getaddress(host)
+      end
+
+      @host ||= "unknown"
+      @ip ||= "unknown"
+
+      @location = location
+      @instance_type = instance_type
     end
 
-    def print_col
-      puts "#{uuid}\t#{host}\t#{status}\t#{binary}\t#{config}"
+    def columns(colors = false)
+      status = @status
+
+      if colors
+        status = case status
+          when "ONLINE" then Colorize::colorize(status, :bright, :green)
+          when "OFFLINE" then Colorize::colorize(status, :bright, :red)
+          when "PROVISIONING" then Colorize::colorize(status, :bright, :blue)
+          else status
+        end
+      end
+
+      return [@agent_id, @host, status, @instance_type, @location].map { |value| value || '' }
     end
   end
 
@@ -54,13 +131,9 @@ module Galaxy
       coordinator_request(filter, options, :get)
     end
 
-    def self.assign(filter, options, args)
+    def self.install(filter, options, args)
       if args.size != 2 then
-        raise CommandError.new(:invalid_usage, "You must specify a binary and config to assign.")
-      end
-      if filter.empty? then
-        raise CommandError.new(:invalid_usage, "You must specify a filter when for assign.")
-
+        raise CommandError.new(:invalid_usage, "You must specify a binary and config to install.")
       end
       if args[0].start_with? '@'
         config = args[0]
@@ -70,21 +143,44 @@ module Galaxy
         config = args[1]
 
       end
-      assignment = {
+      installation = {
           :binary => binary,
           :config => config
       }
-      coordinator_request(filter, options, :put, 'assignment', assignment, true)
+      coordinator_request(filter, options, :post, nil, installation, true)
+    end
+
+    def self.upgrade(filter, options, args)
+      if args.size <= 0 || args.size > 2 then
+        raise CommandError.new(:invalid_usage, "You must specify a binary version or a config version for upgrade.")
+      end
+      if filter.empty? then
+        raise CommandError.new(:invalid_usage, "You must specify a filter when for upgrade.")
+      end
+
+      if args[0].start_with? '@'
+        config_version = args[0][1..-1]
+        binary_version = args[1] if args.size > 1
+      else
+        binary_version = args[0]
+        config_version = args[1][1..-1] if args.size > 1
+
+      end
+      versions = {}
+      versions[:binaryVersion] = binary_version if binary_version
+      versions[:configVersion] = config_version if config_version
+
+      coordinator_request(filter, options, :post, 'assignment', versions, true)
     end
 
-    def self.clear(filter, options, args)
+    def self.terminate(filter, options, args)
       if !args.empty? then
-        raise CommandError.new(:invalid_usage, "You can not pass arguments to clear.")
+        raise CommandError.new(:invalid_usage, "You can not pass arguments to terminate.")
       end
       if filter.empty? then
-        raise CommandError.new(:invalid_usage, "You must specify a filter when for clear.")
+        raise CommandError.new(:invalid_usage, "You must specify a filter when for terminate.")
       end
-      coordinator_request(filter, options, :delete, 'assignment')
+      coordinator_request(filter, options, :delete)
     end
 
     def self.start(filter, options, args)
@@ -130,56 +226,152 @@ module Galaxy
       end
 
       slot = slots.first
-      command = ENV['GALAXY_SSH_COMMAND'] || "ssh"
-      Kernel.system "#{command} #{slot.host}"
+      ssh = ENV['GALAXY_SSH_COMMAND'] || "ssh"
+
+      if slot.path.nil?
+        path = "$HOME"
+      else
+        path = Shell::quote(slot.path)
+      end
+      remote_command = "cd #{path}; #{options[:ssh_command]}"
+      command = "#{ssh} #{slot.host} -t #{Shell::quote(remote_command)}"
+
+      puts command if options[:debug]
+      system(command)
       []
     end
 
-    private
+    def self.reset_to_actual(filter, options, args)
+      if !args.empty? then
+        raise CommandError.new(:invalid_usage, "You can not pass arguments to reset-to-actual.")
+      end
+      if filter.empty? then
+        raise CommandError.new(:invalid_usage, "You must specify a filter when for reset-to-actual.")
+      end
+      coordinator_request(filter, options, :delete, "expected-state")
+    end
 
-    def self.coordinator_request(filter, options, method, sub_path = nil, value = nil, is_json = false)
-      # build the uri
-      uri = options[:coordinator_url]
-      uri += '/' unless uri.end_with? '/'
-      uri += 'v1/slot/'
-      uri += sub_path unless sub_path.nil?
+    def self.agent_show(filter, options, args)
+      if !args.empty? then
+        raise CommandError.new(:invalid_usage, "You can not pass arguments to agent show.")
+      end
+      if !filter.empty? then
+        raise CommandError.new(:invalid_usage, "You can not use filters with agent show.")
+      end
+      coordinator_agent_request(filter, options, :get)
+    end
 
-      # create filter query
-      query = filter.map { |k, v| "#{URI.escape(k.to_s)}=#{URI.escape(v)}" }.join('&')
+    def self.agent_add(filter, options, args)
+      if args.size > 1 then
+        raise CommandError.new(:invalid_usage, "Agent add only accepts one argument.")
+      end
+      if !filter.empty? then
+        raise CommandError.new(:invalid_usage, "You can not use filters with agent show.")
+      end
 
+      if args.size > 0 then
+        instance_type = args[0]
+      end
+
+      provisioning = {}
+      provisioning[:instanceType] = instance_type if instance_type
+      provisioning[:availabilityZone] = options[:availability_zone] if options[:availability_zone]
+
+      query = "count=#{options[:count] || 1}"
+
+      coordinator_agent_request(filter, options, :post, query, nil, provisioning, true)
+    end
+
+    private
+
+    def self.execute_request(method, uri, query, body, options, is_json)
       # encode body as json if necessary
-      body = value
-      headers = {}
+      headers = []
       if is_json
-        body = value.to_json
-        headers['Content-Type'] = 'application/json'
+        body = body.to_json
+        headers += [['Content-Type', 'application/json']]
+      end
+
+      # add date header
+      request_time = Time.now
+      headers += [['Date', request_time.gmtime.strftime('%a, %d %b %Y %H:%M:%S GMT')]]
+
+      # remove query string if empty
+      query = nil if query.to_s.empty?
+
+      # build signature string
+      timestamp = request_time.to_i
+      method = method.to_s.upcase
+      relative_uri = URI.parse(uri).path + ('?' + query if query).to_s
+      body_md5 = Digest::MD5.hexdigest(body.to_s)
+      string_to_sign = [timestamp, method, relative_uri, body_md5].join("\n")
+      puts "string_to_sign:\n#{string_to_sign}\n--" if options[:debug]
+
+      # generate signature headers
+      signer = SSH::Key::Signer.new
+      signatures = signer.sign(string_to_sign)
+      if signatures.empty?
+        raise CommandError.new(:no_signing_identities, "No identities in SSH agent (use ssh-add)")
+      end
+      signatures.each do |sig|
+        fingerprint = sig.identity.fingerprint.delete(':')
+        signature = Base64.strict_encode64(sig.signature)
+        headers += [['Authorization', "Galaxy #{fingerprint}:#{signature}"]]
       end
 
       # log request in as a valid curl command if in debug mode
       if options[:debug]
-        if value then
-          puts "curl -H 'Content-Type: application/json' -X#{method.to_s.upcase} '#{uri}?#{query}' -d '"
-          puts body
-          puts "'"
-        else
-          puts "curl -X#{method.to_s.upcase} '#{uri}?#{query}'"
-        end
+        debug_uri = uri + ('?' + query if query).to_s
+        print "curl"
+        headers.each { |k,v| print " -H '#{k}: #{v}'" }
+        print " -X #{method}"
+        print " -d '#{body}'" if body
+        puts " '#{debug_uri}'"
+        puts
       end
 
       # execute request
-      response = HTTPClient.new.request(method, uri, query, body, headers).body
-
-      # parse response as json
-      slots_json = JSON.parse(response)
+      response = HTTPClient.new.request(method, uri, query, body, headers)
+      response_body = response.body if response.respond_to?(:body)
+      response_body = response_body.content if response_body.respond_to?(:content)
 
       # log response if in debug mode
       if options[:debug]
-        puts slots_json
+        puts
+        puts "#{response.status} #{response.reason}"
+        puts "--\n#{response_body.rstrip}\n--"
+      end
+
+      # check response code
+      if response.status / 100 != 2 then
+        raise CommandError.new(:general_error, response.reason || "Unknown error executing command")
       end
+      response_body
+    end
+
+    def self.coordinator_request(filter, options, method, sub_path = nil, value = nil, is_json = false)
+      # build the uri
+      uri = options[:coordinator_url]
+      uri += '/' unless uri.end_with? '/'
+      uri += 'v1/slot/'
+      uri += sub_path unless sub_path.nil?
+
+      # create filter query
+      params = filter
+      # todo this arbitrary rename here is just wrong
+      params["limit"] <<= options[:count] unless options[:count].nil?
+      params["pretty"] <<= "true" unless options[:debug].nil?
+      query = params.map { |k, v| v.map { |v1| "#{k}=#{URI.escape(v1)}" }.join('&') }.join('&')
+
+      # execute request
+      response_body = execute_request(method, uri, query, value, options, is_json)
+
+      # parse response as json
+      slots_json = JSON.parse(response_body)
 
       # convert parsed json into slot objects
       slots = slots_json.map do |slot_json|
-        Slot.new(slot_json['id'], slot_json['self'], slot_json['binary'], slot_json['config'], slot_json['status'])
+        Slot.new(slot_json['id'], slot_json['shortId'], slot_json['self'], slot_json['binary'], slot_json['config'], slot_json['status'], slot_json['statusMessage'], slot_json['installPath'])
       end
 
       # verify response
@@ -189,19 +381,43 @@ module Galaxy
 
       slots
     end
+
+    def self.coordinator_agent_request(filter, options, method, query = '', sub_path = nil, value = nil, is_json = false)
+      # build the uri
+      uri = options[:coordinator_url]
+      uri += '/' unless uri.end_with? '/'
+      uri += 'v1/admin/agent'
+      uri += sub_path unless sub_path.nil?
+
+      query = '&pretty' unless options[:debug].nil?
+
+      # execute request
+      response_body = execute_request(method, uri, query, value, options, is_json)
+
+      # parse response as json
+      agents_json = JSON.parse(response_body)
+
+      # convert parsed json into slot objects
+      agents = agents_json.map do |agent_json|
+        Agent.new(agent_json['agentId'], agent_json['state'], agent_json['self'], agent_json['location'], agent_json['instanceType'])
+      end
+
+      return agents
+    end
   end
 
   class CLI
 
-    COMMANDS = [:show, :assign, :clear, :start, :stop, :restart, :ssh]
+    COMMANDS = [:show, :install, :upgrade, :terminate, :start, :stop, :restart, :ssh, :reset_to_actual, :agent_show, :agent_add]
     INITIAL_OPTIONS = {
-        :coordinator_url => ENV['GALAXY_COORDINATOR'] || 'http://localhost:64000'
+        :coordinator_url => ENV['GALAXY_COORDINATOR'] || 'http://localhost:64000',
+        :ssh_command => "bash --login"
     }
 
     def self.parse_command_line(args)
       options = INITIAL_OPTIONS
 
-      filter = Hash.new
+      filter = Hash.new{[]}
 
       option_parser = OptionParser.new do |opts|
         opts.banner = "Usage: #{File.basename($0)} [options] <command>"
@@ -231,35 +447,46 @@ module Galaxy
         opts.separator 'Filters:'
 
         opts.on("-b", "--binary BINARY", "Select slots with a given binary") do |arg|
-          filter[:binary] = arg
+          filter[:binary] <<= arg
         end
 
         opts.on("-c", "--config CONFIG", "Select slots with given configuration") do |arg|
-          filter[:config] = arg
+          filter[:config] <<= arg
         end
 
         opts.on("-i", "--host HOST", "Select slots on the given hostname") do |arg|
-          filter[:host] = arg
+          filter[:host] <<= arg
         end
 
         opts.on("-I", "--ip IP", "Select slots at the given IP address") do |arg|
-          filter[:ip] = arg
+          filter[:ip] <<= arg
         end
 
         opts.on("-u", "--uuid SLOT_UUID", "Select slots with given slot uuid") do |arg|
-          filter[:uuid] = arg
+          filter[:uuid] <<= arg
+        end
+
+        opts.on("--count count", "Number of instances to install or agents to provision") do |arg|
+          options[:count] = arg
         end
 
-        opts.on("-s", "--state STATE", "Select 'r{unning}', 's{topped}', 'u{assigned}' or 'unknown' slots", [:running, :r, :stopped, :s, :unassigned, :u, :unknown]) do |arg|
+        opts.on("--availability-zone AVAILABILITY_ZONE", "Availability zone to agent provisioning") do |arg|
+          options[:availability_zone] = arg
+        end
+
+        # todo find a better command line argument
+        opts.on("-x", "--ssh-command SSH_COMMAND", "Command to execute with ssh") do |arg|
+          options[:ssh_command] = arg
+        end
+
+        opts.on("-s", "--state STATE", "Select 'r{unning}', 's{topped}' or 'unknown' slots", [:running, :r, :stopped, :s, :unknown]) do |arg|
           case arg
             when :running, :r then
-              filter[:state] = 'running'
+              filter[:state] <<= 'running'
             when :stopped, :s then
-              filter[:state] = 'stopped'
-            when :unassigned, :u then
-              filter[:state] = 'unassigned'
+              filter[:state] <<= 'stopped'
             when :unknown then
-              filter[:state] = 'unknown'
+              filter[:state] <<= 'unknown'
           end
         end
 
@@ -281,15 +508,27 @@ NOTES
       option_parser.parse!(args)
 
       puts options.map { |k, v| "#{k}=#{v}" }.join("\n") if options[:debug]
-      puts filter.map { |k, v| "#{k}=#{v}" }.join("\n") if options[:debug]
+      puts filter.map { |k, v| v.map { |v1| "#{k}=#{URI.escape(v1)}" }.join('\n') }.join('\n') if options[:debug]
 
-      if args.length == 0
+      if args.length == 0 then
         puts option_parser
         exit EXIT_CODES[:success]
       end
 
+      args[0] = args[0].gsub('-', '_');
       command = args[0].to_sym
 
+      is_agent = false
+      if command == :agent then
+        args = args.drop(1)
+        if args.length == 0 then
+          puts option_parser
+          exit EXIT_CODES[:success]
+        end
+        command = "#{command}_#{args[0]}".to_sym
+        is_agent = true
+      end
+
       unless COMMANDS.include?(command)
         raise CommandError.new(:invalid_usage, "Unsupported command: #{command}")
       end
@@ -298,18 +537,34 @@ NOTES
         raise CommandError.new(:invalid_usage, "You must set Galaxy coordinator host by passing --coordinator COORDINATOR or by setting the GALAXY_COORDINATOR environment variable.")
       end
 
-      return [command, filter, options, args.drop(1)]
+      return [command, filter, options, is_agent, args.drop(1)]
     end
 
     def self.execute(args)
       begin
-        (command, filter, options, command_args) = parse_command_line(args)
-        slots = Commands.send(command, filter, options, command_args)
-        slots = slots.sort_by { |slot| "#{slot.ip}|#{slot.binary}|#{slot.config}|#{slot.uuid}" }
-        puts '' if options[:debug]
-        puts "uuid\tip\tstatus\tbinary\tconfig"
-        slots.each { |slot| slot.print_col } unless slots.nil?
+        (command, filter, options, is_agent, command_args) = parse_command_line(args)
+        result = Commands.send(command, filter, options, command_args)
+        if !is_agent  then
+          slots = result.sort_by { |slot| [slot.ip, slot.binary || '', slot.config || '', slot.uuid] }
+          puts '' if options[:debug]
+
+          table = Table.new(['uuid', 'ip', 'status', 'binary', 'config', ''].map { |h| Colorize::colorize(h, :bright, :cyan) })
+          slots.each { |slot| table << slot.columns(STDOUT.tty?) }
+          puts table.render(STDOUT.tty?)
+        else
+          agents = result.sort_by { |agent| [agent.ip, agent.agent_id] }
+          puts '' if options[:debug]
+
+          table = Table.new(['id', 'ip', 'status', 'type', 'location'].map { |h| Colorize::colorize(h, :bright, :cyan) })
+          agents.each { |agent| table << agent.columns(STDOUT.tty?) }
+          puts table.render(STDOUT.tty?)
+        end
+
+
         exit EXIT_CODES[:success]
+      rescue Errno::ECONNREFUSED
+        puts "Coordinator refused connection: #{options[:coordinator_url]}"
+        exit EXIT_CODES[:general_error]
       rescue CommandError => e
         puts e.message
         if e.code == :invalid_usage
@@ -320,9 +575,15 @@ NOTES
           puts ''
           puts "exit: #{e.code}"
         end
-        exit EXIT_CODES[e.code]
+        exit EXIT_CODES[e.code || :general_error]
       end
     end
 
   end
 end
+
+
+if __FILE__ == $0 then
+  $:.unshift File.join(File.dirname(__FILE__), *%w[.. lib])
+  Galaxy::CLI.execute(ARGV)
+end

data/lib/galaxy/colorize.rb

@@ -0,0 +1,23 @@
+module Colorize
+  module ANSIColors
+    COLORS = {
+      :normal => "\e[0m",
+      :bright => "\e[1m",
+
+      :red => "\e[31m",
+      :green => "\e[32m",
+      :blue => "\e[34m",
+      :cyan => "\e[36m"
+    }
+  end
+
+  def self.colorize(string, *colors)
+    color_string = colors.map { |color| ANSIColors::COLORS[color] }.join
+
+    return "#{color_string}#{string}#{ANSIColors::COLORS[:normal]}"
+  end
+
+  def self.strip_colors(string)
+    string.gsub(/\e\[\d+m/, '')
+  end
+end

data/lib/galaxy/shell.rb

@@ -0,0 +1,6 @@
+module Shell
+  def self.quote(string)
+    string = string.gsub("'", %q('\\\''))
+    "'#{string}'"
+  end
+end

data/lib/galaxy/table.rb

@@ -0,0 +1,50 @@
+require 'galaxy/colorize'
+
+class Table
+  def initialize(headers = nil, rows = nil)
+    @headers = headers
+    @rows = rows || []
+  end
+
+  def <<(row)
+    @rows << row
+  end
+
+  def calculate_widths
+    (@rows + [@headers]).
+            map { |cols| cols.map { |col| Colorize::strip_colors(col) }.map(&:size) }.
+            transpose.
+            slice(0..-2).  # don't pad last column
+            map(&:max)
+  end
+
+  def render(tty = true)
+    widths = calculate_widths
+
+    rows = @rows
+    rows = [@headers] + rows if tty
+
+    rows.map { |row| render_row(row, widths, tty) }.join("\n")
+  end
+
+  def render_row(row, col_widths, tty)
+    if tty
+      row.zip(col_widths).map do |value, width|
+        width ||= Colorize.strip_colors(value).length
+        value + ' ' * (width - Colorize.strip_colors(value).length)
+      end.join('  ')
+    else
+      row.map { |value| Colorize.strip_colors(value) }.join("\t")
+    end
+  end
+end
+
+
+if __FILE__ == $0
+  table = Table.new(['uuid', 'ip', 'status', 'binary', 'config', ''])
+
+  table <<  [Colorize::colorize('1' * 6, :green), '2' * 5, '3' * 10, '4' * 10, '5' * 5, '6' * 3]
+  table <<  ['1' * 6, '2' * 5, '3' * 15, '4' * 7, '5' * 5, '6' * 10]
+
+  puts table.render(true)
+end

metadata

@@ -2,7 +2,7 @@
 name: com.proofpoint.galaxy.cli
 version: !ruby/object:Gem::Version 
   prerelease: 
-  version: "0.6"
+  version: "0.7"
 platform: ruby
 authors: 
 - Dain Sundstrom
@@ -10,7 +10,7 @@ autorequire:
 bindir: bin
 cert_chain: []
 
-date: 2011-06-22 00:00:00 -07:00
+date: 2012-01-11 00:00:00 -08:00
 default_executable: 
 dependencies: 
 - !ruby/object:Gem::Dependency 
@@ -35,6 +35,17 @@ dependencies:
         version: 1.5.1
   type: :runtime
   version_requirements: *id002
+- !ruby/object:Gem::Dependency 
+  name: sshkeyauth
+  prerelease: false
+  requirement: &id003 !ruby/object:Gem::Requirement 
+    none: false
+    requirements: 
+    - - ">="
+      - !ruby/object:Gem::Version 
+        version: 0.0.11
+  type: :runtime
+  version_requirements: *id003
 description: Galaxy command line interface
 email: 
 - dain@iq80.com
@@ -46,6 +57,9 @@ extra_rdoc_files: []
 
 files: 
 - bin/galaxy
+- lib/galaxy/colorize.rb
+- lib/galaxy/shell.rb
+- lib/galaxy/table.rb
 - lib/galaxy/version.rb
 - lib/galaxy.rb
 has_rdoc: true