# Copyright 2015-2017 Capital One Services, LLC
#
# 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.
from __future__ import absolute_import, division, print_function, unicode_literals

import base64
import itertools
import operator
import random
import re
import zlib

import six
from botocore.exceptions import ClientError
from dateutil.parser import parse
from concurrent.futures import as_completed
import jmespath

from c7n.actions import (
    ActionRegistry, BaseAction, ModifyVpcSecurityGroupsAction
)

from c7n.exceptions import PolicyValidationError
from c7n.filters import (
    FilterRegistry, AgeFilter, ValueFilter, Filter, DefaultVpcBase
)
from c7n.filters.offhours import OffHour, OnHour
import c7n.filters.vpc as net_filters

from c7n.manager import resources
from c7n import query, utils
from c7n.utils import type_schema, filter_empty

from c7n.resources.iam import CheckPermissions
from c7n.resources.securityhub import PostFinding

RE_ERROR_INSTANCE_ID = re.compile("'(?P<instance_id>i-.*?)'")

filters = FilterRegistry('ec2.filters')
actions = ActionRegistry('ec2.actions')


@resources.register('ec2')
class EC2(query.QueryResourceManager):

    class resource_type(query.TypeInfo):
        service = 'ec2'
        arn_type = 'instance'
        enum_spec = ('describe_instances', 'Reservations[].Instances[]', None)
        id = 'InstanceId'
        filter_name = 'InstanceIds'
        filter_type = 'list'
        name = 'PublicDnsName'
        date = 'LaunchTime'
        dimension = 'InstanceId'
        config_type = "AWS::EC2::Instance"

        default_report_fields = (
            'CustodianDate',
            'InstanceId',
            'tag:Name',
            'InstanceType',
            'LaunchTime',
            'VpcId',
            'PrivateIpAddress',
        )

    filter_registry = filters
    action_registry = actions

    # if we have to do a fallback scenario where tags don't come in describe
    permissions = ('ec2:DescribeTags',)

    def __init__(self, ctx, data):
        super(EC2, self).__init__(ctx, data)
        self.queries = QueryFilter.parse(self.data.get('query', []))

    def resources(self, query=None):
        q = self.resource_query()
        if q is not None:
            query = query or {}
            query['Filters'] = q
        return super(EC2, self).resources(query=query)

    def resource_query(self):
        qf = []
        qf_names = set()
        # allow same name to be specified multiple times and append the queries
        # under the same name
        for q in self.queries:
            qd = q.query()
            if qd['Name'] in qf_names:
                for qf in qf:
                    if qd['Name'] == qf['Name']:
                        qf['Values'].extend(qd['Values'])
            else:
                qf_names.add(qd['Name'])
                qf.append(qd)
        return qf

    def get_source(self, source_type):
        if source_type == 'describe':
            return DescribeEC2(self)
        elif source_type == 'config':
            return query.ConfigSource(self)
        raise ValueError('invalid source %s' % source_type)


class DescribeEC2(query.DescribeSource):

    def augment(self, resources):
        """EC2 API and AWOL Tags

        While ec2 api generally returns tags when doing describe_x on for
        various resources, it may also silently fail to do so unless a tag
        is used as a filter.

        See footnote on for official documentation.
        https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/Using_Tags.html#Using_Tags_CLI

        Apriori we may be using custodian to ensure tags (including
        name), so there isn't a good default to ensure that we will
        always get tags from describe_x calls.
        """
        # First if we're in event based lambda go ahead and skip this,
        # tags can't be trusted in ec2 instances immediately post creation.
        if not resources or self.manager.data.get(
                'mode', {}).get('type', '') in (
                    'cloudtrail', 'ec2-instance-state'):
            return resources

        # AWOL detector, so we don't make extraneous api calls.
        resource_count = len(resources)
        search_count = min(int(resource_count % 0.05) + 1, 5)
        if search_count > resource_count:
            search_count = resource_count
        found = False
        for r in random.sample(resources, search_count):
            if 'Tags' in r:
                found = True
                break

        if found:
            return resources

        # Okay go and do the tag lookup
        client = utils.local_session(self.manager.session_factory).client('ec2')
        tag_set = self.manager.retry(
            client.describe_tags,
            Filters=[{'Name': 'resource-type',
                      'Values': ['instance']}])['Tags']
        resource_tags = {}
        for t in tag_set:
            t.pop('ResourceType')
            rid = t.pop('ResourceId')
            resource_tags.setdefault(rid, []).append(t)

        m = self.manager.get_model()
        for r in resources:
            r['Tags'] = resource_tags.get(r[m.id], ())
        return resources


