#!/usr/bin/python
#
# This is a free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This Ansible library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this library.  If not, see <http://www.gnu.org/licenses/>.

DOCUMENTATION = '''
---
module: ec2_eni
short_description: Create and optionally attach an Elastic Network Interface (ENI) to an instance
description:
    - Create and optionally attach an Elastic Network Interface (ENI) to an
      instance. If an ENI ID is provided, an attempt is made to update the
      existing ENI. By passing state=detached, an ENI can be detached from its instance.
version_added: "2.0"
author: "Rob White (@wimnat)"
options:
  eni_id:
    description:
      - The ID of the ENI
    required: false
    default: null
  instance_id:
    description:
      - Instance ID that you wish to attach ENI to, if None the new ENI will be
        created in detached state, existing ENI will keep current attachment state.
    required: false
    default: null
  private_ip_address:
    description:
      - Private IP address.
    required: false
    default: null
  subnet_id:
    description:
      - ID of subnet in which to create the ENI. Only required when state=present.
    required: true
  description:
    description:
      - Optional description of the ENI.
    required: false
    default: null
  security_groups:
    description:
      - List of security groups associated with the interface. Only used when state=present. Since version 2.2, you \
      can specify security groups by ID or by name or a combination of both. Prior to 2.2, you can specify only by ID.
    required: false
    default: null
  state:
    description:
      - Create, delete or detach ENI from its instance.
    required: false
    default: present
    choices: [ 'present', 'absent', 'detached' ]
  device_index:
    description:
      - The index of the device for the network interface attachment on the instance.
    required: false
    default: 0
  force_detach:
    description:
      - Force detachment of the interface. This applies either when explicitly detaching the interface by setting instance_id to None or when deleting an interface with state=absent.
    required: false
    default: no
  delete_on_termination:
    description:
      - Delete the interface when the instance it is attached to is terminated. You can only specify this flag when the interface is being modified, not on creation.
    required: false
  source_dest_check:
    description:
      - By default, interfaces perform source/destination checks. NAT instances however need this check to be disabled. You can only specify this flag when the interface is being modified, not on creation.
    required: false
  secondary_private_ip_addresses:
    description:
      - A list of IP addresses to assign as secondary IP addresses to the network interface. This option is mutually exclusive of secondary_private_ip_address_count
    required: false
    version_added: 2.2
  secondary_private_ip_address_count:
    description:
      - The number of secondary IP addresses to assign to the network interface. This option is mutually exclusive of secondary_private_ip_addresses
    required: false
    version_added: 2.2
extends_documentation_fragment:
    - aws
    - ec2
'''

EXAMPLES = '''
# Note: These examples do not set authentication details, see the AWS Guide for details.

# Create an ENI. As no security group is defined, ENI will be created in default security group
- ec2_eni:
    private_ip_address: 172.31.0.20
    subnet_id: subnet-xxxxxxxx
    state: present

# Create an ENI and attach it to an instance
- ec2_eni:
    instance_id: i-xxxxxxx
    device_index: 1
    private_ip_address: 172.31.0.20
    subnet_id: subnet-xxxxxxxx
    state: present

# Create an ENI with two secondary addresses
- ec2_eni:
    subnet_id: subnet-xxxxxxxx
    state: present
    secondary_private_ip_address_count: 2

# Assign a secondary IP address to an existing ENI
# This will purge any existing IPs
- ec2_eni:
    subnet_id: subnet-xxxxxxxx
    eni_id: eni-yyyyyyyy
    state: present
    secondary_private_ip_addresses:
      - 172.16.1.1

# Remove any secondary IP addresses from an existing ENI
- ec2_eni:
    subnet_id: subnet-xxxxxxxx
    eni_id: eni-yyyyyyyy
    state: present
    secondary_private_ip_addresses:
      -

# Destroy an ENI, detaching it from any instance if necessary
- ec2_eni:
    eni_id: eni-xxxxxxx
    force_detach: yes
    state: absent

# Update an ENI
- ec2_eni:
    eni_id: eni-xxxxxxx
    description: "My new description"
    state: present

# Detach an ENI from an instance
- ec2_eni:
    eni_id: eni-xxxxxxx
    instance_id: None
    state: present

### Delete an interface on termination
# First create the interface
- ec2_eni:
    instance_id: i-xxxxxxx
    device_index: 1
    private_ip_address: 172.31.0.20
    subnet_id: subnet-xxxxxxxx
    state: present
  register: eni

# Modify the interface to enable the delete_on_terminaton flag
- ec2_eni:
    eni_id: {{ "eni.interface.id" }}
    delete_on_termination: true

'''

