<commit_before># -*- coding: utf-8 -*-
"""
Manage AWS SSM document creation and execution
"""
from __future__ import print_function
import os
import logging
import time
import json

import boto3
from botocore.exceptions import ClientError
from boto.exception import BotoServerError

from .disco_config import read_config
from .resource_helper import throttled_call, wait_for_state_boto3
from .exceptions import TimeoutError

logger = logging.getLogger(__name__)


SSM_DOCUMENTS_DIR = "ssm/documents"
SSM_EXT = ".ssm"
SSM_WAIT_TIMEOUT = 5 * 60
SSM_WAIT_SLEEP_INTERVAL = 15
AWS_DOCUMENT_PREFIX = "AWS-"
SSM_OUTPUT_ERROR_DELIMITER = "----------ERROR-------"


class DiscoSSM(object):
    """
    A simple class to manage SSM documents
    """

    S3_BUCKET_TAG = "ssm"

    def __init__(self, environment_name=None, config_aws=None):
        self.config_aws = config_aws or read_config(environment=environment_name)

        if environment_name:
            self.environment_name = environment_name.lower()
        else:
            self.environment_name = self.config_aws.environment

        self._conn = None  # Lazily initialized
        self._s3 = None  # Lazily initialized

    @property
    def conn(self):
        """The boto3 ssm connection object"""
        if not self._conn:
            self._conn = boto3.client('ssm')
        return self._conn

    # Pylint thinks 's3' isn't long enough, but it's actually a good descriptor for s3 conn...
    # pylint: disable=invalid-name
    @property
    def s3(self):
        """The boto3 s3 connection object"""
        if not self._s3:
            self._s3 = boto3.client('s3')
        return self._s3

    def get_s3_bucket_name(self):
        """Convenience method for returning the configured s3 bucket for SSM"""
        return self.config_aws.get_asiaq_s3_bucket_name(self.S3_BUCKET_TAG)

    def execute(self, instance_ids, document_name, parameters=None, comment=None, desired_status='Success'):
        """
        Executes the given SSM document against a given list of instance ids.

        Optionally takes parameters to pass to the SSM document and an audit comment to indicate why this
        command was run.

        The parameters object is expected to be in the form of:
        {
            "key": ["values"...],
        }
        """
        bucket_name = self.get_s3_bucket_name()

        arguments = {
            "InstanceIds": instance_ids,
            "DocumentName": document_name
        }

        if parameters is not None:
            arguments["Parameters"] = parameters

        if comment is not None:
            arguments["Comment"] = comment

        if bucket_name is not None:
            try:
                # Head bucket checks if a bucket exists and throws an exception if it doesn't
                self.s3.head_bucket(Bucket=bucket_name)
                arguments["OutputS3BucketName"] = bucket_name
            except ClientError:
                logger.warning(
                    "Unable to access S3 bucket '%s', output limited to 2500 characters",
                    bucket_name
                )

        logger.info(
            "Executing document '%s' against instances %s",
            document_name,
            instance_ids
        )

        try:
            command = self._send_command(**arguments)
            command_id = command["Command"]["CommandId"]

            is_successful = self._wait_for_ssm_command(command_id=command_id, desired_status=desired_status)

            output = self.get_ssm_command_output(command_id=command_id)

            self._print_ssm_output(output)

            return is_successful
        except (ClientError, BotoServerError):
            logger.exception(
                "Unable to execute document '%s' against instances %s",
                document_name,
                instance_ids
            )
            return False

    def _print_ssm_output(self, output):
        """Convenience method for printing output from an SSM command"""
        for instance, instance_output in output.iteritems():
            print("Output for instance: {}".format(instance))
            for plugin in instance_output:
                try:
                    print(u"Plugin: {}\n\n".format(plugin.get('name', '-')))
                    print(u"STDOUT:\n{}\n\n".format(plugin.get('stdout', '-')))
                    print(u"STDERR:\n{}\n\n".format(plugin.get('stderr', '-')))
                    print(u"Exit Code: {}".format(plugin.get('exit_code', 1)))
                except UnicodeEncodeError:
                    logger.exception("Encountered error while printing SSM output")

    def _wait_for_ssm_command(self, command_id, desired_status='Success', timeout=600):
        """
        Method for waiting for the completion of a given command. Requires the command_id as well as an
        optional desired_status and a timeout value in seconds.

        Defaults to a desired status of 'Success' and a timeout of 600 seconds.

        See http://docs.aws.amazon.com/ssm/latest/APIReference/API_Command.html#EC2-Type-Command-Status
        for the valid values of desired_status.

        Note that this method only waits for the desired status to NOT be 'Pending' or 'InProgress'. In other
        words, once the command terminates this method will either return True if the status of the command
        equals the desired status, or False otherwise. For example, the command could be cancelled before it
        completes, it could timeout, or it could return a non-zero exit code.
        """
        stop_time = time.time() + timeout
        while time.time() < stop_time:
            command = self._list_commands(
                CommandId=command_id
            )
            status = command["Commands"][0]["Status"]
            document_name = command["Commands"][0]["DocumentName"]
            instance_ids = command["Commands"][0]["InstanceIds"]
            # If the command is not waiting to execute or executing, let's see if we got the status we wanted
            if status not in ['Pending', 'InProgress']:
                logger.info(
                    "Execution of document '%s' against instances %s completed as '%s'",
                    document_name,
                    instance_ids,
                    status
                )
                return status == desired_status
            logger.info(
                "Waiting for execution of document '%s' against instances %s",
                document_name,
                instance_ids
            )
            time.sleep(5)
        raise TimeoutError(
            "Timed out waiting for execution of document '{0}' against instances {1} after '{2}' "
            "seconds".format(
                document_name,
                instance_ids,
                timeout
            )
        )

    def get_ssm_command_output(self, command_id):
        """
        Method for getting the output of a given command. Requires the command_id of the desired command.

        Returns a dictionary object, in the form of:

        {
            "i-c3dfed1e": [
                {
                    "name": <plugin name>,
                    "stdout": <stdout>,
                    "stderr": <stderr>,
                    "exit_code": <exit code>
                },
                ...
            ],
            ...
        }

        """
        command_invocations = self._list_command_invocations(
            CommandId=command_id,
            Details=True
        )
