#!/usr/bin/python
# This file is part of Ansible
#
# Ansible is 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.
#
# Ansible 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 Ansible.  If not, see <http://www.gnu.org/licenses/>.

DOCUMENTATION = """
---
module: ec2_elb_lb
short_description: Creates or destroys Amazon ELB.
  - Returns information about the load balancer.
  - Will be marked changed when called only if state is changed.
version_added: "1.4"
requirements: [ "boto" ]
author: Jim Dalton
options:
  state:
    description:
      - Create or destroy the ELB
    required: true
  name:
    description:
      - The name of the ELB
    required: true
  listeners:
    description:
      - List of ports/protocols for this ELB to listen on (see example)
    required: false
  purge_listeners:
    description:
      - Purge existing listeners on ELB that are not found in listeners
    required: false
    default: true
  zones:
    description:
      - List of availability zones to enable on this ELB
    required: false
  purge_zones:
    description:
      - Purge existing availability zones on ELB that are not found in zones
    required: false
    default: false
  health_check:
    description:
      - An associative array of health check configuration settigs (see example)
    require: false
    default: None
  aws_secret_key:
    description:
      - AWS secret key. If not set then the value of the AWS_SECRET_KEY environment variable is used. 
    required: false
    default: None
    aliases: ['ec2_secret_key', 'secret_key']
  aws_access_key:
    description:
      - AWS access key. If not set then the value of the AWS_ACCESS_KEY environment variable is used.
    required: false
    default: None
    aliases: ['ec2_access_key', 'access_key']
  region:
    description:
      - The AWS region to use. If not specified then the value of the EC2_REGION environment variable, if any, is used.
    required: false
    aliases: ['aws_region', 'ec2_region']

"""

EXAMPLES = """
# Note: None of these examples set aws_access_key, aws_secret_key, or region.
# It is assumed that their matching environment variables are set.

# Basic provisioning example
- local_action:
    module: ec2_elb_lb
    name: "test-please-delete"
    state: present
    zones:
      - us-east-1a
      - us-east-1d
    listeners:
      - protocol: http # options are http, https, ssl, tcp
        load_balancer_port: 80
        instance_port: 80
      - protocol: https
        load_balancer_port: 443
        instance_protocol: http # optional, defaults to value of protocol setting
        instance_port: 80
        # ssl certificate required for https or ssl
        ssl_certificate_id: "arn:aws:iam::123456789012:server-certificate/company/servercerts/ProdServerCert"

# Configure a health check
- local_action:
    module: ec2_elb_lb
    name: "test-please-delete"
    state: present
    zones:
      - us-east-1d
    listeners:
      - protocol: http
        load_balancer_port: 80
        instance_port: 80
    health_check:
        ping_protocol: http # options are http, https, ssl, tcp
        ping_port: 80
        ping_path: "/index.html" # not required for tcp or ssl
        response_timeout: 5 # seconds
        interval: 30 # seconds
        unhealthy_threshold: 2
        healthy_threshold: 10

# Ensure ELB is gone
- local_action:
    module: ec2_elb_lb
    name: "test-please-delete"
    state: absent

# Normally, this module will purge any listeners that exist on the ELB
# but aren't specified in the listeners parameter. If purge_listeners is
# false it leaves them alone
- local_action:
    module: ec2_elb_lb
    name: "test-please-delete"
    state: present
    zones:
      - us-east-1a
      - us-east-1d
    listeners:
      - protocol: http
        load_balancer_port: 80
        instance_port: 80
    purge_listeners: no

# Normally, this module will leave availability zones that are enabled
# on the ELB alone. If purge_zones is true, then any extreneous zones
# will be removed
- local_action:
    module: ec2_elb_lb
    name: "test-please-delete"
    state: present
    zones:
      - us-east-1a
      - us-east-1d
    listeners:
      - protocol: http
        load_balancer_port: 80
        instance_port: 80
    purge_zones: yes
"""

import sys
import os

AWS_REGIONS = ['ap-northeast-1',
               'ap-southeast-1',
               'ap-southeast-2',
               'eu-west-1',
               'sa-east-1',
               'us-east-1',
               'us-west-1',
               'us-west-2']

