hibernate 0.1.1 → 0.1.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 71a95bed5e00b8b00d0ec93a5f5f3e00faee39b1de4bc2e5f3616c6358a4c98b
-  data.tar.gz: ba6cee2ff815f36090f51e802f361079d486bce36f04ba79eee6f719c32a3437
+  metadata.gz: 21e7dde6cd66464fa7aff8bc01d5f1338e65ba5f180d27c6d7fc176fcf9dbff5
+  data.tar.gz: 84a9f6eb6fd06373cbe68b686048107636cb2821eee5284b9a408062378e05f8
 SHA512:
-  metadata.gz: 57e1625b9990bfd59c947a8f894495be4b0d1e5ed9e6601d3e84ee67bc1d8cb760111d917211b8f68b4fa6b4db6d491f621c168679ded265415edf4010ab97cd
-  data.tar.gz: 1cbe3c34db01b4235f31544ad87452cd722d165b47bf73a50506ccea0b0b1e083c0adde28a70bf84d1b03564ca2860b8f42a39fecab886c7d0f927635be65a89
+  metadata.gz: 250624e8f6984ca9cee0685c8d99a990e6fb45c66a88a16a9fd1bc68f430d252c1377d55ca31f1b95806250748e1b8be683fef73f374e1dc708bbba77b65f8e0
+  data.tar.gz: 5495ddec8a33acf9c2b4a8eeb8026f762d8f37716912bb238d6c797587403e62f0bb8023202d1755b522f6e09292146bc298e97e1bd335febed73124d54a7658

data/bin/hibernate

@@ -12,28 +12,57 @@ class HibernateCLI
     if ARGV.include?('setup')
       command = :setup
       ARGV.delete('setup')
-    elsif ARGV.include?('node')
+    elsif ARGV.include?('rule')
       command = :manage_ec2
-      ARGV.delete('node')
-    elsif ARGV.include?('remove')
-      command = :remove_ec2
-      ARGV.delete('remove')
+      ARGV.delete('rule')
     end
 
     parser = OptionParser.new do |parser|
       parser.banner = "Usage: hibernate [command] [options]"
 
-      parser.on('--in=<INSTANCE_NAME>', 'Specify the EC2 instance name') do |instance_name|
+      parser.on('--instance-name=<INSTANCE_NAME>', 'Specify the EC2 instance name') do |instance_name|
         options[:instance_name] = instance_name
       end
 
-      parser.on('--start_instance=<START_CRON>', 'Specify the cron expression to start the instance') do |start_cron|
+      parser.on('--start-expression=<START_CRON>', 'Specify the cron expression to start the instance') do |start_cron|
         options[:start_cron] = start_cron
       end
 
-      parser.on('--stop_instance=<STOP_CRON>', 'Specify the cron expression to stop the instance') do |stop_cron|
+      parser.on('--stop-expression=<STOP_CRON>', 'Specify the cron expression to stop the instance') do |stop_cron|
         options[:stop_cron] = stop_cron
       end
+
+      parser.on('--start-instance=<true/false>', 'Filter start rules') do |start_instance|
+        options[:start_instance] = start_instance == 'true'
+      end
+
+      parser.on('--stop-instance=<true/false>', 'Filter stop rules') do |stop_instance|
+        options[:stop_instance] = stop_instance == 'true'
+      end
+
+      parser.on('--rule-name=<RULE Name>', 'Specify the rule name to remove or update') do |rule_name|
+        options[:rule_name] = rule_name
+      end
+
+      parser.on('--update', 'Update an existing rule') do
+        options[:update] = true
+      end
+
+      parser.on('--remove', 'Remove an existing rule') do
+        options[:remove] = true
+      end
+
+      parser.on('--create', 'Create a new rule') do
+        options[:create] = true
+      end
+
+      parser.on('--list', 'List rules') do
+        options[:list] = true
+      end
+
+      parser.on('--state=<enable/disable>', 'Set rule state to either enable or disable') do |state|
+        options[:state] = state
+      end
     end
 
     begin
