""" Copyright 2018 EPAM Systems, Inc. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License 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 inspect import json import time from pathlib import Path from typing import Generator, Optional, List, Iterable, Any import botocore from botocore.exceptions import ClientError from boto3 import client from syndicate.commons.log_helper import get_logger from syndicate.connection.helper import apply_methods_decorator, retry from syndicate.core.constants import EC2_LT_RESOURCE_TAGS _LOG = get_logger(__name__) def create_permissions(ranges): ip_ranges = [{'CidrIp': ip_range} for ip_range in ranges] return [{ 'IpProtocol': '-1', 'FromPort': -1, 'ToPort': -1, 'IpRanges': ip_ranges }] def preserve_default_permission(group_id, permissions): for permission in permissions: if 'UserIdGroupPairs' in permission and permission['UserIdGroupPairs']: for pair in permission['UserIdGroupPairs']: if pair['GroupId'] == group_id: permission.pop('UserIdGroupPairs') @apply_methods_decorator(retry()) class EC2Connection(object): """ EC2 connection class.""" def __init__(self, region=None, aws_access_key_id=None, aws_secret_access_key=None, aws_session_token=None): self.client = client('ec2', region, aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, aws_session_token=aws_session_token) _LOG.debug('Opened new EC2 connection.') def describe_security_groups(self, name=None, sg_id=None, vpc_id=None): filters = [] if name: if isinstance(name, list): filters.append({'Name': 'group-name', 'Values': name}) elif isinstance(name, str): filters.append({'Name': 'group-name', 'Values': [name]}) else: _LOG.warning('Unacceptable name type: %s', type(name)) if sg_id: if isinstance(sg_id, list): filters.append({'Name': 'group-id', 'Values': sg_id}) elif isinstance(sg_id, str): filters.append({'Name': 'group-id', 'Values': [sg_id]}) else: _LOG.warning( f'Unacceptable security group id type: {type(sg_id)}') if vpc_id: filters.append({'Name': 'vpc-id', 'Values': [vpc_id]}) return self.client.describe_security_groups(Filters=filters)[ 'SecurityGroups'] def describe_regions(self, name=None): filters = [] if name: if isinstance(name, list): filters.append({'Name': 'region-name', 'Values': name}) elif isinstance(name, str): filters.append({'Name': 'region-name', 'Values': [name]}) else: _LOG.warning('Unacceptable name type: %s', type(name)) return self.client.describe_regions(Filters=filters)['Regions'] def get_default_vpc_id(self): for vpc in self.client.describe_vpcs()['Vpcs']: if vpc['IsDefault']: return vpc['VpcId'] def create_sg(self, name, desc, vpc_id): return self.client.create_security_group(GroupName=name, Description=desc, VpcId=vpc_id)['GroupId'] def authorize_ingress(self, group_id, group_name, ranges): if not group_id: group_id = self.get_sg_id(group_name) permissions = create_permissions(ranges) self.client.authorize_security_group_ingress(GroupId=group_id, GroupName=group_name, IpPermissions=permissions) def revoke_ingress(self, group_id, group_name, ranges): permissions = create_permissions(ranges) self.revoke_ingress_permissions(group_id, group_name, permissions) def revoke_ingress_permissions(self, group_id, group_name, permissions): if not group_id: group_id = self.get_sg_id(group_name) preserve_default_permission(group_id, permissions) if not permissions: return self.client.revoke_security_group_ingress(GroupId=group_id, GroupName=group_name, IpPermissions=permissions) def delete_sg(self, group_name): self.client.delete_security_group(GroupName=group_name) def get_sg_id(self, group_name, vpc_id=None): if not vpc_id: vpc_id = self.get_default_vpc_id() group, = self.describe_security_groups(group_name, vpc_id) if group: return group['GroupId'] def get_key_pairs( self, dry_run: bool = False, key_names: list | None = None, filters: list | None = None, ) -> dict | list[dict]: """ :type dry_run: bool :type key_names: list :type filters: list :return: """ params = dict(DryRun=dry_run) if key_names: params['KeyNames'] = key_names if filters: params['Filters'] = filters return self.client.describe_key_pairs(**params) def if_key_pair_exists(self, key_name): key_pairs = self.get_key_pairs().get('KeyPairs') if key_pairs: for each in key_pairs: if each['KeyName'] == key_name: return True def list_vpcs(self, dry_run=None, vpc_ids=None, filters=None): params = dict() if dry_run: params['DryRun'] = dry_run if vpc_ids: params['VpcIds'] = vpc_ids if filters: params['Filters'] = filters response = self.client.describe_vpcs(**params) if response: return response['Vpcs'] def list_subnets(self, dry_run=None, subnet_ids=None, filters=None): params = dict() if dry_run: params['DryRun'] = dry_run if subnet_ids: params['SubnetIds'] = subnet_ids if filters: params['Filters'] = filters response = self.client.describe_subnets(**params) if response: return response['Subnets'] def get_azs(self): response = self.client.describe_availability_zones() return [az['ZoneName'] for az in response['AvailabilityZones']] def describe_image(self, image_id): params = dict(ImageIds=[image_id]) response = self.client.describe_images(**params) return response['Images'] def describe_instances(self, filters, instance_ids=None): params = {} result_list = [] if filters: params['Filters'] = filters if instance_ids: params['InstanceIds'] = instance_ids response = self.client.describe_instances(**params) result_list.extend([reservation['Instances'][0] for reservation in response['Reservations']]) token = response.get('NextToken') while token: # value is 'null' if there is no token params['NextToken'] = token response = self.client.describe_instances(**params) result_list.extend([reservation['Instances'][0] for reservation in response['Reservations']]) token = response.get('NextToken') return result_list def terminate_instances(self, instance_ids): self.client.terminate_instances( InstanceIds=instance_ids ) def launch_instance(self, image_id, instance_type, security_groups_names=None, security_group_ids=None, iam_instance_profile=None, name=None, key_name=None, user_data=None, tags=None, subnet_id=None, availability_zone=None): tags = tags or [] if iam_instance_profile: if not iam_instance_profile.get('Arn') \ and not iam_instance_profile.get('Name'): raise AssertionError('Provided instance profile {0}' 'is not well-formed. ' 'Arn or Name nodes required.' .format(iam_instance_profile)) if name: tags.append({ 'Key': 'Name', 'Value': name }) instance_parameters = { 'ImageId': image_id, 'InstanceType': instance_type, 'MinCount': 1, 'MaxCount': 1 } if tags: instance_parameters['TagSpecifications'] = [{ 'ResourceType': 'instance', 'Tags': tags }] if availability_zone: instance_parameters['Placement'] = { 'AvailabilityZone': availability_zone } if key_name: instance_parameters['KeyName'] = key_name if user_data: instance_parameters['UserData'] = user_data if iam_instance_profile: instance_parameters['IamInstanceProfile'] = iam_instance_profile if subnet_id: instance_parameters['SubnetId'] = subnet_id if security_groups_names: instance_parameters['SecurityGroups'] = security_groups_names if security_group_ids: instance_parameters['SecurityGroupIds'] = security_group_ids response = self.client.run_instances(**instance_parameters) # always launch only one instance launched_instances = response['Instances'] if len(launched_instances) < 1: return 'No instances launched' else: return launched_instances[0] def modify_instance_attribute(self, **kwargs): """ log_not_found_error parameter is needed for proper log handling in the retry decorator """ kwargs.pop('log_not_found_error', None) if not kwargs['InstanceId']: raise AssertionError('InstanceId must be specified') self.client.modify_instance_attribute(**kwargs) def deploy_security_groups(self, groups): default_vpc_id = self.get_default_vpc_id() groups_in_default_vpc = self.describe_security_groups( vpc_id=default_vpc_id) group_names = [g['GroupName'] for g in groups_in_default_vpc] verify = [] create = [] for group in groups: if group['n'] in group_names: verify.append(group) else: create.append(group) if create: self._create_security_groups(default_vpc_id, create) if verify: self._verify_security_groups(verify, groups_in_default_vpc) def _delete_security_groups(self, groups): vpc_id = self.client.get_default_vpc_id() security_groups = self.client.describe_security_groups(groups, vpc_id) existent_names = [sg['GroupName'] for sg in security_groups] for sg_name in groups: if sg_name not in existent_names: continue if sg_name == 'default': security_group, = self.client.describe_security_groups( 'default', vpc_id) self.client.revoke_ingress_permissions(group_id=None, group_name='default', permissions= security_group[ 'IpPermissions']) else: self.client.delete_sg(sg_name) def _create_security_groups(self, vpc_id, groups_to_create): for group in groups_to_create: if group['n'] == 'default': continue group_id = self.create_sg(group['n'], group['d'], vpc_id) ranges = group['r'] # add waiting to auth sg time.sleep(5) self.authorize_ingress(group_id, group['n'], ranges) def _verify_security_groups(self, groups, sgs): security_groups = {sg['GroupName']: sg for sg in sgs} for group in groups: name = group['n'] sg = security_groups[name] should_be = group['r'] actual = [] for permission in sg['IpPermissions']: for ip_range in permission['IpRanges']: actual.append(ip_range['CidrIp']) to_revoke = [cidr for cidr in actual if cidr not in should_be] to_authorize = [cidr for cidr in should_be if cidr not in actual] if to_revoke: self.revoke_ingress(sg['GroupId'], name, to_revoke) if to_authorize: # add waiting to auth sg time.sleep(5) self.authorize_ingress(sg['GroupId'], name, to_authorize) def associate_address(self, instance_id=None, public_ip=None, allow_reassociation=False): params = dict(AllowReassociation=allow_reassociation) if instance_id: params['InstanceId'] = instance_id if public_ip: params['PublicIp'] = public_ip return self.client.associate_address(**params) def resolve_resource_tags( # noqa Reason: @apply_methods_decorator(retry()) self, resource_tags: dict[str, dict[str, str]], ) -> list[dict | None]: base_error_message = ( "Failed to process 'resource_tags' for the EC2 launch template. " "This step will be skipped." ) if not isinstance(resource_tags, dict): _LOG.error( f"{base_error_message} Reason: 'resource_tags' should be a " f"dictionary." ) return [] resource_tag_specs = [] for resource_type, tags_dict in resource_tags.items(): if not isinstance(tags_dict, dict): _LOG.error( f"{base_error_message} The value for key '{resource_type}' " f"is not a dictionary as expected" ) return [] if resource_type not in EC2_LT_RESOURCE_TAGS: _LOG.error( f"{base_error_message} " f"Encountered an invalid resource type '{resource_type}'. " f"Valid types include: {EC2_LT_RESOURCE_TAGS}" ) return [] tag_list = [{'Key': k, 'Value': v} for k, v in tags_dict.items()] resource_tag_specs.append({ 'ResourceType': resource_type, 'Tags': tag_list, }) return resource_tag_specs def create_launch_template( self, name: str, lt_data: dict, version_description: str | None = None, tags: list[dict[str, Any]] | None = None, resource_tags: dict[str, dict[str, str]] | None = None, ) -> dict: params = { 'LaunchTemplateName': name, 'LaunchTemplateData': lt_data, } if version_description is not None: params['VersionDescription'] = version_description if tags: params['TagSpecifications'] = [{ 'ResourceType': 'launch-template', 'Tags': tags, }] if resource_tags: resource_tag_specs = self.resolve_resource_tags(resource_tags) if 'TagSpecifications' not in lt_data: lt_data['TagSpecifications'] = [] lt_data['TagSpecifications'].extend(resource_tag_specs) return self.client.create_launch_template(**params) def create_launch_template_version( self, lt_name: str | None = None, lt_id: str | None = None, source_version: str | None = None, lt_data: dict | None = None, version_description: str | None = None, resource_tags: dict[str, dict[str, str]] | None = None, ) -> dict | None: if not lt_name and not lt_id: _LOG.error( 'A launch template version cannot be created without the name ' 'or ID of the launch template' ) return None params = { 'LaunchTemplateId': lt_id or None, 'LaunchTemplateName': lt_name if not lt_id else None, 'SourceVersion': source_version or None, 'VersionDescription': version_description or None, 'LaunchTemplateData': lt_data or {}, } if lt_name and lt_id: _LOG.warning( 'Both the launch template name and ID are specified. The ' 'request will be made by ID' ) if resource_tags: resource_tag_specs = self.resolve_resource_tags(resource_tags) if 'TagSpecifications' not in lt_data: lt_data['TagSpecifications'] = [] lt_data['TagSpecifications'].extend(resource_tag_specs) return self.client.create_launch_template_version( **{k: v for k, v in params.items() if v is not None} ) def describe_launch_templates( self, lt_name: str | None = None, lt_id = None, ) -> list: result_list = list() params = dict() if lt_name is not None and lt_id is not None: _LOG.warning('Both the launch template name and ID are specified. ' 'The request will be made by ID.') if isinstance(lt_id, list): params['LaunchTemplateIds'] = lt_id elif isinstance(lt_id, str): params['LaunchTemplateIds'] = [lt_id] else: _LOG.warning( f'Unsupported launch template ID type {type(lt_id)}') elif lt_name is not None: if isinstance(lt_name, list): params['LaunchTemplateNames'] = lt_name elif isinstance(lt_name, str): params['LaunchTemplateNames'] = [lt_name] else: _LOG.warning( f'Unsupported launch template name type {type(lt_name)}') elif lt_id is not None: if isinstance(lt_id, list): params['LaunchTemplateIds'] = lt_id elif isinstance(lt_id, str): params['LaunchTemplateIds'] = [lt_id] else: _LOG.warning( f'Unsupported launch template ID type {type(lt_id)}') try: response = self.client.describe_launch_templates(**params) if ( params) else self.client.describe_launch_templates() token = response.get('NextToken') result_list.extend(response['LaunchTemplates']) while token: if params: params['NextToken'] = token response = self.client.describe_launch_templates(**params) else: response = self.client.describe_launch_templates() token = response.get('NextToken') result_list.extend(response['LaunchTemplates']) except ClientError as e: if ('InvalidLaunchTemplateName.NotFoundException' in str(e) or 'InvalidLaunchTemplateId.NotFound' in str(e)): dynamic_message = f"by name '{lt_name}'" if lt_name else \ f"by ID '{lt_id}'" _LOG.warning(f"Launch template not found by {dynamic_message}") else: raise e return result_list def delete_launch_template( self, lt_name = None, lt_id = None, log_not_found_error: bool = True, ) -> None: """ log_not_found_error parameter is needed for proper log handling in the retry decorator """ params = dict() if lt_name is not None and lt_id is not None: _LOG.warning( 'Both the launch template name and ID are specified. The ' 'request will be made by ID.' ) params['LaunchTemplateId'] = lt_id elif lt_name is not None: params['LaunchTemplateName'] = lt_name elif lt_id is not None: params['LaunchTemplateId'] = lt_id else: raise AssertionError( 'Either the launch template name or ID must be provided for ' 'removing the launch template.') self.client.delete_launch_template(**params) def modify_launch_template( self, default_version, lt_name=None, lt_id=None, ) -> None: params = dict() params['DefaultVersion'] = default_version if lt_name is not None and lt_id is not None: _LOG.warning( 'Both the launch template name and ID are specified. The ' 'request will be made by ID.' ) params['LaunchTemplateId'] = lt_id elif lt_name is not None: params['LaunchTemplateName'] = lt_name elif lt_id is not None: params['LaunchTemplateId'] = lt_id else: _LOG.error('A launch template modification can not be done ' 'without the name or ID of the launch template.') return return self.client.modify_launch_template(**params) class InstanceTypes: @staticmethod def from_api( region_name: Optional[str] = None, current_generation: Optional[bool] = None, arch: Optional[List] = None, ) -> Generator[str, None, None]: filters = [] if isinstance(current_generation, bool): filters.append({ 'Name': 'current-generation', 'Values': [str(current_generation).lower()] }) if isinstance(arch, list): assert set(arch).issubset({'arm64', 'i386', 'x86_64'}) filters.append({ 'Name': 'processor-info.supported-architecture', 'Values': arch }) params = {} if filters: params['Filters'] = filters ec2 = client('ec2', region_name=region_name) while True: res = ec2.describe_instance_types(**params) yield from (item['InstanceType'] for item in res['InstanceTypes']) _next = res.get('NextToken') if not _next: break params['NextToken'] = _next @staticmethod def from_botocore() -> Generator[str, None, None]: path = Path(inspect.getfile(botocore)).parent with open(Path(path, 'data', 'ec2', '2016-11-15', 'service-2.json'), encoding="utf8") as fp: data = json.load(fp) yield from data['shapes']['InstanceType']['enum'] @staticmethod def instance_type_group(instance_type: str) -> str: return instance_type.split('.')[0] @staticmethod def with_groups(it: Iterable[str]) -> Generator[str, None, None]: """ Before yielding an instance type, yields its group. A group is yielded only once. :param it: Iterable[str] :return: Generator[str, None, None] """ emitted = set() for instance_type in it: group = InstanceTypes.instance_type_group(instance_type) if group not in emitted: yield group emitted.add(group) yield instance_type