ec2ctl 0.7.9 → 0.8.0

data/lib/ec2ctl/client.rb

@@ -1,39 +1,567 @@
 require "aws-sdk"
+require "timeout"
 
 module EC2Ctl
   class Client
-    def initialize(config = {})
-      @ec2_resource = Aws::EC2::Resource.new(config)
-    end
+    PingCheckError           = Class.new RuntimeError
+    CommandNotSucceeded      = Class.new RuntimeError
+    NoInstanceForExecCommand = Class.new RuntimeError
+    InvalidFilter            = Class.new RuntimeError
 
-    def instance_infos(attributes = [], search = {})
-      @ec2_resource.instances.inject Array.new do |acc, instance|
-        instance_info = attributes.inject Hash.new do |acc, attribute|
-          acc[attribute] = case attribute
-          when /\Atag:/i
-            tag = instance.tags.find {|tag| tag.key == attribute.split(":")[1..-1].join(":")}
+    INSTANCE_IDS_LIMIT           = 50
+    COMMAND_STATUS_FINISHED      = %w(Success TimedOut Cancelled Failed).freeze
+    COMMAND_STATUS_NOT_SUCCEEDED = %w(TimedOut Cancelled Failed).freeze
 
-            if tag
-              tag.value
-            else
-              nil
-            end
+    SSM_DOCUMENT_NAMES = {
+      Linux:   "AWS-RunShellScript",
+      Windows: "AWS-RunPowerShellScript",
+    }.freeze
+
+    def initialize(
+      logger:                 nil,
+      load_balancer_name:     nil,
+      platform_type:          "Linux",
+      skip_ping_check:        false,
+      commands:               [],
+      working_directory:      nil,
+      execution_timeout:      3600,
+      skip_command_waits:     false,
+      wait_interval:          5,
+      timeout_seconds:        nil,
+      comment:                "Started at #{Time.now} (#{self.class}@#{VERSION})",
+      output_s3_bucket_name:  nil,
+      output_s3_key_prefix:   nil,
+      service_role_arn:       nil,
+      notification_arn:       nil,
+      notification_events:    nil,
+      notification_type:      nil,
+      rolling_group_size:     1,
+      skip_draining_waits:    false,
+      skip_inservice_waits:   false,
+      inservice_wait_timeout: 180,
+      instance_ids:           [],
+      filters:                [],
+      attributes:             [],
+      search:                 []
+    )
+      @logger = logger
+
+      @ec2_resource = Aws::EC2::Resource.new
+      @elb_client   = Aws::ElasticLoadBalancing::Client.new
+      @ssm_client   = Aws::SSM::Client.new
+
+      @load_balancer_name     = load_balancer_name
+      @platform_type          = platform_type
+      @skip_ping_check        = skip_ping_check
+      @commands               = commands
+      @working_directory      = working_directory
+      @execution_timeout      = execution_timeout
+      @skip_command_waits     = skip_command_waits
+      @wait_interval          = wait_interval
+      @timeout_seconds        = timeout_seconds
+      @comment                = comment
+      @output_s3_bucket_name  = output_s3_bucket_name
+      @output_s3_key_prefix   = output_s3_key_prefix
+      @service_role_arn       = service_role_arn
+      @notification_arn       = notification_arn
+      @notification_events    = notification_events
+      @notification_type      = notification_type
+      @skip_draining_waits    = skip_draining_waits
+      @skip_inservice_waits   = skip_inservice_waits
+      @inservice_wait_timeout = inservice_wait_timeout
+      @instance_ids           = instance_ids
+      @filters                = filters
+      @attributes             = attributes
+      @search                 = search
+
+      if @load_balancer_name
+        @elb_instance_ids = elb_instance_states.map(&:instance_id).select do |instance_id|
+          if instance_ids.empty?
+            true
           else
-            attribute.split(".").map(&:intern).inject instance.data do |acc, method|
-              acc.send method
+            instance_ids.include? instance_id
+          end
+        end
+
+        @rolling_group_size = [rolling_group_size, @elb_instance_ids.size].min
+      end
+    end
+
+    def ec2_list
+      if @logger
+        if @logger.debug?
+          @logger.debug ec2_instances: ec2_instances
+        else
+          ec2_instances_summary = ec2_instances.each_with_object Array.new do |instance, instances_memo|
+            instance_summary = @attributes.each_with_object Hash.new do |attribute, attributes_memo|
+              attributes_memo[attribute] = query_instance_attribute(instance, attribute)
             end
+
+            instances_memo.push instance_summary
           end
 