import time
import re

try:
    import boto.ec2
    import boto.vpc
    from boto.exception import BotoServerError
    HAS_BOTO = True
except ImportError:
    HAS_BOTO = False


def get_eni_info(interface):

    # Private addresses
    private_addresses = []
    for ip in interface.private_ip_addresses:
        private_addresses.append({ 'private_ip_address': ip.private_ip_address, 'primary_address': ip.primary })

    interface_info = {'id': interface.id,
                      'subnet_id': interface.subnet_id,
                      'vpc_id': interface.vpc_id,
                      'description': interface.description,
                      'owner_id': interface.owner_id,
                      'status': interface.status,
                      'mac_address': interface.mac_address,
                      'private_ip_address': interface.private_ip_address,
                      'source_dest_check': interface.source_dest_check,
                      'groups': dict((group.id, group.name) for group in interface.groups),
                      'private_ip_addresses': private_addresses
                      }

    if interface.attachment is not None:
        interface_info['attachment'] = {'attachment_id': interface.attachment.id,
                                        'instance_id': interface.attachment.instance_id,
                                        'device_index': interface.attachment.device_index,
                                        'status': interface.attachment.status,
                                        'attach_time': interface.attachment.attach_time,
                                        'delete_on_termination': interface.attachment.delete_on_termination,
                                        }

    return interface_info


def wait_for_eni(eni, status):

    while True:
        time.sleep(3)
        eni.update()
        # If the status is detached we just need attachment to disappear
        if eni.attachment is None:
            if status == "detached":
                break
        else:
            if status == "attached" and eni.attachment.status == "attached":
                break


def create_eni(connection, vpc_id, module):

    instance_id = module.params.get("instance_id")
    if instance_id == 'None':
        instance_id = None
    device_index = module.params.get("device_index")
    subnet_id = module.params.get('subnet_id')
    private_ip_address = module.params.get('private_ip_address')
    description = module.params.get('description')
    security_groups = get_ec2_security_group_ids_from_names(module.params.get('security_groups'), connection, vpc_id=vpc_id, boto3=False)
    secondary_private_ip_addresses = module.params.get("secondary_private_ip_addresses")
    secondary_private_ip_address_count = module.params.get("secondary_private_ip_address_count")
    changed = False

    try:
        eni = find_eni(connection, module)
        if eni is None:
            eni = connection.create_network_interface(subnet_id, private_ip_address, description, security_groups)
            if instance_id is not None:
                try:
                    eni.attach(instance_id, device_index)
                except BotoServerError:
                    eni.delete()
                    raise
                # Wait to allow creation / attachment to finish
                wait_for_eni(eni, "attached")
                eni.update()

            if secondary_private_ip_address_count is not None:
                try:
                    connection.assign_private_ip_addresses(network_interface_id=eni.id, secondary_private_ip_address_count=secondary_private_ip_address_count)
                except BotoServerError:
                    eni.delete()
                    raise

            if secondary_private_ip_addresses is not None:
                try:
                    connection.assign_private_ip_addresses(network_interface_id=eni.id, private_ip_addresses=secondary_private_ip_addresses)
                except BotoServerError:
                    eni.delete()
                    raise

            changed = True

    except BotoServerError as e:
        module.fail_json(msg=e.message)

    module.exit_json(changed=changed, interface=get_eni_info(eni))


