Added "ecs-session" support for ECS Execute Command

README.md

@@ -13,13 +13,25 @@ Helper tools for AWS Systems Manager: `ssm-session`, `ssm-ssh` and `ssm-tunnel`.
   Wrapper around `aws ssm start-session` that can open
   SSM Session to an instance specified by *Name* or *IP Address*.
 
-  It doesn't need user credentials or even `sshd` running on the instace.
+  It doesn't need user credentials or even `sshd` running on the instance.
 
   Check out *[SSM Sessions the easy
   way](https://aws.nz/projects/ssm-session/)* for an example use.
 
   Works with any Linux or Windows EC2 instance registered in SSM.
 
+* **ecs-session**
+
+  Wrapper around `aws ecs execute-command` that can run a command
+  or open an interactive session to an Exec-enabled ECS container
+  specified by the service, name, IP address, etc.
+
+  It doesn't need user credentials or `sshd` running on the container,
+  however the containers must be configured to allow this access.
+
+  Check out *[Interactive shell in ECS Containers](https://aws.nz/projects/ecs-session/)*
+  for an example use.
+
 * **ssm-tunnel**
 
   Open *IP tunnel* to the SSM instance and to enable *network access*

ecs-session

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+import sys
+from ssm_tools.ecs_session_cli import main
+
+if __name__ == "__main__":
+    sys.exit(main())

setup.py

@@ -14,6 +14,7 @@
 
 SCRIPTS = [
     'ssm-session',
+    'ecs-session',
     'ssm-ssh',
     'ssm-tunnel',
 ]

ssm_tools/ecs_session_cli.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+
+# Convenience wrapper around 'aws ecs execute-command'
+#
+# See https://aws.nz/aws-utils/ecs-session for more info.
+#
+# Author: Michael Ludvig (https://aws.nz)
+
+# The script can list the available containers across all your ECS clusters.
+# In the end it executes 'aws ecs execute-command' with the appropriate parameters.
+# Supports both EC2 and Fargate ECS tasks.
+
+import os
+import sys
+import logging
+import signal
+import argparse
+import botocore.exceptions
+
+from .common import *
+from .resolver import ContainerResolver
+
+logger = logging.getLogger()
+
+def parse_args(argv):
+    """
+    Parse command line arguments.
+    """
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter, add_help=False)
+
+    add_general_parameters(parser)
+
+    group_container = parser.add_argument_group('Container Selection')
+    group_container.add_argument('CONTAINER', nargs='?', help='Task ID, Container Name or IP address')
+    group_container.add_argument('--list', '-l', dest='list', action="store_true", help='List available containers configured for ECS RunTask')
+
+    group_session = parser.add_argument_group('Session Parameters')
+    group_session.add_argument('--command', dest='command', metavar='COMMAND', default='/bin/sh', help='Command to run inside the container. Default: /bin/sh')
+
+    parser.description = 'Execute "ECS Run Task" on a given container'
+    parser.epilog = f'''
+IMPORTANT: containers must have "execute-command" setting enabled or they
+will not be recognised by {parser.prog} nor show up in --list output.
+
+Visit https://aws.nz/aws-utils/ecs-session for more info and usage examples.
+
+Author: Michael Ludvig
+'''
+
+    # Parse supplied arguments
+    args, extras = parser.parse_known_args(argv)
+
+    # If --version do it now and exit
+    if args.show_version:
+        show_version(args)
+
+    # Require exactly one of CONTAINER or --list
+    if bool(args.CONTAINER) + bool(args.list) != 1:
+        parser.error("Specify either CONTAINER or --list")
+
+    return args, extras
+
+def start_session(container, args, command):
+    aws_args = ""
+    if args.profile:
+        aws_args += f"--profile {args.profile} "
+    if args.region:
+        aws_args += f"--region {args.region} "
+
+    command = (
+        f'aws {aws_args} ecs execute-command '
+            f'--cluster {container["cluster_arn"]} '
+            f'--task {container["task_arn"]} '
+            f'--container {container["container_name"]} '
+            f'--command \'{command}\' --interactive '
+    )
+    logger.info("Running: %s", command)
+    os.system(command)
+
+def main():
+    ## Split command line to main args and optional command to run
+    args, extras = parse_args(sys.argv[1:])
+
+    global logger
+    logger = configure_logging("ecs-session", args.log_level)
+
+    try:
+        if args.list:
+            ContainerResolver(args).print_list()
+            quit(0)
+
+        container = ContainerResolver(args).resolve_container(args.CONTAINER)
+
+        if not container:
+            logger.warning("Could not find any container matching '%s'", args.CONTAINER)
+            quit(1)
+
+        start_session(container, args, args.command)
+
+    except (botocore.exceptions.BotoCoreError,
+            botocore.exceptions.ClientError) as e:
+        logger.error(e)
+        quit(1)
+
+if __name__ == "__main__":
+    main()

ssm_tools/resolver.py

@@ -9,18 +9,23 @@
 
 logger = logging.getLogger("ssm-session")
 
-class InstanceResolver():
+class CommonResolver():
     def __init__(self, args):
         # aws-cli compatible MFA cache
         cli_cache = os.path.join(os.path.expanduser('~'),'.aws/cli/cache')
 
         # Construct boto3 session with MFA cache
-        session = boto3.session.Session(profile_name=args.profile, region_name=args.region)
-        session._session.get_component('credential_provider').get_provider('assume-role').cache = botocore.credentials.JSONFileCache(cli_cache)
+        self.session = boto3.session.Session(profile_name=args.profile, region_name=args.region)
+        self.session._session.get_component('credential_provider').get_provider('assume-role').cache = botocore.credentials.JSONFileCache(cli_cache)
+
+
+class InstanceResolver(CommonResolver):
+    def __init__(self, args):
+        super().__init__(args)
 
         # Create boto3 clients from session
-        self.ssm_client = session.client('ssm')
-        self.ec2_client = session.client('ec2')
+        self.ssm_client = self.session.client('ssm')
+        self.ec2_client = self.session.client('ec2')
 
     def get_list(self):
         def _try_append(_list, _dict, _key):
@@ -160,3 +165,109 @@ def resolve_instance(self, instance):
         # Found only one instance - return it
         return instances[0]
 
+class ContainerResolver(CommonResolver):
+    def __init__(self, args):
+        super().__init__(args)
+
+        # Create boto3 clients from session
+        self.ecs_client = self.session.client('ecs')
+
+        self.containers = []
+        self._tasks = {}
+
+    def add_container(self, container):
+        _task_parsed = container['taskArn'].split(":")[-1].split("/")
+        self.containers.append({
+            "cluster_name": _task_parsed[1],
+            "task_id": _task_parsed[2],
+            "cluster_arn": self._tasks[container['taskArn']]['clusterArn'],
+            "task_arn": container['taskArn'],
+            "group_name": self._tasks[container['taskArn']]['group'],
+            "container_name": container['name'],
+            "container_ip": container['networkInterfaces'][0]['privateIpv4Address'],
+        })
+
+    def get_list(self):
+        def _try_append(_list, _dict, _key):
+            if _key in _dict:
+                _list.append(_dict[_key])
+
+        items = {}
+
+        # List ECS Clusters
+        clusters = []
+        logger.debug("Listing ECS Clusters")
+        paginator = self.ecs_client.get_paginator('list_clusters')
+        for page in paginator.paginate():
+            clusters.extend(page['clusterArns'])
+
+        if not clusters:
+            logger.warning("No ECS Clusters found.")
+            return []
+
+        # List tasks in each cluster
+        paginator = self.ecs_client.get_paginator('list_tasks')
+        for cluster in clusters:
+            logger.debug("Listing tasks in cluster: %s", cluster)
+
+            # maxResults must be <= 100 because describe_tasks() doesn't accept more than that
+            for page in paginator.paginate(cluster=cluster, maxResults=100):
+                response = self.ecs_client.describe_tasks(cluster=cluster, tasks=page['taskArns'])
+
+                # Filter containers that have a running ExecuteCommandAgent
+                for task in response['tasks']:
+                    self._tasks[task['taskArn']] = task
+                    for container in task['containers']:
+                        if not 'managedAgents' in container:
+                            continue
+                        for agent in container['managedAgents']:
+                            if agent['name'] == 'ExecuteCommandAgent' and agent['lastStatus'] == 'RUNNING':
+                                self.add_container(container)
+
+        return self.containers
+
+    def print_containers(self, containers):
+        max_len = {}
+        for container in containers:
+            for key in container.keys():
+                if not key in max_len:
+                    max_len[key] = len(container[key])
+                else:
+                    max_len[key] = max(max_len[key], len(container[key]))
+        containers.sort(key = lambda x: [x['cluster_name'], x['container_name']])
+        for container in containers:
+            print(f"{container['cluster_name']:{max_len['cluster_name']}}  {container['group_name']:{max_len['group_name']}}  {container['task_id']:{max_len['task_id']}}  {container['container_name']:{max_len['container_name']}}  {container['container_ip']:{max_len['container_ip']}}")
+
+    def print_list(self):
+        containers = self.get_list()
+
+        if not containers:
+            logger.warning("No Execute-Command capable contaianers found!")
+            quit(1)
+
+        self.print_containers(containers)
+
+    def resolve_container(self, keyword):
+        containers = self.get_list()
+
+        if not containers:
+            logger.warning("No Execute-Command capable contaianers found!")
+            quit(1)
+
+        logger.debug("Searching for: %s", keyword)
+
+        candidates = []
+        for container in containers:
+            if keyword in (container['group_name'], container['task_id'], container['container_name'], container['container_ip']):
+                candidates.append(container)
+        if not candidates:
+            logger.warning("No container found for: %s", keyword)
+            quit(1)
+        elif len(candidates) == 1:
+            self.print_containers(candidates)
+            return candidates[0]
+        else:
+            logger.warning("Found %d instances for: %s", len(containers), keyword)
+            logger.warning("Use Container IP or Task ID to connect to a specific one")
+            self.print_containers(candidates)
+            quit(1)