@filters.register('security-group')
class SecurityGroupFilter(net_filters.SecurityGroupFilter):

    RelatedIdsExpression = "SecurityGroups[].GroupId"


@filters.register('subnet')
class SubnetFilter(net_filters.SubnetFilter):

    RelatedIdsExpression = "SubnetId"


@filters.register('vpc')
class VpcFilter(net_filters.VpcFilter):

    RelatedIdsExpression = "VpcId"


@filters.register('check-permissions')
class ComputePermissions(CheckPermissions):

    def get_iam_arns(self, resources):
        profile_arn_map = {
            r['IamInstanceProfile']['Arn']: r['IamInstanceProfile']['Id']
            for r in resources if 'IamInstanceProfile' in r}

        # py2 compat on dict ordering
        profile_arns = list(profile_arn_map.items())
        profile_role_map = {
            arn: profile['Roles'][0]['Arn']
            for arn, profile in zip(
                [p[0] for p in profile_arns],
                self.manager.get_resource_manager(
                    'iam-profile').get_resources(
                        [p[1] for p in profile_arns]))}
        return [
            profile_role_map.get(r.get('IamInstanceProfile', {}).get('Arn'))
            for r in resources]


@filters.register('state-age')
class StateTransitionAge(AgeFilter):
    """Age an instance has been in the given state.

    .. code-block:: yaml

        policies:
          - name: ec2-state-running-7-days
            resource: ec2
            filters:
              - type: state-age
                op: ge
                days: 7
    """
    RE_PARSE_AGE = re.compile(r"\(.*?\)")

    # this filter doesn't use date_attribute, but needs to define it
    # to pass AgeFilter's validate method
    date_attribute = "dummy"

    schema = type_schema(
        'state-age',
        op={'$ref': '#/definitions/filters_common/comparison_operators'},
        days={'type': 'number'})

    def get_resource_date(self, i):
        v = i.get('StateTransitionReason')
        if not v:
            return None
        dates = self.RE_PARSE_AGE.findall(v)
        if dates:
            return parse(dates[0][1:-1])
        return None


class StateTransitionFilter(object):
    """Filter instances by state.

    Try to simplify construction for policy authors by automatically
    filtering elements (filters or actions) to the instances states
    they are valid for.

    For more details see
     https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-lifecycle.html

    """
    valid_origin_states = ()

    def filter_instance_state(self, instances, states=None):
        states = states or self.valid_origin_states
        orig_length = len(instances)
        results = [i for i in instances
                   if i['State']['Name'] in states]
        self.log.info("%s %d of %d instances" % (
            self.__class__.__name__, len(results), orig_length))
        return results


@filters.register('ebs')
class AttachedVolume(ValueFilter):
    """EC2 instances with EBS backed volume

    Filters EC2 instances with EBS backed storage devices (non ephemeral)

    :Example:

    .. code-block:: yaml

        policies:
          - name: ec2-encrypted-ebs-volumes
            resource: ec2
            filters:
              - type: ebs
                key: Encrypted
                value: true
    """

    schema = type_schema(
        'ebs', rinherit=ValueFilter.schema,
        **{'operator': {'enum': ['and', 'or']},
           'skip-devices': {'type': 'array', 'items': {'type': 'string'}}})
    schema_alias = False

    def get_permissions(self):
        return self.manager.get_resource_manager('ebs').get_permissions()

    def process(self, resources, event=None):
        self.volume_map = self.get_volume_mapping(resources)
        self.skip = self.data.get('skip-devices', [])
        self.operator = self.data.get(
            'operator', 'or') == 'or' and any or all
        return list(filter(self, resources))

    def get_volume_mapping(self, resources):
        volume_map = {}
        manager = self.manager.get_resource_manager('ebs')
        for instance_set in utils.chunks(resources, 200):
            volume_ids = []
            for i in instance_set:
                for bd in i.get('BlockDeviceMappings', ()):
                    if 'Ebs' not in bd:
                        continue
                    volume_ids.append(bd['Ebs']['VolumeId'])
            for v in manager.get_resources(volume_ids):
                if not v['Attachments']:
                    continue
                volume_map.setdefault(
                    v['Attachments'][0]['InstanceId'], []).append(v)
        return volume_map

    def __call__(self, i):
        volumes = self.volume_map.get(i['InstanceId'])
        if not volumes:
            return False
        if self.skip:
            for v in list(volumes):
                for a in v.get('Attachments', []):
                    if a['Device'] in self.skip:
                        volumes.remove(v)
        return self.operator(map(self.match, volumes))


