add the ability to spawn a console for any service container

Gemfile.lock

@@ -1,8 +1,9 @@
 PATH
   remote: .
   specs:
-    simplygenius-atmos (0.11.3)
+    simplygenius-atmos (0.11.4)
       activesupport (~> 5.2.1)
+      aws-sdk-cloudwatchlogs
       aws-sdk-core
       aws-sdk-ecr
       aws-sdk-ecs
@@ -38,19 +39,22 @@ GEM
     addressable (2.7.0)
       public_suffix (>= 2.0.2, < 5.0)
     aws-eventstream (1.0.3)
-    aws-partitions (1.237.0)
-    aws-sdk-core (3.76.0)
+    aws-partitions (1.243.0)
+    aws-sdk-cloudwatchlogs (1.27.0)
+      aws-sdk-core (~> 3, >= 3.71.0)
+      aws-sigv4 (~> 1.1)
+    aws-sdk-core (3.80.0)
       aws-eventstream (~> 1.0, >= 1.0.2)
-      aws-partitions (~> 1, >= 1.228.0)
+      aws-partitions (~> 1, >= 1.239.0)
       aws-sigv4 (~> 1.1)
       jmespath (~> 1.0)
     aws-sdk-ecr (1.23.0)
       aws-sdk-core (~> 3, >= 3.71.0)
       aws-sigv4 (~> 1.1)
-    aws-sdk-ecs (1.52.0)
+    aws-sdk-ecs (1.53.0)
       aws-sdk-core (~> 3, >= 3.71.0)
       aws-sigv4 (~> 1.1)
-    aws-sdk-iam (1.31.0)
+    aws-sdk-iam (1.32.0)
       aws-sdk-core (~> 3, >= 3.71.0)
       aws-sigv4 (~> 1.1)
     aws-sdk-kms (1.25.0)
@@ -59,11 +63,11 @@ GEM
     aws-sdk-organizations (1.34.0)
       aws-sdk-core (~> 3, >= 3.71.0)
       aws-sigv4 (~> 1.1)
-    aws-sdk-s3 (1.53.0)
-      aws-sdk-core (~> 3, >= 3.71.0)
+    aws-sdk-s3 (1.57.0)
+      aws-sdk-core (~> 3, >= 3.77.0)
       aws-sdk-kms (~> 1)
       aws-sigv4 (~> 1.1)
-    aws-sdk-ssm (1.59.0)
+    aws-sdk-ssm (1.63.0)
       aws-sdk-core (~> 3, >= 3.71.0)
       aws-sigv4 (~> 1.1)
     aws-sigv4 (1.1.0)

atmos.gemspec

@@ -51,6 +51,7 @@ Gem::Specification.new do |spec|
   spec.add_dependency "aws-sdk-ssm"
   spec.add_dependency "aws-sdk-ecr"
   spec.add_dependency "aws-sdk-ecs"
+  spec.add_dependency "aws-sdk-cloudwatchlogs"
   spec.add_dependency "os", "~> 1.0.0"
   spec.add_dependency "rotp", "~> 3.3.1"
   spec.add_dependency "clipboard", "~> 1.1.2"

lib/simplygenius/atmos/commands/container.rb

@@ -14,14 +14,14 @@ def self.description
           "Manages containers in the cloud provider"
         end
 
-        subcommand "push", "Only push a container image without activating it" do
+        option ["-c", "--cluster"],
+               "CLUSTER", "The cluster name\n",
+               required: true
 
-          option ["-c", "--cluster"],
-                 "CLUSTER", "The cluster name\n",
-                 required: true
+        option ["-r", "--role"],
+               "ROLE", "The role to assume when deploying\n"
 
-          option ["-r", "--role"],
-                 "ROLE", "The role to assume when deploying\n"
+        subcommand "push", "Only push a container image without activating it" do
 
           option ["-i", "--image"],
                  "IMAGE", "The local container image to deploy\nDefaults to service/task name"
@@ -54,15 +54,8 @@ def execute
 
         subcommand "activate", "Activate a container that has already been pushed" do
 
-          option ["-c", "--cluster"],
-                 "CLUSTER", "The cluster name\n",
-                 required: true
-
-          option ["-r", "--role"],
-                 "ROLE", "The role to assume when deploying\n"
-
           option ["-v", "--revision"],
-                 "REVISION", "Use the given revision of the pushed image to activate\n"
+                 "REVISION", "Use the given revision of the pushed image\nto activate\n"
 
           option ["-l", "--list"],
                  :flag, "List the most recent pushed images\n"