def modify_eni(connection, vpc_id, module, eni):

    instance_id = module.params.get("instance_id")
    do_detach = module.params.get('state') == 'detached'
    device_index = module.params.get("device_index")
    description = module.params.get('description')
    security_groups = module.params.get('security_groups')
    force_detach = module.params.get("force_detach")
    source_dest_check = module.params.get("source_dest_check")
    delete_on_termination = module.params.get("delete_on_termination")
    secondary_private_ip_addresses = module.params.get("secondary_private_ip_addresses")
    secondary_private_ip_address_count = module.params.get("secondary_private_ip_address_count")
    changed = False

    try:
        if description is not None:
            if eni.description != description:
                connection.modify_network_interface_attribute(eni.id, "description", description)
                changed = True
        if len(security_groups) > 0:
            groups = get_ec2_security_group_ids_from_names(security_groups, connection, vpc_id=vpc_id, boto3=False)
            if sorted(get_sec_group_list(eni.groups)) != sorted(groups):
                connection.modify_network_interface_attribute(eni.id, "groupSet", groups)
                changed = True
        if source_dest_check is not None:
            if eni.source_dest_check != source_dest_check:
                connection.modify_network_interface_attribute(eni.id, "sourceDestCheck", source_dest_check)
                changed = True
        if delete_on_termination is not None and eni.attachment is not None:
            if eni.attachment.delete_on_termination is not delete_on_termination:
                connection.modify_network_interface_attribute(eni.id, "deleteOnTermination", delete_on_termination, eni.attachment.id)
                changed = True

        current_secondary_addresses = [i.private_ip_address for i in eni.private_ip_addresses if not i.primary]
        if secondary_private_ip_addresses is not None:
            secondary_addresses_to_remove = list(set(current_secondary_addresses) - set(secondary_private_ip_addresses))
            if secondary_addresses_to_remove:
                connection.unassign_private_ip_addresses(network_interface_id=eni.id, private_ip_addresses=list(set(current_secondary_addresses) - set(secondary_private_ip_addresses)), dry_run=False)
            connection.assign_private_ip_addresses(network_interface_id=eni.id, private_ip_addresses=secondary_private_ip_addresses, secondary_private_ip_address_count=None, allow_reassignment=False, dry_run=False)
        if secondary_private_ip_address_count is not None:
            current_secondary_address_count = len(current_secondary_addresses)

            if secondary_private_ip_address_count > current_secondary_address_count:
                connection.assign_private_ip_addresses(network_interface_id=eni.id, private_ip_addresses=None, secondary_private_ip_address_count=(secondary_private_ip_address_count - current_secondary_address_count), allow_reassignment=False, dry_run=False)
                changed = True
            elif secondary_private_ip_address_count < current_secondary_address_count:
                # How many of these addresses do we want to remove
                secondary_addresses_to_remove_count = current_secondary_address_count - secondary_private_ip_address_count
                connection.unassign_private_ip_addresses(network_interface_id=eni.id, private_ip_addresses=current_secondary_addresses[:secondary_addresses_to_remove_count], dry_run=False)

        if instance_id is not None:
            eni.attach(instance_id, device_index)
            wait_for_eni(eni, "attached")
            changed = True

    except BotoServerError as e:
        module.fail_json(msg=e.message)

    eni.update()
    module.exit_json(changed=changed, interface=get_eni_info(eni))


def delete_eni(connection, module):

    eni_id = module.params.get("eni_id")
    force_detach = module.params.get("force_detach")

    try:
        eni_result_set = connection.get_all_network_interfaces(eni_id)
        eni = eni_result_set[0]

        if force_detach is True:
            if eni.attachment is not None:
                eni.detach(force_detach)
                # Wait to allow detachment to finish
                wait_for_eni(eni, "detached")
                eni.update()
            eni.delete()
            changed = True
        else:
            eni.delete()
            changed = True

        module.exit_json(changed=changed)
    except BotoServerError as e:
        regex = re.compile('The networkInterface ID \'.*\' does not exist')
        if regex.search(e.message) is not None:
            module.exit_json(changed=False)
        else:
            module.fail_json(msg=e.message)


