executecommand.py

# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
import logging
import json
import errno

from subprocess import check_call
from awscli.compat import ignore_user_entered_signals
from awscli.clidriver import ServiceOperation, CLIOperationCaller

logger = logging.getLogger(__name__)

ERROR_MESSAGE = (
    'SessionManagerPlugin is not found. ',
    'Please refer to SessionManager Documentation here: ',
    'http://docs.aws.amazon.com/console/systems-manager/',
    'session-manager-plugin-not-found'
)

TASK_NOT_FOUND = (
    'The task provided in the request was '
    'not found.'
)


class ECSExecuteCommand(ServiceOperation):

    def create_help_command(self):
        help_command = super(ECSExecuteCommand, self).create_help_command()
        # change the output shape because the command provides no output.
        self._operation_model.output_shape = None
        return help_command


def get_container_runtime_id(client, container_name, task_id, cluster_name):
    describe_tasks_params = {
        "cluster": cluster_name,
        "tasks": [task_id]
    }
    describe_tasks_response = client.describe_tasks(**describe_tasks_params)
    # need to fail here if task has failed in the intermediate time
    tasks = describe_tasks_response['tasks']
    if not tasks:
        raise ValueError(TASK_NOT_FOUND)
    response = describe_tasks_response['tasks'][0]['containers']
    for container in response:
        if container_name == container['name']:
            return container['runtimeId']


def build_ssm_request_paramaters(response, client):
    cluster_name = response['clusterArn'].split('/')[-1]
    task_id = response['taskArn'].split('/')[-1]
    container_name = response['containerName']
    # in order to get container run-time id
    # we need to make a call to describe-tasks
    container_runtime_id = \
        get_container_runtime_id(client, container_name,
                                 task_id, cluster_name)
    target = "ecs:{}_{}_{}".format(cluster_name, task_id,
                                   container_runtime_id)
    ssm_request_params = {"Target": target}
    return ssm_request_params


class ExecuteCommandCaller(CLIOperationCaller):
    def invoke(self, service_name, operation_name, parameters, parsed_globals):
        try:
            # making an execute-command call to connect to an
            # active session on a container would require
            # session-manager-plugin to be installed on the client machine.
            # Hence, making this empty session-manager-plugin call
            # before calling execute-command to ensure that
            # session-manager-plugin is installed
            # before execute-command-command is made
            check_call(["session-manager-plugin"])
            client = self._session.create_client(
                service_name, region_name=parsed_globals.region,
                endpoint_url=parsed_globals.endpoint_url,
                verify=parsed_globals.verify_ssl)
            response = client.execute_command(**parameters)
            region_name = client.meta.region_name
            profile_name = self._session.profile \
                if self._session.profile is not None else ''
            endpoint_url = client.meta.endpoint_url
            ssm_request_params = build_ssm_request_paramaters(response, client)
            # ignore_user_entered_signals ignores these signals
            # because if signals which kills the process are not
            # captured would kill the foreground process but not the
            # background one. Capturing these would prevents process
            # from getting killed and these signals are input to plugin
            # and handling in there
            with ignore_user_entered_signals():
                # call executable with necessary input
                check_call(["session-manager-plugin",
                            json.dumps(response['session']),
                            region_name,
                            "StartSession",
                            profile_name,
                            json.dumps(ssm_request_params),
                            endpoint_url])
            return 0
        except OSError as ex:
            if ex.errno == errno.ENOENT:
                logger.debug('SessionManagerPlugin is not present',
                             exc_info=True)
                raise ValueError(''.join(ERROR_MESSAGE))