try:
    import boto
    import boto.ec2.elb
    from boto.ec2.elb.healthcheck import HealthCheck
    from boto.regioninfo import RegionInfo
except ImportError:
    print "failed=True msg='boto required for this module'"
    sys.exit(1)


class ElbManager(object):
    """Handles ELB creation and destruction"""

    def __init__(self, module, name, listeners=None, purge_listeners=None,
                 zones=None, purge_zones=None, health_check=None,
                 aws_access_key=None, aws_secret_key=None, region=None):
        self.module = module
        self.name = name
        self.listeners = listeners
        self.purge_listeners = purge_listeners
        self.zones = zones
        self.purge_zones = purge_zones
        self.health_check = health_check

        self.aws_access_key = aws_access_key
        self.aws_secret_key = aws_secret_key
        self.region = region

        self.changed = False
        self.status = 'gone'
        self.elb_conn = self._get_elb_connection()
        self.elb = self._get_elb()

    def ensure_ok(self):
        """Create the ELB"""
        if not self.elb:
            # Zones and listeners will be added at creation
            self._create_elb()
        else:
            self._set_zones()
            self._set_elb_listeners()
        self._set_health_check()

    def ensure_gone(self):
        """Destroy the ELB"""
        if self.elb:
            self._delete_elb()

    def get_info(self):
        if not self.elb:
            info = {
                'name': self.name,
                'status': self.status
            }
        else:
            info = {
                'name': self.elb.name,
                'dns_name': self.elb.dns_name,
                'zones': self.elb.availability_zones,
                'status': self.status
            }

            if self.elb.health_check:
                info['health_check'] = {
                    'target': self.elb.health_check.target,
                    'interval': self.elb.health_check.interval,
                    'timeout': self.elb.health_check.timeout,
                    'healthy_threshold': self.elb.health_check.healthy_threshold,
                    'unhealthy_threshold': self.elb.health_check.unhealthy_threshold,
                }

            if self.elb.listeners:
                info['listeners'] = [l.get_complex_tuple()
                                     for l in self.elb.listeners]
            elif self.status == 'created':
                # When creating a new ELB, listeners don't show in the
                # immediately returned result, so just include the
                # ones that were added
                info['listeners'] = [self._listener_as_tuple(l)
                                     for l in self.listeners]
            else:
                info['listeners'] = []
        return info

    def _get_elb(self):
        elbs = self.elb_conn.get_all_load_balancers()
        for elb in elbs:
            if self.name == elb.name:
                self.status = 'ok'
                return elb

    def _get_elb_connection(self):
        try:
            endpoint = "elasticloadbalancing.%s.amazonaws.com" % self.region
            connect_region = RegionInfo(name=self.region, endpoint=endpoint)
            return boto.ec2.elb.ELBConnection(self.aws_access_key,
                                              self.aws_secret_key,
                                              region=connect_region)
        except boto.exception.NoAuthHandlerFound, e:
            self.module.fail_json(msg=str(e))

    def _delete_elb(self):
        # True if succeeds, exception raised if not
        result = self.elb_conn.delete_load_balancer(name=self.name)
        if result:
            self.changed = True
            self.status = 'deleted'

    def _create_elb(self):
        listeners = [self._listener_as_tuple(l) for l in self.listeners]
        self.elb = self.elb_conn.create_load_balancer(name=self.name,
                                                      zones=self.zones,
                                                      complex_listeners=listeners)
        if self.elb:
            self.changed = True
            self.status = 'created'

    def _create_elb_listeners(self, listeners):
        """Takes a list of listener tuples and creates them"""
        # True if succeeds, exception raised if not
        self.changed = self.elb_conn.create_load_balancer_listeners(self.name,
                                                                    complex_listeners=listeners)

    def _delete_elb_listeners(self, listeners):
        """Takes a list of listener tuples and deletes them from the elb"""
        ports = [l[0] for l in listeners]

        # True if succeeds, exception raised if not
        self.changed = self.elb_conn.delete_load_balancer_listeners(self.name,
                                                                    ports)

    def _set_elb_listeners(self):
        """
        Creates listeners specified by self.listeners; overwrites existing
        listeners on these ports; removes extraneous listeners
        """
        listeners_to_add = []
        listeners_to_remove = []
        listeners_to_keep = []

        # Check for any listeners we need to create or overwrite
        for listener in self.listeners:
            listener_as_tuple = self._listener_as_tuple(listener)

            # First we loop through existing listeners to see if one is
            # already specified for this port
            existing_listener_found = None
            for existing_listener in self.elb.listeners:
                # Since ELB allows only one listener on each incoming port, a
                # single match on the incomping port is all we're looking for
                if existing_listener[0] == listener['load_balancer_port']:
                    existing_listener_found = existing_listener.get_complex_tuple()
                    break

            if existing_listener_found:
                # Does it match exactly?
                if listener_as_tuple != existing_listener_found:
                    # The ports are the same but something else is different,
                    # so we'll remove the exsiting one and add the new one
                    listeners_to_remove.append(existing_listener_found)
                    listeners_to_add.append(listener_as_tuple)
                else:
                    # We already have this listener, so we're going to keep it
                    listeners_to_keep.append(existing_listener_found)
            else:
                # We didn't find an existing listener, so just add the new one
                listeners_to_add.append(listener_as_tuple)

        # Check for any extraneous listeners we need to remove, if desired
        if self.purge_listeners:
            for existing_listener in self.elb.listeners:
                existing_listener_tuple = existing_listener.get_complex_tuple()
                if existing_listener_tuple in listeners_to_remove:
                    # Already queued for removal
                    continue
                if existing_listener_tuple in listeners_to_keep:
                    # Keep this one around
                    continue
                # Since we're not already removing it and we don't need to keep
                # it, let's get rid of it
                listeners_to_remove.append(existing_listener_tuple)

        if listeners_to_remove:
            self._delete_elb_listeners(listeners_to_remove)

        if listeners_to_add:
            self._create_elb_listeners(listeners_to_add)

    def _listener_as_tuple(self, listener):
        """Formats listener as a 4- or 5-tuples, in the order specified by the
        ELB API"""
        # N.B. string manipulations on protocols below (str(), upper()) is to
        # ensure format matches output from ELB API
        listener_list = [
            listener['load_balancer_port'],
            listener['instance_port'],
            str(listener['protocol'].upper()),
        ]

        # Instance protocol is not required by ELB API; it defaults to match
        # load balancer protocol. We'll mimic that behavior here
        if 'instance_protocol' in listener:
            listener_list.append(str(listener['instance_protocol'].upper()))
        else:
            listener_list.append(str(listener['protocol'].upper()))

        if 'ssl_certificate_id' in listener:
            listener_list.append(str(listener['ssl_certificate_id']))

        return tuple(listener_list)

    def _enable_zones(self, zones):
        self.elb_conn.enable_availability_zones(self.name, zones)
        self.changed = True

    def _disable_zones(self, zones):
        self.elb_conn.disable_availability_zones(self.name, zones)
        self.changed = True

    def _set_zones(self):
        """Determine which zones need to be enabled or disabled on the ELB"""
        if self.purge_zones:
            zones_to_disable = list(set(self.elb.availability_zones) -
                                    set(self.zones))
            zones_to_enable = list(set(self.zones) -
                                   set(self.elb.availability_zones))
        else:
            zones_to_disable = None
            zones_to_enable = list(set(self.zones) -
                                   set(self.elb.availability_zones))
        if zones_to_enable:
            self._enable_zones(zones_to_enable)
        # N.B. This must come second, in case it would have removed all zones
        if zones_to_disable:
            self._disable_zones(zones_to_disable)

    def _set_health_check(self):
        """Set health check values on ELB as needed"""
        if self.health_check:
            # This just makes it easier to compare each of the attributes
            # and look for changes. Keys are attributes of the current
            # health_check; values are desired values of new health_check
            health_check_config = {
                "target": self._get_health_check_target(),
                "timeout": self.health_check['response_timeout'],
                "interval": self.health_check['interval'],
                "unhealthy_threshold": self.health_check['unhealthy_threshold'],
                "healthy_threshold": self.health_check['healthy_threshold'],
            }

            update_health_check = False

            # The health_check attribute is *not* set on newly created
            # ELBs! So we have to create our own.
            if not self.elb.health_check:
                self.elb.health_check = HealthCheck()

            for attr, desired_value in health_check_config.iteritems():
                if getattr(self.elb.health_check, attr) != desired_value:
                    setattr(self.elb.health_check, attr, desired_value)
                    update_health_check = True

            if update_health_check:
                self.elb.configure_health_check(self.elb.health_check)
                self.changed = True

    def _get_health_check_target(self):
        """Compose target string from healthcheck parameters"""
        protocol = self.health_check['ping_protocol'].upper()
        path = ""

        if protocol in ['HTTP', 'HTTPS'] and 'ping_path' in self.health_check:
            path = self.health_check['ping_path']

        return "%s:%s%s" % (protocol, self.health_check['ping_port'], path)