@filters.register('termination-protected')
class DisableApiTermination(Filter):
    """EC2 instances with ``disableApiTermination`` attribute set

    Filters EC2 instances with ``disableApiTermination`` attribute set to true.

    :Example:

    .. code-block:: yaml

        policies:
          - name: termination-protection-enabled
            resource: ec2
            filters:
              - type: termination-protected

    :Example:

    .. code-block:: yaml

        policies:
          - name: termination-protection-NOT-enabled
            resource: ec2
            filters:
              - not:
                - type: termination-protected
    """

    schema = type_schema('termination-protected')
    permissions = ('ec2:DescribeInstanceAttribute',)

    def get_permissions(self):
        perms = list(self.permissions)
        perms.extend(self.manager.get_permissions())
        return perms

    def process(self, resources, event=None):
        client = utils.local_session(
            self.manager.session_factory).client('ec2')
        return [r for r in resources
                if self.is_termination_protection_enabled(client, r)]

    def is_termination_protection_enabled(self, client, inst):
        attr_val = self.manager.retry(
            client.describe_instance_attribute,
            Attribute='disableApiTermination',
            InstanceId=inst['InstanceId']
        )
        return attr_val['DisableApiTermination']['Value']


class InstanceImageBase(object):

    def prefetch_instance_images(self, instances):
        image_ids = [i['ImageId'] for i in instances if 'c7n:instance-image' not in i]
        self.image_map = self.get_local_image_mapping(image_ids)

    def get_base_image_mapping(self):
        return {i['ImageId']: i for i in
                self.manager.get_resource_manager('ami').resources()}

    def get_instance_image(self, instance):
        image = instance.get('c7n:instance-image', None)
        if not image:
            image = instance['c7n:instance-image'] = self.image_map.get(instance['ImageId'], None)
        return image

    def get_local_image_mapping(self, image_ids):
        base_image_map = self.get_base_image_mapping()
        resources = {i: base_image_map[i] for i in image_ids if i in base_image_map}
        missing = list(set(image_ids) - set(resources.keys()))
        if missing:
            loaded = self.manager.get_resource_manager('ami').get_resources(missing, False)
            resources.update({image['ImageId']: image for image in loaded})
        return resources


@filters.register('image-age')
class ImageAge(AgeFilter, InstanceImageBase):
    """EC2 AMI age filter

    Filters EC2 instances based on the age of their AMI image (in days)

    :Example:

    .. code-block:: yaml

        policies:
          - name: ec2-ancient-ami
            resource: ec2
            filters:
              - type: image-age
                op: ge
                days: 90
    """

    date_attribute = "CreationDate"

    schema = type_schema(
        'image-age',
        op={'$ref': '#/definitions/filters_common/comparison_operators'},
        days={'type': 'number'})

    def get_permissions(self):
        return self.manager.get_resource_manager('ami').get_permissions()

    def process(self, resources, event=None):
        self.prefetch_instance_images(resources)
        return super(ImageAge, self).process(resources, event)

    def get_resource_date(self, i):
        image = self.get_instance_image(i)
        if image:
            return parse(image['CreationDate'])
        else:
            return parse("2000-01-01T01:01:01.000Z")


@filters.register('image')
class InstanceImage(ValueFilter, InstanceImageBase):

    schema = type_schema('image', rinherit=ValueFilter.schema)
    schema_alias = False

    def get_permissions(self):
        return self.manager.get_resource_manager('ami').get_permissions()

    def process(self, resources, event=None):
        self.prefetch_instance_images(resources)
        return super(InstanceImage, self).process(resources, event)

    def __call__(self, i):
        image = self.get_instance_image(i)
        # Finally, if we have no image...
        if not image:
            self.log.warning(
                "Could not locate image for instance:%s ami:%s" % (
                    i['InstanceId'], i["ImageId"]))
            # Match instead on empty skeleton?
            return False
        return self.match(image)


@filters.register('offhour')
class InstanceOffHour(OffHour, StateTransitionFilter):
    """Custodian OffHour filter

    Filters running EC2 instances with the intent to stop at a given hour of
    the day. A list of days to excluded can be included as a list of strings
    with the format YYYY-MM-DD. Alternatively, the list (using the same syntax)
    can be taken from a specified url.

    :Example:

    .. code-block:: yaml

        policies:
          - name: offhour-evening-stop
            resource: ec2
            filters:
              - type: offhour
                tag: custodian_downtime
                default_tz: et
                offhour: 20
            actions:
              - stop

          - name: offhour-evening-stop-skip-holidays
            resource: ec2
            filters:
              - type: offhour
                tag: custodian_downtime
                default_tz: et
                offhour: 20
                skip-days: ['2017-12-25']
            actions:
              - stop

          - name: offhour-evening-stop-skip-holidays-from
            resource: ec2
            filters:
              - type: offhour
                tag: custodian_downtime
                default_tz: et
                offhour: 20
                skip-days-from:
                  expr: 0
                  format: csv
                  url: 's3://location/holidays.csv'
            actions:
              - stop
    """

    valid_origin_states = ('running',)

    def process(self, resources, event=None):
        return super(InstanceOffHour, self).process(
            self.filter_instance_state(resources))