@@ -122,13 +115,6 @@ def execute
 
         subcommand "deploy", "Push and activate a container" do
 
-          option ["-c", "--cluster"],
-                 "CLUSTER", "The cluster name\n",
-                 required: true
-
-          option ["-r", "--role"],
-                 "ROLE", "The role to assume when deploying\n"
-
           option ["-i", "--image"],
                  "IMAGE", "The local container image to deploy\nDefaults to service/task name"
 
@@ -164,6 +150,44 @@ def execute
           end
         end
 
+        subcommand "console", "Spawn a console and attach to it" do
+
+          option ["-p", "--persist"],
+                 :flag, "Leave the task running after disconnect\n"
+
+          parameter "NAME",
+                    "The name of the service (or task) to attach\nthe console to"
+
+          def execute
+            Atmos.config.provider.auth_manager.authenticate(ENV, role: role) do |auth_env|
+              ClimateControl.modify(auth_env) do
+                mgr = Atmos.config.provider.container_manager
+                remote_command = Array(Atmos.config['atmos.container.console.remote_command'])
+                remote_persist_command = Array(Atmos.config['atmos.container.console.remote_persist_command'])
+                log_pattern = Atmos.config['atmos.container.console.remote_log_pattern']
+                local_command = Atmos.config['atmos.container.console.local_command']
+
+                cmd = persist? ? remote_persist_command : remote_command
+                logger.debug "Running remote command: #{cmd.join(" ")}"
+                result = mgr.run_task(cluster, name, command: cmd, waiter_log_pattern: log_pattern)
+                logger.debug "Run task result: #{result}"
+                begin
+                  match = result[:log_match]
+                  local_command = local_command.collect {|c| match.named_captures.each {|n, v| c = c.gsub("<#{n}>", v) }; c }
+                  system(*local_command)
+                ensure
+                  if persist?
+                    logger.info "Console disconnected, you can reconnect with: #{local_command.join(" ")}"
+                  else
+                    logger.info "Console complete, stopping task"
+                    mgr.stop_task(cluster, result[:task_id])
+                  end
+                end
+              end
+            end
+          end
+        end
+
       end
 
     end

lib/simplygenius/atmos/providers/aws/container_manager.rb

@@ -1,6 +1,7 @@
 require_relative '../../../atmos'
 require 'aws-sdk-ecs'
 require 'aws-sdk-ecr'
+require 'aws-sdk-cloudwatchlogs'
 require 'open3'
 
 module SimplyGenius
@@ -159,7 +160,87 @@ def list_image_tags(cluster, name)
             return result
           end
 
-          private
+          def run_task(cluster, name, command:, waiter_log_pattern: nil, launch_type: "FARGATE")
+            result = {}
+
+            ecs = ::Aws::ECS::Client.new
+            resp = nil
+
+            task_opts = {
+                count: 1,
+                cluster: cluster,
+                task_definition: name,
+                launch_type: launch_type,
+                overrides: {container_overrides: [{name: name, command: command}]}
+            }
+
+            defn_arn = nil
+            resp = ecs.describe_services(cluster: cluster, services: [name])
+            if resp.services.size > 0
+              svc = resp.services.first
+              task_opts[:launch_type] = svc.launch_type
+              task_opts[:network_configuration] = svc.network_configuration.to_h
+              defn_arn = svc.task_definition
+              logger.info "Running service task as '#{task_opts[:launch_type]}'"
+            else
+              resp = ecs.list_task_definitions(family_prefix: name, sort: 'DESC')
+              defn_arn = resp.task_definition_arns.first
+              logger.info "Running task as '#{task_opts[:launch_type]}'"
+            end
+
+            resp = ecs.describe_task_definition(task_definition: defn_arn)
+            defn = resp.task_definition
+            raise "Invalid Launch type '#{launch_type}'" unless (defn.requires_compatibilities + defn.compatibilities).include?(launch_type)
+
+            log_config = defn.container_definitions.first.log_configuration
+            log_group = nil
+            log_stream_prefix = nil
+            if log_config && log_config.log_driver == "awslogs"
+              log_group = log_config.options["awslogs-group"]
+              log_stream_prefix = log_config.options["awslogs-stream-prefix"]
+            end
+            if waiter_log_pattern && log_group.nil?
+              logger.error "Cannot wait on a log unless task definition uses cloudwatch for logging"
+              waiter_log_pattern = nil
+            end
+
+            resp = ecs.run_task(**task_opts)
+            task_arn = result[:task_arn] = resp.tasks.first.task_arn
+            task_id = result[:task_id] = task_arn.split('/').last
+
+            logger.info "Waiting for task to start"
+            ecs.wait_until(:tasks_running, cluster: cluster, tasks: [task_id])
+
+            if waiter_log_pattern
+              cwl = ::Aws::CloudWatchLogs::Client.new
+
+              waiter_regexp = Regexp.new(waiter_log_pattern)
+              log_stream = "#{log_stream_prefix}/#{name}/#{task_id}"
+              logger.info "Task started, looking for log pattern in group=#{log_group} stream=#{log_stream}"
+              log_token = nil
+              10.times do
+                resp = cwl.get_log_events(log_group_name: log_group, log_stream_name: log_stream, start_from_head: true, next_token: log_token)
+                resp.events.each do |e|
+                  logger.debug("Task log #{e.timestamp}: #{e.message}")
+                  if e.message =~ waiter_regexp
+                    result[:log_match] = Regexp.last_match
+                    return result # return, not break due to doubly nested iterator
+                  end
+                end
+                log_token = resp.next_forward_token
+                sleep 1
+              end
+            end
+
+            return result
+          end
+
+          def stop_task(cluster, task)
+            ecs = ::Aws::ECS::Client.new
+            resp = ecs.stop_task(cluster: cluster, task: task)
+          end
+
+        private
 
           def run(*args, **opts)
             logger.debug("Running: #{args}")