<commit_msg>Remove the timeout for executing SSM documents<commit_after># -*- coding: utf-8 -*-
"""
Manage AWS SSM document creation and execution
"""
from __future__ import print_function
import os
import logging
import time
import json

import boto3
from botocore.exceptions import ClientError
from boto.exception import BotoServerError

from .disco_config import read_config
from .resource_helper import throttled_call, wait_for_state_boto3
from .exceptions import TimeoutError

logger = logging.getLogger(__name__)


SSM_DOCUMENTS_DIR = "ssm/documents"
SSM_EXT = ".ssm"
SSM_WAIT_TIMEOUT = 5 * 60
SSM_WAIT_SLEEP_INTERVAL = 15
AWS_DOCUMENT_PREFIX = "AWS-"
SSM_OUTPUT_ERROR_DELIMITER = "----------ERROR-------"


class DiscoSSM(object):
    """
    A simple class to manage SSM documents
    """

    S3_BUCKET_TAG = "ssm"

    def __init__(self, environment_name=None, config_aws=None):
        self.config_aws = config_aws or read_config(environment=environment_name)

        if environment_name:
            self.environment_name = environment_name.lower()
        else:
            self.environment_name = self.config_aws.environment

        self._conn = None  # Lazily initialized
        self._s3 = None  # Lazily initialized

    @property
    def conn(self):
        """The boto3 ssm connection object"""
        if not self._conn:
            self._conn = boto3.client('ssm')
        return self._conn

    # Pylint thinks 's3' isn't long enough, but it's actually a good descriptor for s3 conn...
    # pylint: disable=invalid-name
    @property
    def s3(self):
        """The boto3 s3 connection object"""
        if not self._s3:
            self._s3 = boto3.client('s3')
        return self._s3

    def get_s3_bucket_name(self):
        """Convenience method for returning the configured s3 bucket for SSM"""
        return self.config_aws.get_asiaq_s3_bucket_name(self.S3_BUCKET_TAG)

    def execute(self, instance_ids, document_name, parameters=None, comment=None, desired_status='Success'):
        """
        Executes the given SSM document against a given list of instance ids.

        Optionally takes parameters to pass to the SSM document and an audit comment to indicate why this
        command was run.

        The parameters object is expected to be in the form of:
        {
            "key": ["values"...],
        }
        """
        bucket_name = self.get_s3_bucket_name()

        arguments = {
            "InstanceIds": instance_ids,
            "DocumentName": document_name
        }

        if parameters is not None:
            arguments["Parameters"] = parameters

        if comment is not None:
            arguments["Comment"] = comment

        if bucket_name is not None:
            try:
                # Head bucket checks if a bucket exists and throws an exception if it doesn't
                self.s3.head_bucket(Bucket=bucket_name)
                arguments["OutputS3BucketName"] = bucket_name
            except ClientError:
                logger.warning(
                    "Unable to access S3 bucket '%s', output limited to 2500 characters",
                    bucket_name
                )

        logger.info(
            "Executing document '%s' against instances %s",
            document_name,
            instance_ids
        )

        try:
            command = self._send_command(**arguments)
            command_id = command["Command"]["CommandId"]

            is_successful = self._wait_for_ssm_command(command_id=command_id, desired_status=desired_status)

            output = self.get_ssm_command_output(command_id=command_id)

            self._print_ssm_output(output)

            return is_successful
        except (ClientError, BotoServerError):
            logger.exception(
                "Unable to execute document '%s' against instances %s",
                document_name,
                instance_ids
            )
            return False

    def _print_ssm_output(self, output):
        """Convenience method for printing output from an SSM command"""
        for instance, instance_output in output.iteritems():
            print("Output for instance: {}".format(instance))
            for plugin in instance_output:
                try:
                    print(u"Plugin: {}\n\n".format(plugin.get('name', '-')))
                    print(u"STDOUT:\n{}\n\n".format(plugin.get('stdout', '-')))
                    print(u"STDERR:\n{}\n\n".format(plugin.get('stderr', '-')))
                    print(u"Exit Code: {}".format(plugin.get('exit_code', 1)))
                except UnicodeEncodeError:
                    logger.exception("Encountered error while printing SSM output")

    def _wait_for_ssm_command(self, command_id, desired_status='Success'):
        """
        Method for waiting for the completion of a given command. Requires the command_id as well as an
        optional desired_status.

        Defaults to a desired status of 'Success'.

        See http://docs.aws.amazon.com/ssm/latest/APIReference/API_Command.html#EC2-Type-Command-Status
        for the valid values of desired_status.

        Note that this method only waits for the desired status to NOT be 'Pending' or 'InProgress'. In other
        words, once the command terminates this method will either return True if the status of the command
        equals the desired status, or False otherwise. For example, the command could be cancelled before it
        completes, or it could return a non-zero exit code.
        """
        while True:
            command = self._list_commands(
                CommandId=command_id
            )
            status = command["Commands"][0]["Status"]
            document_name = command["Commands"][0]["DocumentName"]
            instance_ids = command["Commands"][0]["InstanceIds"]
            # If the command is not waiting to execute or executing, let's see if we got the status we wanted
            if status not in ['Pending', 'InProgress']:
                logger.info(
                    "Execution of document '%s' against instances %s completed as '%s'",
                    document_name,
                    instance_ids,
                    status
                )
                return status == desired_status
            logger.info(
                "Waiting for execution of document '%s' against instances %s",
                document_name,
                instance_ids
            )
            time.sleep(5)

    def get_ssm_command_output(self, command_id):
        """
        Method for getting the output of a given command. Requires the command_id of the desired command.

        Returns a dictionary object, in the form of:

        {
            "i-c3dfed1e": [
                {
                    "name": <plugin name>,
                    "stdout": <stdout>,
                    "stderr": <stderr>,
                    "exit_code": <exit code>
                },
                ...
            ],
            ...
        }

        """
        command_invocations = self._list_command_invocations(
            CommandId=command_id,
            Details=True
        )

{'fraction_non_alphanumeric': 0.08670994192005467, 'fraction_numerical': 0.0063546293132900585, 'mean_word_length': 1.9790352127009974, 'pattern_counts': {'":': 16, '<': 12, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2470203', 'n_tokens_mistral': 3938, 'n_tokens_neox': 3758, 'n_words': 1149}