@filters.register('network-location')
class EC2NetworkLocation(net_filters.NetworkLocation, StateTransitionFilter):

    valid_origin_states = ('pending', 'running', 'shutting-down', 'stopping',
                           'stopped')

    def process(self, resources, event=None):
        resources = self.filter_instance_state(resources)
        if not resources:
            return []
        return super(EC2NetworkLocation, self).process(resources)


@filters.register('onhour')
class InstanceOnHour(OnHour, StateTransitionFilter):
    """Custodian OnHour filter

    Filters stopped EC2 instances with the intent to start at a given hour of
    the day. A list of days to excluded can be included as a list of strings
    with the format YYYY-MM-DD. Alternatively, the list (using the same syntax)
    can be taken from a specified url.

    :Example:

    .. code-block:: yaml

        policies:
          - name: onhour-morning-start
            resource: ec2
            filters:
              - type: onhour
                tag: custodian_downtime
                default_tz: et
                onhour: 6
            actions:
              - start

          - name: onhour-morning-start-skip-holidays
            resource: ec2
            filters:
              - type: onhour
                tag: custodian_downtime
                default_tz: et
                onhour: 6
                skip-days: ['2017-12-25']
            actions:
              - start

          - name: onhour-morning-start-skip-holidays-from
            resource: ec2
            filters:
              - type: onhour
                tag: custodian_downtime
                default_tz: et
                onhour: 6
                skip-days-from:
                  expr: 0
                  format: csv
                  url: 's3://location/holidays.csv'
            actions:
              - start
    """

    valid_origin_states = ('stopped',)

    def process(self, resources, event=None):
        return super(InstanceOnHour, self).process(
            self.filter_instance_state(resources))


@filters.register('ephemeral')
class EphemeralInstanceFilter(Filter):
    """EC2 instances with ephemeral storage

    Filters EC2 instances that have ephemeral storage (an instance-store backed
    root device)

    :Example:

    .. code-block:: yaml

        policies:
          - name: ec2-ephemeral-instances
            resource: ec2
            filters:
              - type: ephemeral

    http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/InstanceStorage.html
    """

    schema = type_schema('ephemeral')

    def __call__(self, i):
        return self.is_ephemeral(i)

    @staticmethod
    def is_ephemeral(i):
        for bd in i.get('BlockDeviceMappings', []):
            if bd['DeviceName'] in ('/dev/sda1', '/dev/xvda', 'xvda'):
                if 'Ebs' in bd:
                    return False
                return True
        return True


@filters.register('instance-uptime')
class UpTimeFilter(AgeFilter):

    date_attribute = "LaunchTime"

    schema = type_schema(
        'instance-uptime',
        op={'$ref': '#/definitions/filters_common/comparison_operators'},
        days={'type': 'number'})


@filters.register('instance-age')
class InstanceAgeFilter(AgeFilter):
    """Filters instances based on their age (in days)

    :Example:

    .. code-block:: yaml

        policies:
          - name: ec2-30-days-plus
            resource: ec2
            filters:
              - type: instance-age
                op: ge
                days: 30
    """

    date_attribute = "LaunchTime"
    ebs_key_func = operator.itemgetter('AttachTime')

    schema = type_schema(
        'instance-age',
        op={'$ref': '#/definitions/filters_common/comparison_operators'},
        days={'type': 'number'},
        hours={'type': 'number'},
        minutes={'type': 'number'})

    def get_resource_date(self, i):
        # LaunchTime is basically how long has the instance
        # been on, use the oldest ebs vol attach time
        ebs_vols = [
            block['Ebs'] for block in i['BlockDeviceMappings']
            if 'Ebs' in block]
        if not ebs_vols:
            # Fall back to using age attribute (ephemeral instances)
            return super(InstanceAgeFilter, self).get_resource_date(i)
        # Lexographical sort on date
        ebs_vols = sorted(ebs_vols, key=self.ebs_key_func)
        return ebs_vols[0]['AttachTime']


@filters.register('default-vpc')
class DefaultVpc(DefaultVpcBase):
    """ Matches if an ec2 database is in the default vpc
    """

    schema = type_schema('default-vpc')

    def __call__(self, ec2):
        return ec2.get('VpcId') and self.match(ec2.get('VpcId')) or False