def main():
    module = AnsibleModule(
        argument_spec=dict(
            state={'required': True, 'choices': ['present', 'absent']},
            name={'required': True},
            listeners={'default': None, 'required': False, 'type': 'list'},
            purge_listeners={'default': True, 'required': False,
                             'choices': BOOLEANS, 'type': 'bool'},
            zones={'default': None, 'required': False, 'type': 'list'},
            purge_zones={'default': False, 'required': False,
                         'choices': BOOLEANS, 'type': 'bool'},
            health_check={'default': None, 'required': False, 'type': 'dict'},
            aws_secret_key={'default': None,
                            'aliases': ['ec2_secret_key', 'secret_key'],
                            'no_log': True},
            aws_access_key={'default': None,
                            'aliases': ['ec2_access_key', 'access_key']},
            region={'default': None, 'required': False,
                    'aliases': ['aws_region', 'ec2_region'],
                    'choices': AWS_REGIONS},
        )
    )

    aws_secret_key = module.params['aws_secret_key']
    aws_access_key = module.params['aws_access_key']
    region = module.params['region']
    name = module.params['name']
    state = module.params['state']
    listeners = module.params['listeners']
    purge_listeners = module.params['purge_listeners']
    zones = module.params['zones']
    purge_zones = module.params['purge_zones']
    health_check = module.params['health_check']

    if state == 'present' and not listeners:
        module.fail_json(msg="At least one port is required for ELB creation")

    if state == 'present' and not zones:
        module.fail_json(msg="At least one availability zone is required for ELB creation")

    if not aws_secret_key:
        if 'AWS_SECRET_KEY' in os.environ:
            aws_secret_key = os.environ['AWS_SECRET_KEY']
        elif 'EC2_SECRET_KEY' in os.environ:
            aws_secret_key = os.environ['EC2_SECRET_KEY']

    if not aws_access_key:
        if 'AWS_ACCESS_KEY' in os.environ:
            aws_access_key = os.environ['AWS_ACCESS_KEY']
        elif 'EC2_ACCESS_KEY' in os.environ:
            aws_access_key = os.environ['EC2_ACCESS_KEY']

    if not region:
        if 'AWS_REGION' in os.environ:
            region = os.environ['AWS_REGION']
        elif 'EC2_REGION' in os.environ:
            region = os.environ['EC2_REGION']

    if not region:
        module.fail_json(msg=str("Either region or EC2_REGION environment variable must be set."))

    elb_man = ElbManager(module, name, listeners, purge_listeners, zones,
                         purge_zones, health_check, aws_access_key,
                         aws_secret_key, region=region)

    if state == 'present':
        elb_man.ensure_ok()
    elif state == 'absent':
        elb_man.ensure_gone()

    ansible_facts = {'ec2_elb': 'info'}
    ec2_facts_result = dict(changed=elb_man.changed,
                            elb=elb_man.get_info(),
                            ansible_facts=ansible_facts)

    module.exit_json(**ec2_facts_result)

# this is magic, see lib/ansible/module_common.py
#<<INCLUDE_ANSIBLE_MODULE_COMMON>>

main()