<commit_before>from __future__ import print_function
from __future__ import absolute_import
from future import standard_library
standard_library.install_aliases()
from builtins import str
from past.builtins import basestring
from builtins import object
from os import environ
import sys
from os.path import abspath, dirname, join, isfile
import logging
from cloudcompose.exceptions import CloudComposeException
from .iam import InstancePolicyController
from .ebs import EBSController
from .cloudwatch import LogsController
from cloudcompose.util import require_env_var
import boto3
import botocore
from time import sleep
import time, datetime
from retrying import retry
from pprint import pprint
from gzip import GzipFile
from base64 import b64encode
from io import StringIO
from dateutil.parser import parse
import pytz
from dateutil.tz import tzlocal

MAX_CLOUD_INIT_LENGTH = 16000

class CloudController(object):
    def __init__(self, cloud_config, ec2_client=None, asg_client=None, silent=False):
        logging.basicConfig(level=logging.ERROR)
        self.logger = logging.getLogger(__name__)
        self.cloud_config = cloud_config
        self.silent = silent
        self.config_data = cloud_config.config_data('cluster')
        self.aws = self.config_data['aws']
        self.log_driver = self.config_data.get('logging', {}).get('driver')
        self.log_group = self.config_data.get('logging', {}).get('meta', {}).get('group')
        self.log_retention = self.config_data.get('logging', {}).get('meta', {}).get('retention')
        self.instance_policy = self.aws.get('instance_policy')
        self.cluster_name = self.config_data['name']
        self.ec2 = ec2_client or self._get_ec2_client()
        self.asg = asg_client or self._get_asg_client()

    def _get_ec2_client(self):
        return boto3.client('ec2', aws_access_key_id=require_env_var('AWS_ACCESS_KEY_ID'),
                            aws_secret_access_key=require_env_var('AWS_SECRET_ACCESS_KEY'),
                            region_name=environ.get('AWS_REGION', 'us-east-1'))

    def _get_asg_client(self):
        return boto3.client('autoscaling', aws_access_key_id=require_env_var('AWS_ACCESS_KEY_ID'),
                            aws_secret_access_key=require_env_var('AWS_SECRET_ACCESS_KEY'),
                            region_name=environ.get('AWS_REGION', 'us-east-1'))

    def up(self, cloud_init=None, use_snapshots=True, upgrade_image=False, snapshot_cluster=None, snapshot_time=None):
        if snapshot_time and use_snapshots:
            snapshot_time = self._parse_localized_time(snapshot_time)
            if not self.silent:
                print('restoring from snapshot created on or before %s' % snapshot_time.strftime('%Y-%m-%d %H:%M:%S %Z'))

        self.aws['ami'] = self._resolve_ami_name(upgrade_image)
        block_device_map = self._block_device_map(use_snapshots, snapshot_cluster, snapshot_time)
        if self.log_driver == 'awslogs':
            self._create_log_group(self.log_group, self.log_retention)
        if self.aws.get('asg'):
            self._create_asg(block_device_map, cloud_init)
        else:
            self._create_instances(block_device_map, cloud_init)

    def _parse_localized_time(self, snapshot_time):
        snapshot_time = parse(snapshot_time)
        if snapshot_time.tzinfo is None or snapshot_time.tzinfo.utcoffset(snapshot_time) is None:
            snapshot_time = snapshot_time.replace(tzinfo=tzlocal())

        snapshot_time = snapshot_time.astimezone(pytz.UTC)
        return snapshot_time

    def down(self, force=False):
        if self.aws.get('asg'):
            asg_name = self.cluster_name
            try:
                self._asg_update_auto_scaling_group(
                                        AutoScalingGroupName=asg_name,
                                        MinSize=0,
                                        MaxSize=0,
                                        DesiredCapacity=0)
                if not self.silent:
                    print('auto scaling group %s size is now 0' % asg_name)
            except botocore.exceptions.ClientError as ex:
                if ex.response["Error"]["Code"] == 'ValidationError':
                    if not self.silent:
                        print('auto scaling group %s does not exist' % asg_name)
                else:
                    raise ex

        else:
            ips = [node['ip'] for node in self.aws.get('nodes', [])]
            instance_ids = self._instance_ids_from_private_ip(ips)
            if len(instance_ids) > 0:
                if force:
                    self._disable_terminate_protection(instance_ids)
                self._ec2_terminate_instances(InstanceIds=instance_ids)
                if not self.silent:
                    print('terminated %s' % ','.join(instance_ids))

    def _disable_terminate_protection(self, instance_ids):
        for instance_id in instance_ids:
            self._ec2_modify_instance_attribute(InstanceId=instance_id, DisableApiTermination={"Value": False})


    def cleanup(self):
        if self.aws.get('asg'):
            asg_name = self.cluster_name
            asg_details = self._describe_asg(asg_name)

            asg_lc = asg_details["AutoScalingGroups"][0]["LaunchConfigurationName"]
            asg_instances = len(asg_details["AutoScalingGroups"][0]["Instances"])

            if asg_instances != 0:
                if not self.silent:
                    print('unable to delete autoscaling group %s because of %s active instances' % (asg_name, asg_instances))
                    print('run cloud-compose cluster down first or wait for instances to terminate')
            else:
                self._delete_asg(asg_name)
                if not self.silent:
                    print('deleted autoscaling group %s' % asg_name)

                self._delete_launch_config(asg_lc)
                if not self.silent:
                    print('deleted launch configuration %s' % asg_lc)
        else:
            if not self.silent:
                print('cleanup has no effect for non-ASG clusters')
                print('use cloud-compose cluster down to remove instances')

    def _resolve_ami_name(self, upgrade_image):
        if self.aws['ami'].startswith('ami-'):
            return self.aws['ami']

        ami = None
        message = None

        if not upgrade_image and not self.aws.get('asg'):
            ami = self._find_ami_on_cluster()
            if ami:
                message = 'as used on other cluster nodes'

        if not ami:
            ami, creation_date = self._find_ami_by_name_tag()
            if ami:
                message = 'created on %s' % creation_date

        if ami:
            if not self.silent:
                print('ami %s resolves to %s %s' % (self.aws['ami'], ami, message))
        else:
            raise CloudComposeException('Unable to resolve AMI %s' % self.aws['ami'])<commit_msg>add ami name so we can reference it later when it gets overridden<commit_after>from __future__ import print_function