spec/commands/container_spec.rb

@@ -249,6 +249,58 @@ module Commands
 
         end
 
+        describe "console" do
+
+          it "requires a cluster" do
+            expect(Atmos.config.provider.auth_manager).to_not receive(:authenticate)
+            expect(Atmos.config.provider.container_manager).to_not receive(:push)
+            expect { cli.run(["console", "bar"]) }.to raise_error(Clamp::UsageError, /'-c' is required/)
+          end
+
+          it "requires a name" do
+            expect(Atmos.config.provider.auth_manager).to_not receive(:authenticate)
+            expect(Atmos.config.provider.container_manager).to_not receive(:push)
+            expect { cli.run(["console", "-c", "foo"]) }.to raise_error(Clamp::UsageError, /NAME.*no value provided/)
+          end
+
+          it "runs a task" do
+            env = Hash.new
+            remote_command = Atmos.config["atmos.container.console.remote_command"] = ["run", "server"]
+            log_pattern = Atmos.config["atmos.container.console.remote_log_pattern"] = "^ssh (?<token>\\w+)@foo.com$"
+            Atmos.config["atmos.container.console.local_command"] = ["run", "client", "<token>"]
+
+            fake_match = Regexp.new(log_pattern).match("ssh abcxyz@foo.com")
+            expect(Atmos.config.provider.auth_manager).
+                to receive(:authenticate).with(ENV, role: nil).and_yield(env)
+            expect(Atmos.config.provider.container_manager).
+                to receive(:run_task).with("foo", "bar", command: remote_command, waiter_log_pattern: log_pattern).and_return(log_match: fake_match, task_id: "tid")
+            expect_any_instance_of(described_class.find_subcommand_class("console")).to receive(:system).with("run", "client", "abcxyz")
+            expect(Atmos.config.provider.container_manager).
+                to receive(:stop_task).with("foo", "tid")
+
+            cli.run(["console", "-c", "foo", "bar"])
+          end
+
+          it "runs a persistant task" do
+            env = Hash.new
+            Atmos.config["atmos.container.console.remote_command"] = ["run", "server"]
+            remote_persist_command = Atmos.config["atmos.container.console.remote_persist_command"] = ["run", "persist"]
+            log_pattern = Atmos.config["atmos.container.console.remote_log_pattern"] = "^ssh (?<token>\\w+)@foo.com$"
+            Atmos.config["atmos.container.console.local_command"] = ["run", "client", "<token>"]
+
+            fake_match = Regexp.new(log_pattern).match("ssh abcxyz@foo.com")
+            expect(Atmos.config.provider.auth_manager).
+                to receive(:authenticate).with(ENV, role: nil).and_yield(env)
+            expect(Atmos.config.provider.container_manager).
+                to receive(:run_task).with("foo", "bar", command: remote_persist_command, waiter_log_pattern: log_pattern).and_return(log_match: fake_match, task_id: "tid")
+            expect_any_instance_of(described_class.find_subcommand_class("console")).to receive(:system).with("run", "client", "abcxyz")
+            expect(Atmos.config.provider.container_manager).
+                to_not receive(:stop_task).with("foo", "tid")
+
+            cli.run(["console", "-c", "foo", "-p", "bar"])
+          end
+
+        end
       end
 
     end