-          acc
+          @logger.info ec2_instances_summary: ec2_instances_summary
         end
+      end
+
+      {ec2_instances: ec2_instances}
+    end
 
-        if search.empty?
-          acc.push instance_info
+    def ec2_execute
+      ping_check ec2_instances.map(&:instance_id) unless @skip_ping_check
+      execute_commands ec2_instances.map(&:instance_id)
+    end
+
+    def elb_list
+      if @logger
+        if @logger.debug?
+          @logger.debug load_balancer_descriptions: load_balancer_descriptions
         else
-          acc.push instance_info if search.all? {|sk, sv| instance_info[sk].match Regexp.new(sv)}
+          @logger.info load_balancer_descriptions_summary: load_balancer_descriptions.map {|lb|
+            {
+              load_balancer_name: lb.load_balancer_name,
+              dns_name:           lb.dns_name,
+              instances:          lb.instances.size,
+            }
+          }
+        end
+      end
+
+      {load_balancer_descriptions: load_balancer_descriptions}
+    end
+
+    def elb_status
+      load_balancer_description = load_balancer_descriptions.first
+
+      if @logger
+        @logger.info elb_instance_counts: elb_instance_counts
+        @logger.info elb_instance_states: elb_instance_states
+
+        instance_ids = elb_instance_states.map(&:instance_id)
+
+        @logger.debug ssm_instance_ping_counts:  ssm_instance_ping_counts(instance_ids)
+        @logger.debug ssm_instance_states:       ssm_instance_states(instance_ids)
+        @logger.debug load_balancer_description: load_balancer_description
+        @logger.debug load_balancer_attributes:  load_balancer_attributes
+      end
+
+      {
+        load_balancer_description: load_balancer_description,
+        load_balancer_attributes:  load_balancer_attributes,
+        elb_instance_states:       elb_instance_states,
+      }
+    end
+
+    def elb_attach
+      attach_instances @instance_ids
+    end
+
+    def elb_detach
+      detach_instances @instance_ids
+    end
+
+    def elb_execute
+      ping_check @elb_instance_ids unless @skip_ping_check
+      execute_commands @elb_instance_ids
+    end
+
+    def elb_graceful
+      ping_check @elb_instance_ids unless @skip_ping_check
+
+      @elb_instance_ids.each_slice(@rolling_group_size).to_a.tap do |instance_id_groups|
+        instance_id_groups.each.with_index do |instance_id_group, group_index|
+          @logger.info(progress: {
+            completed: @rolling_group_size * group_index,
+            remaining: @elb_instance_ids.size - @rolling_group_size * group_index,
+            total:     @elb_instance_ids.size,
+          }) if @logger
+
+          detach_instances instance_id_group
+          wait_draining instance_id_group unless @skip_draining_waits
+          execute_commands instance_id_group
+          attach_instances instance_id_group
+          wait_inservice instance_id_group unless @skip_inservice_waits
+        end
+
+        @logger.info "Everything done!" if @logger
+      end
+    end
+
+    def detach_instances(instance_ids = [])
+      response = @elb_client.deregister_instances_from_load_balancer(
+        load_balancer_name: @load_balancer_name,
+        instances:          instance_ids.map {|i| {instance_id: i}},
+      )
+
+      status = {
+        detached:   instance_ids,
+        registered: response.instances.map(&:instance_id),
+      }
+
+      if @logger
+        if @logger.debug?
+          @logger.debug status: status
+        else
+          @logger.info detached: instance_ids
+        end
+      end
+
+      status
+    end
+
+    def attach_instances(instance_ids = [])
+      response = @elb_client.register_instances_with_load_balancer(
+        load_balancer_name: @load_balancer_name,
+        instances:          instance_ids.map {|i| {instance_id: i}},
+      )
+
+      status = {
+        attached:   instance_ids,
+        registered: response.instances.map(&:instance_id),
+      }
+
+      if @logger
+        if @logger.debug?
+          @logger.debug status: status
+        else
+          @logger.info attached: instance_ids
+        end
+      end
+
+      status
+    end
+
+    def wait_draining(instance_ids = [])
+      connection_draining = load_balancer_attributes.connection_draining
+
+      unless connection_draining.enabled
+        @logger.info wait_draining_timeout: "Disabled.".freeze if @logger
+        return
+      end
+
+      if @logger
+        @logger.debug detached_instances: elb_instance_states!(instance_ids)
+        @logger.info wait_draining_timeout: connection_draining.timeout
+      end
+
+      sleep connection_draining.timeout
+    end
+
+    def wait_inservice(instance_ids = [])
+      s = elb_instance_states! instance_ids
+
+      @logger.info wait_instance_inservice: s if @logger
+
+      Timeout.timeout @inservice_wait_timeout do
+        loop do
+          sleep @wait_interval
+
+          s = elb_instance_states! instance_ids
+          @logger.debug wait_instance_inservice: s if @logger
+
+          break if s.all? {|i| i.state == "InService".freeze}
+        end
+      end
+
+      @logger.info wait_instance_inservice: s if @logger
+    end
+
+    private
+
+    def elb_instance_states!(instance_ids = nil)
+      @elb_instance_states = @elb_client.describe_instance_health(
+        load_balancer_name: @load_balancer_name
+      ).instance_states
+
+      if instance_ids
+        elb_instance_states.select do |i|
+          instance_ids.include? i.instance_id
+        end
+      else
+        @elb_instance_states
+      end
+    end
+
+    def command_console_url(command_id)
+      [
+        "https://".freeze,
+        @ssm_client.config.region,
+        ".console.aws.amazon.com/ec2/v2/home?region=".freeze,
+        @ssm_client.config.region,
+        "#Commands:CommandId=".freeze,
+        command_id,
+      ].join
+    end
+
+    def wait_command_finish(command_id)
+      Timeout.timeout(@execution_timeout + @wait_interval * 3) do
+        loop do
+          command_invocations = get_command_invocations(command_id)
+
+          @logger.debug command_invocations: command_invocations if @logger
+
+          return command_invocations if command_invocations.all? {|i| COMMAND_STATUS_FINISHED.include? i.status}
+
+          sleep @wait_interval
+        end
+      end
+    end
+
+    def get_command_invocations(command_id)
+      command_invocations = []
+      next_token           = nil
+
+      loop do
+        response = @ssm_client.list_command_invocations(
+          command_id: command_id,
+          details:    true,
+          next_token: next_token,
+        )
+
+        command_invocations += response.command_invocations
+        next_token = response.next_token
+
+        return command_invocations unless next_token
+      end
+    end
+
+    def command_params
+      return @command_params if @command_params
+
+      @command_params = {
+        commands: @commands,
+      }
+
+      @command_params.update workingDirectory: [@working_directory]      if @working_directory
+      @command_params.update executionTimeout: [@execution_timeout.to_s] if @execution_timeout
+      @command_params
+    end
+
+    def nortification_config
+      return @nortification_config if @nortification_config
+
+      @nortification_config = {}
+      @nortification_config.update notification_arn:    @notification_arn    if @notification_arn
+      @nortification_config.update notification_events: @notification_events if @notification_events
+      @nortification_config.update notification_type:   @notification_type   if @notification_type
+      @nortification_config
+    end
+
+    def send_command_params(instance_ids)
+      {
+        document_name:         SSM_DOCUMENT_NAMES[@platform_type.intern],
+        instance_ids:          instance_ids,
+        parameters:            command_params,
+        comment:               @comment,
+        output_s3_bucket_name: @output_s3_bucket_name,
+        output_s3_key_prefix:  @output_s3_key_prefix,
+        service_role_arn:      @service_role_arn,
+        notification_config:   @notification_config,
+      }
+    end
+
+    def ping_check(instance_ids = [])
+      messages        = []
+
+      check = instance_ids.each_with_object(Hash.new(0)) do |instance_id, memo|
+        instance = ssm_instance_states(instance_ids).find {|si| si.instance_id == instance_id}
+        status   = instance.nil? ? "NotFound" : instance.ping_status
+
+        messages.push "#{instance_id}'s SSM agent is #{status}." unless status == "Online"
+
+        memo[status] += 1
+      end
+
+      @logger.debug ping_check: check if @logger
+
+      unless check.reject {|k, v| k == "Online"}.empty?
+        fail PingCheckError, messages.join(" ")
+      end
+    end
+
+    def execute_commands(instance_ids = [])
+      fail NoInstanceForExecCommand, "No instance for executing commands." if instance_ids.empty?
+
+      instance_ids.each_slice INSTANCE_IDS_LIMIT do |instance_ids_slice|
+        send_command_result = @ssm_client.send_command(send_command_params(instance_ids_slice))
+        command_id          = send_command_result.command.command_id
+
+        if @logger
+          if @logger.debug?
+            @logger.debug command: send_command_result.command
+          else
+            @logger.info(
+              command_summary: {
+                command_id:   command_id,
+                console_url:  command_console_url(command_id),
+                commands:     @commands,
+                instance_ids: instance_ids_slice,
+              }
+            )
+          end
         end
 