@@ -54,9 +83,15 @@ class HibernateCLI
     when :setup
       create_lambda_function
     when :manage_ec2
-      manage_ec2_command(options)
-    when :remove_ec2
-      remove_ec2_command(options)
+      if options[:update]
+        update_ec2_rule_command(options)
+      elsif options[:remove]
+        remove_ec2_rule_command(options)
+      elsif options[:create]
+        create_ec2_rule_command(options)
+      elsif options[:list]
+        list_ec2_rule_command(options)
+      end
     end
   end
 
@@ -64,25 +99,50 @@ class HibernateCLI
     LambdaSetup.new.run
   end
 
-  def self.manage_ec2_command(options)
+  def self.create_ec2_rule_command(options)
     if options[:instance_name].nil? || (options[:start_cron].nil? && options[:stop_cron].nil?)
       puts "Please provide the instance name, and at least one cron expression (start or stop)."
-      puts "Usage: hibernate node --in=<INSTANCE_NAME> --start_instance=<START_CRON> --stop_instance=<STOP_CRON>"
+      puts "Usage: hibernate rule --create --instance-name=<INSTANCE_NAME> --start-expression=<START_CRON> --stop-expression=<STOP_CRON>"
       exit
     else
-      ec2_manager = EC2Manager.new(options[:instance_name], options[:start_cron], options[:stop_cron])
-      ec2_manager.create_event_rule
+      ec2_manager = EC2Manager.new(options[:instance_name])
+      ec2_manager.create_event_rule(options[:start_cron], options[:stop_cron])
     end
   end
 
-  def self.remove_ec2_command(options)
-    if options[:instance_name].nil? || (options[:start_cron].nil? && options[:stop_cron].nil?)
-      puts "Please provide the instance name, and at least one cron expression (start or stop) to remove."
-      puts "Usage: hibernate remove --in=<INSTANCE_NAME> --start_instance=<START_CRON> --stop_instance=<STOP_CRON>"
+  def self.remove_ec2_rule_command(options)
+    if options[:rule_name].nil?
+      puts "Please provide the rule name to remove."
+      puts "Usage: hibernate rule --remove --rule-name=<RULE NAME>"
+      exit
+    else
+      ec2_manager = EC2Manager.new(options[:instance_name])
+      ec2_manager.remove_event_rule(options[:rule_name])
+    end
+  end
+
+  def self.list_ec2_rule_command(options)
+    ec2_manager = EC2Manager.new(options[:instance_name])
+    ec2_manager.list_event_rules(options)
+  end
+
+  def self.update_ec2_rule_command(options)
+    if options[:rule_name].nil?
+      puts "Please provide the rule name to update."
+      puts "Usage: hibernate rule --update --rule-name=<RULE_NAME> [--start-expression=<START_CRON>] [--stop-expression=<STOP_CRON>] [--state=<enable/disable>]"
+      exit
+    elsif options[:start_cron].nil? && options[:stop_cron].nil? && options[:state].nil?
+      puts "Please provide atleast one attribute to update."
+      puts "Usage: hibernate rule --update --rule-name=<RULE_NAME> [--start-expression=<START_CRON>] [--stop-expression=<STOP_CRON>] [--state=<enable/disable>]"
       exit
     else
-      ec2_manager = EC2Manager.new(options[:instance_name], options[:start_cron], options[:stop_cron])
-      ec2_manager.remove_event_rule
+      ec2_manager = EC2Manager.new(options[:instance_name])
+      ec2_manager.update_event_rule(
+        rule_name: options[:rule_name],
+        start_cron: options[:start_cron],
+        stop_cron: options[:stop_cron],
+        state: options[:state]
+      )
     end
   end
 end

data/lib/hibernate/cloud_watch_event_manager.rb

@@ -1,9 +1,10 @@
 require 'aws-sdk-cloudwatchevents'
 require 'json'
 require 'dotenv/load'
