Shortcuts for awsesh

bin/awsesh.rb

@@ -14,15 +14,19 @@
   gem "prx-ruby-aws-creds"
 end
 
-OPTS = Slop.parse do |o|
-  o.string "--profile", "AWS profile", default: "prx-legacy"
-  o.string "--region", 'Region (e.g., "us-east-1")', default: "us-east-1"
-  o.string "-i", "--instance", "Instance ID (e.g., i-06d0f11e24baaddg7)"
-  o.on "-h", "--help" do
-    puts o
-    exit
-  end
+opts = Slop::Options.new
+opts.banner = "usage:"
+opts.separator "  awsesh [options]"
+opts.separator "  awsesh stag|prod service command [options]"
+opts.separator "  awsesh stag|prod service web|worker command [options]"
+opts.string "--profile", "AWS profile", default: "prx-legacy"
+opts.string "--region", 'Region (e.g., "us-east-1")', default: "us-east-1"
+opts.string "-i", "--instance", "Instance ID (e.g., i-06d0f11e24baaddg7)"
+opts.on "-h", "--help" do
+  puts opts
+  exit
 end
+OPTS = Slop::Parser.new(opts).parse(ARGV)
 
 def colorize_label(label)
   colorized_label = label
@@ -38,6 +42,56 @@ def colorize_label(label)
   colorized_label
 end
 
+def connect_to_shared_task(env, name, type = nil, command = nil)
+  ecs = Aws::ECS::Client.new(region: OPTS[:region], credentials: PrxRubyAwsCreds.client_credentials(OPTS[:profile]), retry_mode: "adaptive")
+
+  cluster = nil
+  ecs.list_clusters.find do |res|
+    cluster ||= res.cluster_arns.find { |arn| arn.include?("-Shared") && arn.include?("-#{env}") }
+  end
+  abort "Shared #{env} ECS cluster not found!".red unless cluster
+
+  services = ecs.list_services(cluster:).map do |resp|
+    resp.service_arns.select do |arn|
+      if name.downcase.start_with?("router")
+        arn.downcase.include?("-dovetail#{name.downcase}")
+      else
+        arn.downcase.include?("-#{name.downcase}")
+      end
+    end
+  end.flatten
+  abort "Service matching \"#{name}\" not found".red unless services.any?
+
+  service_name = nil
+  if type
+    service_name = services.find { |arn| arn.downcase.include?("-#{type.downcase}") }
+    abort "Service type \"#{type}\" not found".red unless service_name
+  else
+    service_name = services.find { |arn| arn.downcase.include?("-worker") } || services.first
+  end
+
+  tasks = ecs.list_tasks(cluster:, service_name:).task_arns.take(1)
+  abort "No running tasks for service".red unless tasks.any?
+
+  if command == "host"
+    container_instances = ecs.describe_tasks(cluster:, tasks:).tasks.map(&:container_instance_arn).take(1)
+    instance_id = ecs.describe_container_instances(cluster:, container_instances:).container_instances.first.ec2_instance_id
+    puts
+    puts "Connecting to EC2 instance #{instance_id.greenish}"
+    puts
+    exec("aws ssm start-session --target #{instance_id} --region #{OPTS[:region]} --profile #{OPTS[:profile]}")
+  else
+    puts
+    puts "Connecting to a task for #{service_name.split("/").last.greenish}"
+    puts
+    if command
+      exec(%(aws ecs execute-command --region #{OPTS[:region]} --profile #{OPTS[:profile]} --cluster "#{cluster}" --task "#{tasks[0]}" --interactive --command "bin/application #{command}"))
+    else
+      exec(%(aws ecs execute-command --region #{OPTS[:region]} --profile #{OPTS[:profile]} --cluster "#{cluster}" --task "#{tasks[0]}" --interactive --command "/bin/bash"))
+    end
+  end
+end
+
 def show_wizard
   puts
   puts "-------------------------------------------------------------"
@@ -181,6 +235,10 @@ def show_wizard
 if OPTS[:instance]
   # Connect directly to an instance if --instance is provided
   exec("aws ssm start-session --target #{OPTS[:instance]} --region #{OPTS[:region]} --profile #{OPTS[:profile]}")
+elsif ["stag", "prod"].include?(ARGV[0]) && ["web", "worker"].include?(ARGV[2])
+  connect_to_shared_task(ARGV[0], ARGV[1], ARGV[2], ARGV[3])
+elsif ["stag", "prod"].include?(ARGV[0])
+  connect_to_shared_task(ARGV[0], ARGV[1], nil, ARGV[2])
 else
   # Otherwise, show the wizard
   show_wizard