def detach_eni(connection, module):

    eni = find_eni(connection, module)
    if eni.attachment is not None:
        eni.detach(force_detach)
        wait_for_eni(eni, "detached")
        eni.update()
        module.exit_json(changed=True, interface=get_eni_info(eni))
    else:
        module.exit_json(changed=False, interface=get_eni_info(eni))


def find_eni(connection, module):

    eni_id = module.params.get("eni_id")
    subnet_id = module.params.get('subnet_id')
    private_ip_address = module.params.get('private_ip_address')

    try:
        filters = {}
        if private_ip_address:
            filters['private-ip-address'] = private_ip_address
        if subnet_id:
            filters['subnet-id'] = subnet_id

        eni_result = connection.get_all_network_interfaces(eni_id, filters=filters)
        if len(eni_result) > 0:
            return eni_result[0]
        else:
            return None

    except BotoServerError as e:
        module.fail_json(msg=e.message)

    return None


def get_sec_group_list(groups):

    # Build list of remote security groups
    remote_security_groups = []
    for group in groups:
        remote_security_groups.append(group.id.encode())

    return remote_security_groups


def _get_vpc_id(connection, module, subnet_id):

    try:
        return connection.get_all_subnets(subnet_ids=[subnet_id])[0].vpc_id
    except BotoServerError as e:
        module.fail_json(msg=e.message)


def main():
    argument_spec = ec2_argument_spec()
    argument_spec.update(
        dict(
            eni_id=dict(default=None, type='str'),
            instance_id=dict(default=None, type='str'),
            private_ip_address=dict(type='str'),
            subnet_id=dict(type='str'),
            description=dict(type='str'),
            security_groups=dict(default=[], type='list'),
            device_index=dict(default=0, type='int'),
            state=dict(default='present', choices=['present', 'absent', 'detached']),
            force_detach=dict(default='no', type='bool'),
            source_dest_check=dict(default=None, type='bool'),
            delete_on_termination=dict(default=None, type='bool'),
            secondary_private_ip_addresses=dict(default=None, type='list'),
            secondary_private_ip_address_count=dict(default=None, type='int')
        )
    )

    module = AnsibleModule(argument_spec=argument_spec,
                           mutually_exclusive=[
                               ['secondary_private_ip_addresses', 'secondary_private_ip_address_count']
                            ],
                           required_if=([
                               ('state', 'present', ['subnet_id']),
                               ('state', 'absent', ['eni_id']),
                               ('state', 'detached', ['eni_id']),
                            ])
                           )

    if not HAS_BOTO:
        module.fail_json(msg='boto required for this module')

    region, ec2_url, aws_connect_params = get_aws_connection_info(module)

    if region:
        try:
            connection = connect_to_aws(boto.ec2, region, **aws_connect_params)
            vpc_connection = connect_to_aws(boto.vpc, region, **aws_connect_params)
        except (boto.exception.NoAuthHandlerFound, AnsibleAWSError), e:
            module.fail_json(msg=str(e))
    else:
        module.fail_json(msg="region must be specified")

    state = module.params.get("state")
    eni_id = module.params.get("eni_id")
    private_ip_address = module.params.get('private_ip_address')

    if state == 'present':
        subnet_id = module.params.get("subnet_id")
        vpc_id = _get_vpc_id(vpc_connection, module, subnet_id)

        eni = find_eni(connection, module)
        if eni is None:
            create_eni(connection, vpc_id, module)
        else:
            modify_eni(connection, vpc_id, module, eni)

    elif state == 'absent':
        delete_eni(connection, module)

    elif state == 'detached':
        detach_eni(connection, module)

from ansible.module_utils.basic import *
from ansible.module_utils.ec2 import *

if __name__ == '__main__':
    main()