+require 'digest'
 
 class CloudWatchEventManager
-  def initialize(events_client, instance_id, instance_name, lambda_function_arn)
+  def initialize(events_client, instance_id = nil, instance_name = nil, lambda_function_arn)
     @events_client = events_client
     @instance_id = instance_id
     @instance_name = instance_name
@@ -13,38 +14,110 @@ class CloudWatchEventManager
     @lambda_client = Aws::Lambda::Client.new(region: @aws_region)
   end
 
+  attr_writer :instance_id, :instance_name
+
   def create_start_rule(cron_expression)
+    rule_name = "StartInstanceRule-#{@instance_id}-#{cron_expression_hash(cron_expression)}"
     create_rule(
-      "StartInstanceRule-#{@instance_id}",
+      rule_name,
       cron_expression,
       { instance_id: @instance_id, action: 'start' },
-      "start"
+      'start'
     )
   end
 
   def create_stop_rule(cron_expression)
+    rule_name = "StopInstanceRule-#{@instance_id}-#{cron_expression_hash(cron_expression)}"
     create_rule(
-      "StopInstanceRule-#{@instance_id}",
+      rule_name,
       cron_expression,
       { instance_id: @instance_id, action: 'stop' },
-      "stop"
+      'stop'
     )
   end
 
-  def remove_start_rule
-    rule_name = "StartInstanceRule-#{@instance_id}"
-    remove_rule(rule_name)
+  def remove_rule(rule_name)
+    remove_rule_by_name(rule_name)
     remove_lambda_permission(rule_name)
   end
 
-  def remove_stop_rule
-    rule_name = "StopInstanceRule-#{@instance_id}"
-    remove_rule(rule_name)
-    remove_lambda_permission(rule_name)
+  def list_event_rules(options)
+    next_token = nil
+
+    column_widths = {
+      rule_name: 40,
+      instance_id: 22,
+      schedule: 30,
+      state: 10,
+      action: 10
+    }
+
+    print_header(column_widths)
+
+    loop do
+      response = @events_client.list_rules(next_token: next_token)
+
+      response.rules.each do |rule|
+        process_rule(rule, options, column_widths)
+      end
+
+      next_token = response.next_token
+      break if next_token.nil?
+    end
+
+    print_footer(column_widths)
+  end
+
+  def update_rule_state(rule_name, state)
+    state = state == 'enable' ? 'ENABLED' : 'DISABLED'
+
+    rule_details = @events_client.describe_rule({ name: rule_name })
+
+    params = {
+      name: rule_name,
+      state: state,
+      description: rule_details.description
+    }
+
+    if rule_details.schedule_expression
+      params[:schedule_expression] = rule_details.schedule_expression
+    elsif rule_details.event_pattern
+      params[:event_pattern] = rule_details.event_pattern
+    else
+      puts "No ScheduleExpression or EventPattern found for rule '#{rule_name}'."
+      exit 1
+    end
+  
+    @events_client.put_rule(params)
+  end
+
+  def get_instance_id_from_rule(rule_name)
+    response = @events_client.list_targets_by_rule({ rule: rule_name })
+    return nil if response.targets.empty?
+
+    target_input = response.targets[0].input
+    parsed_input = JSON.parse(target_input)
+    parsed_input['instance_id'] if parsed_input.key?('instance_id')
+  rescue Aws::CloudWatchEvents::Errors::ResourceNotFoundException => e
+    puts "Error fetching targets for rule: #{rule_name} - #{e.message}"
+    nil
+  end
+
+  def rule_exists?(rule_name)
+    begin
+      response = @events_client.describe_rule({ name: rule_name })
+      return true unless response.nil?
+    rescue Aws::CloudWatchEvents::Errors::ResourceNotFoundException
+      return false
+    end
   end
 
   private
 