def deserialize_user_data(user_data):
    data = base64.b64decode(user_data)
    # try raw and compressed
    try:
        return data.decode('utf8')
    except UnicodeDecodeError:
        return zlib.decompress(data, 16).decode('utf8')


@filters.register('user-data')
class UserData(ValueFilter):
    """Filter on EC2 instances which have matching userdata.
    Note: It is highly recommended to use regexes with the ?sm flags, since Custodian
    uses re.match() and userdata spans multiple lines.

        :example:

        .. code-block:: yaml

            policies:
              - name: ec2_userdata_stop
                resource: ec2
                filters:
                  - type: user-data
                    op: regex
                    value: (?smi).*password=
                actions:
                  - stop
    """

    schema = type_schema('user-data', rinherit=ValueFilter.schema)
    schema_alias = False
    batch_size = 50
    annotation = 'c7n:user-data'
    permissions = ('ec2:DescribeInstanceAttribute',)

    def __init__(self, data, manager):
        super(UserData, self).__init__(data, manager)
        self.data['key'] = '"c7n:user-data"'

    def process(self, resources, event=None):
        client = utils.local_session(self.manager.session_factory).client('ec2')
        results = []
        with self.executor_factory(max_workers=3) as w:
            futures = {}
            for instance_set in utils.chunks(resources, self.batch_size):
                futures[w.submit(
                    self.process_instance_set,
                    client, instance_set)] = instance_set

            for f in as_completed(futures):
                if f.exception():
                    self.log.error(
                        "Error processing userdata on instance set %s", f.exception())
                results.extend(f.result())
        return results

    def process_instance_set(self, client, resources):
        results = []
        for r in resources:
            if self.annotation not in r:
                try:
                    result = client.describe_instance_attribute(
                        Attribute='userData',
                        InstanceId=r['InstanceId'])
                except ClientError as e:
                    if e.response['Error']['Code'] == 'InvalidInstanceId.NotFound':
                        continue
                if 'Value' not in result['UserData']:
                    r[self.annotation] = None
                else:
                    r[self.annotation] = deserialize_user_data(
                        result['UserData']['Value'])
            if self.match(r):
                results.append(r)
        return results


@filters.register('singleton')
class SingletonFilter(Filter, StateTransitionFilter):
    """EC2 instances without autoscaling or a recover alarm

    Filters EC2 instances that are not members of an autoscaling group
    and do not have Cloudwatch recover alarms.

    :Example:

    .. code-block:: yaml

        policies:
          - name: ec2-recover-instances
            resource: ec2
            filters:
              - singleton
            actions:
              - type: tag
                key: problem
                value: instance is not resilient

    https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-recover.html
    """

    schema = type_schema('singleton')

    permissions = ('cloudwatch:DescribeAlarmsForMetric',)

    valid_origin_states = ('running', 'stopped', 'pending', 'stopping')

    in_asg = ValueFilter({
        'key': 'tag:aws:autoscaling:groupName',
        'value': 'not-null'}).validate()

    def process(self, instances, event=None):
        return super(SingletonFilter, self).process(
            self.filter_instance_state(instances))

    def __call__(self, i):
        if self.in_asg(i):
            return False
        else:
            return not self.has_recover_alarm(i)

    def has_recover_alarm(self, i):
        client = utils.local_session(self.manager.session_factory).client('cloudwatch')
        alarms = client.describe_alarms_for_metric(
            MetricName='StatusCheckFailed_System',
            Namespace='AWS/EC2',
            Dimensions=[
                {
                    'Name': 'InstanceId',
                    'Value': i['InstanceId']
                }
            ]
        )

        for i in alarms['MetricAlarms']:
            for a in i['AlarmActions']:
                if (
                    a.startswith('arn:aws:automate:') and
                    a.endswith(':ec2:recover')
                ):
                    return True

        return False


@EC2.filter_registry.register('ssm')
class SsmStatus(ValueFilter):
    """Filter ec2 instances by their ssm status information.

    :Example:

    Find ubuntu 18.04 instances are active with ssm.

    .. code-block:: yaml

        policies:
          - name: ec2-ssm-check
            resource: ec2
            filters:
              - type: ssm
                key: PingStatus
                value: Online
              - type: ssm
                key: PlatformName
                value: Ubuntu
              - type: ssm
                key: PlatformVersion
                value: 18.04
    """
    schema = type_schema('ssm', rinherit=ValueFilter.schema)
    schema_alias = False
    permissions = ('ssm:DescribeInstanceInformation',)
    annotation = 'c7n:SsmState'

    def process(self, resources, event=None):