from __future__ import absolute_import
from future import standard_library
standard_library.install_aliases()
from builtins import str
from past.builtins import basestring
from builtins import object
from os import environ
import sys
from os.path import abspath, dirname, join, isfile
import logging
from cloudcompose.exceptions import CloudComposeException
from .iam import InstancePolicyController
from .ebs import EBSController
from .cloudwatch import LogsController
from cloudcompose.util import require_env_var
import boto3
import botocore
from time import sleep
import time, datetime
from retrying import retry
from pprint import pprint
from gzip import GzipFile
from base64 import b64encode
from io import StringIO
from dateutil.parser import parse
import pytz
from dateutil.tz import tzlocal

MAX_CLOUD_INIT_LENGTH = 16000

class CloudController(object):
    def __init__(self, cloud_config, ec2_client=None, asg_client=None, silent=False):
        logging.basicConfig(level=logging.ERROR)
        self.logger = logging.getLogger(__name__)
        self.cloud_config = cloud_config
        self.silent = silent
        self.config_data = cloud_config.config_data('cluster')
        self.aws = self.config_data['aws']
        self.log_driver = self.config_data.get('logging', {}).get('driver')
        self.log_group = self.config_data.get('logging', {}).get('meta', {}).get('group')
        self.log_retention = self.config_data.get('logging', {}).get('meta', {}).get('retention')
        self.instance_policy = self.aws.get('instance_policy')
        self.cluster_name = self.config_data['name']
        self.ec2 = ec2_client or self._get_ec2_client()
        self.asg = asg_client or self._get_asg_client()

    def _get_ec2_client(self):
        return boto3.client('ec2', aws_access_key_id=require_env_var('AWS_ACCESS_KEY_ID'),
                            aws_secret_access_key=require_env_var('AWS_SECRET_ACCESS_KEY'),
                            region_name=environ.get('AWS_REGION', 'us-east-1'))

    def _get_asg_client(self):
        return boto3.client('autoscaling', aws_access_key_id=require_env_var('AWS_ACCESS_KEY_ID'),
                            aws_secret_access_key=require_env_var('AWS_SECRET_ACCESS_KEY'),
                            region_name=environ.get('AWS_REGION', 'us-east-1'))

    def up(self, cloud_init=None, use_snapshots=True, upgrade_image=False, snapshot_cluster=None, snapshot_time=None):
        if snapshot_time and use_snapshots:
            snapshot_time = self._parse_localized_time(snapshot_time)
            if not self.silent:
                print('restoring from snapshot created on or before %s' % snapshot_time.strftime('%Y-%m-%d %H:%M:%S %Z'))

        self.aws['ami'] = self._resolve_ami_name(upgrade_image)
        block_device_map = self._block_device_map(use_snapshots, snapshot_cluster, snapshot_time)
        if self.log_driver == 'awslogs':
            self._create_log_group(self.log_group, self.log_retention)
        if self.aws.get('asg'):
            self._create_asg(block_device_map, cloud_init)
        else:
            self._create_instances(block_device_map, cloud_init)

    def _parse_localized_time(self, snapshot_time):
        snapshot_time = parse(snapshot_time)
        if snapshot_time.tzinfo is None or snapshot_time.tzinfo.utcoffset(snapshot_time) is None:
            snapshot_time = snapshot_time.replace(tzinfo=tzlocal())

        snapshot_time = snapshot_time.astimezone(pytz.UTC)
        return snapshot_time

    def down(self, force=False):
        if self.aws.get('asg'):
            asg_name = self.cluster_name
            try:
                self._asg_update_auto_scaling_group(
                                        AutoScalingGroupName=asg_name,
                                        MinSize=0,
                                        MaxSize=0,
                                        DesiredCapacity=0)
                if not self.silent:
                    print('auto scaling group %s size is now 0' % asg_name)
            except botocore.exceptions.ClientError as ex:
                if ex.response["Error"]["Code"] == 'ValidationError':
                    if not self.silent:
                        print('auto scaling group %s does not exist' % asg_name)
                else:
                    raise ex

        else:
            ips = [node['ip'] for node in self.aws.get('nodes', [])]
            instance_ids = self._instance_ids_from_private_ip(ips)
            if len(instance_ids) > 0:
                if force:
                    self._disable_terminate_protection(instance_ids)
                self._ec2_terminate_instances(InstanceIds=instance_ids)
                if not self.silent:
                    print('terminated %s' % ','.join(instance_ids))

    def _disable_terminate_protection(self, instance_ids):
        for instance_id in instance_ids:
            self._ec2_modify_instance_attribute(InstanceId=instance_id, DisableApiTermination={"Value": False})


    def cleanup(self):
        if self.aws.get('asg'):
            asg_name = self.cluster_name
            asg_details = self._describe_asg(asg_name)

            asg_lc = asg_details["AutoScalingGroups"][0]["LaunchConfigurationName"]
            asg_instances = len(asg_details["AutoScalingGroups"][0]["Instances"])

            if asg_instances != 0:
                if not self.silent:
                    print('unable to delete autoscaling group %s because of %s active instances' % (asg_name, asg_instances))
                    print('run cloud-compose cluster down first or wait for instances to terminate')
            else:
                self._delete_asg(asg_name)
                if not self.silent:
                    print('deleted autoscaling group %s' % asg_name)

                self._delete_launch_config(asg_lc)
                if not self.silent:
                    print('deleted launch configuration %s' % asg_lc)
        else:
            if not self.silent:
                print('cleanup has no effect for non-ASG clusters')
                print('use cloud-compose cluster down to remove instances')

    def _resolve_ami_name(self, upgrade_image):
        if self.aws['ami'].startswith('ami-'):
            return self.aws['ami']
        # Set the ami name so we can reference it downstream
        self.aws['ami_name'] = self.aws['ami']
        ami = None
        message = None

        if not upgrade_image and not self.aws.get('asg'):
            ami = self._find_ami_on_cluster()
            if ami:
                message = 'as used on other cluster nodes'

        if not ami:
            ami, creation_date = self._find_ami_by_name_tag()
            if ami:
                message = 'created on %s' % creation_date

        if ami:
            if not self.silent:
                print('ami %s resolves to %s %s' % (self.aws['ami'], ami, message))
        else:
            raise CloudComposeException('Unable to resolve AMI %s' % self.aws['ami'])
{'fraction_non_alphanumeric': 0.08971164115343538, 'fraction_numerical': 0.004271982912068352, 'mean_word_length': 2.396034816247582, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '122327', 'n_tokens_mistral': 4103, 'n_tokens_neox': 4001, 'n_words': 913}