+  def cron_expression_hash(cron_expression)
+    Digest::SHA256.hexdigest(cron_expression)[0..7] 
+  end
+
   def create_rule(rule_name, cron_expression, input_data, action)
     @events_client.put_rule({
       name: rule_name,
@@ -85,7 +158,7 @@ class CloudWatchEventManager
     end
   end
 
-  def remove_rule(rule_name)
+  def remove_rule_by_name(rule_name)
     @events_client.remove_targets({
       rule: rule_name,
       ids: ['1']
@@ -94,7 +167,7 @@ class CloudWatchEventManager
       name: rule_name
     })
 
-    puts "Removed rule '#{rule_name}' for instance '#{@instance_name}' (ID: #{@instance_id})."
+    puts "Removed rule '#{rule_name}'"
   end
 
   def remove_lambda_permission(rule_name)
@@ -108,4 +181,52 @@ class CloudWatchEventManager
       puts "Permission not found: #{e.message}"
     end
   end
+
+  def print_header(column_widths)
+    total_width = column_widths.values.sum + 8
+    puts "-" * total_width
+    puts "| #{'Rule Name'.ljust(column_widths[:rule_name])} | " \
+         "#{ 'Instance ID'.ljust(column_widths[:instance_id])} | " \
+         "#{ 'Schedule (UTC)'.ljust(column_widths[:schedule])} | " \
+         "#{ 'State'.ljust(column_widths[:state])} | " \
+         "#{ 'Action'.ljust(column_widths[:action])} |"
+    puts "-" * total_width
+  end
+
+  def print_footer(column_widths)
+    total_width = column_widths.values.sum + 8
+    puts '-' * total_width
+  end
+
+  def process_rule(rule, options, column_widths)
+    targets = @events_client.list_targets_by_rule(rule: rule.name).targets
+    target = targets.find { |t| t.arn == @lambda_function_arn }
+
+    return unless target
+
+    input = JSON.parse(target.input)
+    action = input['action']
+    rule_instance_id = input['instance_id']
+
+    if matches_criteria?(rule_instance_id, action, options)
+      print_rule(rule, rule_instance_id, action, column_widths)
+    end
+  end
+
+  def matches_criteria?(rule_instance_id, action, options)
+    instance_id_match = options[:instance_id].nil? || options[:instance_id] == rule_instance_id
+    action_match = (options[:start_instance] && action == 'start') ||
+                   (options[:stop_instance] && action == 'stop') ||
+                   (options[:start_instance].nil? && options[:stop_instance].nil?)
+
+    instance_id_match && action_match
+  end
+
+  def print_rule(rule, rule_instance_id, action, column_widths)
+    puts "| #{rule.name.ljust(column_widths[:rule_name])} | " \
+         "#{rule_instance_id.ljust(column_widths[:instance_id])} | " \
+         "#{rule.schedule_expression.ljust(column_widths[:schedule])} | " \
+         "#{rule.state.ljust(column_widths[:state])} | " \
+         "#{action.capitalize.ljust(column_widths[:action])} |"
+  end
 end

data/lib/hibernate/ec2_manager.rb

@@ -4,10 +4,8 @@ require 'json'
 require_relative 'cloud_watch_event_manager' # Adjust the path to where the new class is located
 
 class EC2Manager
-  def initialize(instance_name, start_cron, stop_cron)
+  def initialize(instance_name = nil)
     @instance_name = instance_name
-    @start_cron = start_cron
-    @stop_cron = stop_cron
     @aws_region = ENV['AWS_REGION']
     @account_id = ENV['ACCOUNT_ID']
 
@@ -16,11 +14,13 @@ class EC2Manager
 
     @lambda_function_name = "ec2_auto_shutdown_start_function"
     @lambda_function_arn = construct_lambda_function_arn
-    @instance_id = get_instance_id_by_name
+    @instance_id = get_instance_id_by_name unless @instance_name.nil?
 
     @cloudwatch_event_manager = CloudWatchEventManager.new(@events_client, @instance_id, @instance_name, @lambda_function_arn)
   end
 
+  attr_writer :instance_name, :instance_id
+
   def get_instance_id_by_name
     response = @ec2_client.describe_instances({
       filters: [
@@ -39,16 +39,58 @@ class EC2Manager
     instance_id
   end
 
-  def create_event_rule
-    @cloudwatch_event_manager.create_start_rule(@start_cron) unless @start_cron.nil?
-    @cloudwatch_event_manager.create_stop_rule(@stop_cron) unless @stop_cron.nil?
+  def create_event_rule(start_cron, stop_cron)
+    @cloudwatch_event_manager.create_start_rule(start_cron) unless start_cron.nil?
+    @cloudwatch_event_manager.create_stop_rule(stop_cron) unless stop_cron.nil?
     puts "CloudWatch Events created for instance '#{@instance_name}' (ID: #{@instance_id})."
   end
 
-  def remove_event_rule
-    @cloudwatch_event_manager.remove_start_rule unless @start_cron.nil?
-    @cloudwatch_event_manager.remove_stop_rule unless @stop_cron.nil?
-    puts "CloudWatch Events removed for instance '#{@instance_name}' (ID: #{@instance_id})."
+  def remove_event_rule(rule_name)
+    @cloudwatch_event_manager.remove_rule(rule_name)
+    puts "CloudWatch Events removed for rule: #{rule_name}."
+  end
+
+  def list_event_rules(options)
+    options[:instance_id] = @instance_id unless @instance_id.nil?
+    @cloudwatch_event_manager.list_event_rules(options)
+  end
+
+  def update_event_rule(rule_name:, start_cron:, stop_cron:, state:)
+    rule_exists = @cloudwatch_event_manager.rule_exists?(rule_name)
+
+    unless rule_exists
+      puts "Rule '#{rule_name}' does not exist."
+      exit 1
+    end
+
+    target_instance_id = @cloudwatch_event_manager.get_instance_id_from_rule(rule_name)
+    if target_instance_id.nil?
+      puts "No targets found for the rule '#{rule_name}'."
+      exit 1
+    end
+
+    instance_name = get_instance_name_by_id(target_instance_id)
+
+    @cloudwatch_event_manager.instance_id = target_instance_id
+    @cloudwatch_event_manager.instance_name = instance_name
+
+    self.instance_id = target_instance_id
+    self.instance_name = instance_name
+
+    puts "Found instance ID: #{@instance_id} from the rule: #{rule_name}"
+
+    if start_cron || stop_cron
+      puts "Removing old rule: #{rule_name} as cron expression is being updated."
+      remove_event_rule(rule_name)
+
+      create_event_rule(start_cron, stop_cron)
+      puts "Created new rule with updated cron expression for instance '#{@instance_name}' (ID: #{@instance_id})."
+    end
+
+    return if state.nil?
+
+    @cloudwatch_event_manager.update_rule_state(rule_name, state)
+    puts "Rule '#{rule_name}' has been #{state == 'enable' ? 'enabled' : 'disabled'}."
   end
 
   private
@@ -56,4 +98,21 @@ class EC2Manager
   def construct_lambda_function_arn
     "arn:aws:lambda:#{@aws_region}:#{@account_id}:function:#{@lambda_function_name}"
   end
+
+  def get_instance_name_by_id(instance_id)
+    response = @ec2_client.describe_instances({
+      instance_ids: [instance_id]
+    })
+
+    if response.reservations.empty?
+      puts "No instance found with ID '#{instance_id}'."
+      return nil
+    end
+
+    instance = response.reservations[0].instances[0]
+    tags = instance.tags || []
+
+    name_tag = tags.find { |tag| tag.key == 'Name' }
+    name_tag&.value
+  end
 end

data/lib/hibernate/version.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Hibernate
-  VERSION = "0.1.1"
-end
+  VERSION = '0.1.3'
+end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: hibernate
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.1.3
 platform: ruby
 authors:
 - Manish Sharma
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2024-09-30 00:00:00.000000000 Z
+date: 2024-10-05 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: aws-sdk-ec2