-        acc
+        unless @skip_command_waits
+          command_invocations = wait_command_finish(command_id)
+
+          if @logger
+            if @logger.debug?
+              @logger.debug command_invocations: command_invocations
+            else
+              @logger.info(
+                command_invocations_status_counts: command_invocations.each_with_object(Hash.new(0)) {|invocation, memo|
+                  memo[invocation.status] += 1
+                }
+              )
+
+              @logger.info(
+                command_invocations_summary: command_invocations.map {|invocation|
+                  {
+                    instance_id: invocation.instance_id,
+                    status:      invocation.status,
+                    output:      invocation.command_plugins.first.output,
+                  }
+                }
+              )
+            end
+          end
+
+          if command_invocations.any? {|i| COMMAND_STATUS_NOT_SUCCEEDED.include? i.status}
+            fail CommandNotSucceeded, "One or more command invocation(s) has not succeeded."
+          end
+        end
+
+        get_command_invocations(command_id)
+      end
+    end
+
+    def elb_instance_states
+      @elb_instance_states ||= elb_instance_states!
+    end
+
+    def ssm_instance_states(instance_ids = [])
+      return @ssm_instance_states if @ssm_instance_states
+
+      @ssm_instance_states = []
+
+      loop do
+        next_token = nil
+
+        response = @ssm_client.describe_instance_information(
+          next_token: next_token,
+
+          instance_information_filter_list: [
+            {
+              key:       "PlatformTypes".freeze,
+              value_set: [@platform_type],
+            },
+            {
+              key:       "InstanceIds".freeze,
+              value_set: instance_ids,
+            },
+          ]
+        )
+
+        @ssm_instance_states += response.instance_information_list
+        next_token = response.next_token
+        break unless next_token
+      end
+
+      @ssm_instance_states
+    end
+
+    def load_balancer_descriptions
+      return @load_balancer_descriptions if @load_balancer_descriptions
+
+      @load_balancer_descriptions = []
+
+      loop do
+        next_marker = nil
+
+        params = {marker: next_marker}
+        params.update load_balancer_names: [@load_balancer_name] if @load_balancer_name
+
+        response = @elb_client.describe_load_balancers params
+
+        @load_balancer_descriptions += response.load_balancer_descriptions
+        next_marker = response.next_marker
+        break unless next_marker
+      end
+
+      @load_balancer_descriptions
+    end
+
+    def load_balancer_attributes
+      @load_balancer_attributes ||= @elb_client.describe_load_balancer_attributes(
+        load_balancer_name: @load_balancer_name,
+      ).load_balancer_attributes
+    end
+
+    def elb_instance_counts
+      elb_instance_states.each_with_object Hash.new(0) do |i, memo|
+        memo[i.state] += 1
+      end
+    end
+
+    def ssm_instance_ping_counts(instance_ids = [])
+      ssm_instance_states(instance_ids).each_with_object Hash.new(0) do |i, memo|
+        memo[i.ping_status] += 1
+      end
+    end
+
+    def ec2_instances
+      return @ec2_instances if @ec2_instances
+
+      params = {}
+
+      params.update instance_ids: @instance_ids unless @instance_ids.empty?
+
+      unless @filters.empty?
+        filters = @filters.map do |f|
+          match = f.match(/\A(.*)=(.*)\z/)
+
+          fail InvalidFilter, "Filter should be `key=value` format (got `#{f}`)." unless match
+
+          {
+            name:   match[1],
+            values: [match[2]],
+          }
+        end
+
+        params.update filters: filters
+      end
+
+      @ec2_instances = @ec2_resource.instances(params)
+
+      unless @search.empty?
+        search_hash = @search.each_with_object Hash.new do |s, search_memo|
+          match = s.match(/\A(.*)=(.*)\z/)
+
+          fail InvalidFilter, "Search should be `key=value` format (got `#{s}`)." unless match
+
+          search_memo.update match[1] => match[2]
+        end
+
+        @ec2_instances = @ec2_instances.select do |instance|
+          search_hash.all? do |k, v|
+            Regexp.new(v).match query_instance_attribute(instance, k)
+          end
+        end
+      end
+
+      @ec2_instances
+    end
+
+    def query_instance_attribute(instance, attribute)
+      case attribute
+      when /\Atag:/i
+        tag = instance.tags.find {|t| t.key == attribute.split(":")[1..-1].join(":")}
+
+        tag ? tag.value : nil
+      else
+        attribute.split(".").map(&:intern).inject instance.data do |_acc, method|
+          _acc.send method
+        end
       end
     end
   end