#!/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
short_description: De-registers or registers instances from EC2 ELB(s)
description:
  - This module de-registers or registers an AWS EC2 instance from the ELB(s)
    that it belongs to.
  - Returns fact "ec2_elbs" which is a list of elbs attached to the instance
    if state=absent is passed as an argument.
  - Will be marked changed when called only if there are ELBs found to operate on.
version_added: "1.2"
requirements: [ "boto", "urllib2" ]
author: John Jarvis
options:
  state:
    description:
      - register or deregister the instance
    required: true

  instance_id:
    description:
      - EC2 Instance ID
    required: true

  ec2_elbs:
    description:
      - List of ELB names, required for registration. The ec2_elbs fact should be used if there was a previous de-register.
    required: false
    default: None
  ec2_secret_key:
    description:
      - AWS Secret API key
    required: false
    default: None
  ec2_access_key:
    description:
      - AWS Access API key
    required: false
    default: None
  ec2_region:
    description:
      - AWS region of your load balancer.  If not set then the region in which
        this module is running will be used.
    required: false

"""

EXAMPLES = """
# basic pre_task and post_task example
pre_tasks:
  - name: Gathering ec2 facts
    ec2_facts:
  - name: Instance De-register
    local_action: ec2_elb
    args:
      instance_id: "{{ ansible_ec2_instance_id }}"
      state: 'absent'
roles:
  - myrole
post_tasks:
  - name: Instance Register
    local_action: ec2_elb
    args:
      instance_id: "{{ ansible_ec2_instance_id }}"
      ec2_elbs: "{{ ec2_elbs }}"
      state: 'present'
"""

import time
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.regioninfo import RegionInfo
except ImportError:
    print "failed=True msg='boto required for this module'"
    sys.exit(1)

try:
    import urllib2
except ImportError:
    print "failed=True msg='urllib2 required for this module'"
    sys.exit(1)

class ElbManager:
    """Handles EC2 instance ELB registration and de-registration"""

    def __init__(self, module, instance_id=None, ec2_elbs=None,
                 ec2_access_key=None, ec2_secret_key=None, ec2_region=None):
        self.ec2_access_key = ec2_access_key
        self.ec2_secret_key = ec2_secret_key
        self.module = module
        self.instance_id = instance_id
        self.ec2_region = ec2_region
        self.lbs = self._get_instance_lbs(ec2_elbs)

        # if there are no ELBs to operate on
        # there will be no changes made
        if len(self.lbs) > 0:
            self.changed = True
        else:
            self.changed = False

    def deregister(self):
        """De-register the instance from all ELBs and wait for the ELB
        to report it out-of-service"""

        for lb in self.lbs:
            lb.deregister_instances([self.instance_id])
            self._await_elb_instance_state(lb, 'OutOfService')

    def register(self):
        """Register the instance for all ELBs and wait for the ELB
        to report the instance in-service"""

        for lb in self.lbs:
            lb.register_instances([self.instance_id])
            self._await_elb_instance_state(lb, 'InService')

    def exists(self, lbtest):
        """ Verify that the named ELB actually exists """
        
        found = False
        for lb in self.lbs:
            if lb.name == lbtest:
                found=True
                break
        return found


    def _await_elb_instance_state(self, lb, awaited_state):
        """Wait for an ELB to change state
        lb: load balancer
        awaited_state : state to poll for (string)"""

        while True:
            state = lb.get_instance_health([self.instance_id])[0].state
            if state == awaited_state:
                break
            else:
                time.sleep(1)

    def _get_instance_lbs(self, ec2_elbs=None):
        """Returns a list of ELBs attached to self.instance_id
        ec2_elbs: an optional list of elb names that will be used
                  for elb lookup instead of returning what elbs
                  are attached to self.instance_id"""

        try:
            endpoint="elasticloadbalancing.%s.amazonaws.com" % self.ec2_region
            connect_region = RegionInfo(name=self.ec2_region, endpoint=endpoint)
            elb = boto.ec2.elb.ELBConnection(self.ec2_access_key, self.ec2_secret_key, region=connect_region)
        except boto.exception.NoAuthHandlerFound, e:
            self.module.fail_json(msg=str(e))

        elbs = elb.get_all_load_balancers()

        if ec2_elbs:
            lbs = sorted(lb for lb in elbs if lb.name in ec2_elbs)
        else:
            lbs = []
            for lb in elbs:
                for info in lb.instances:
                    if self.instance_id == info.id:
                        lbs.append(lb)
        return lbs


def main():

    module = AnsibleModule(
        argument_spec=dict(
            state={'required': True,
                    'choices': ['present', 'absent']},
            instance_id={'required': True},
            ec2_elbs={'default': None, 'required': False},
            ec2_secret_key={'default': None, 'aliases': ['EC2_SECRET_KEY']},
            ec2_access_key={'default': None, 'aliases': ['EC2_ACCESS_KEY']},
            ec2_region={'default': None, 'required': False, 'choices':AWS_REGIONS}
        )
    )

    ec2_secret_key = module.params['ec2_secret_key']
    ec2_access_key = module.params['ec2_access_key']
    ec2_elbs = module.params['ec2_elbs']
    ec2_region = module.params['ec2_region']

    if module.params['state'] == 'present' and 'ec2_elbs' not in module.params:
        module.fail_json(msg="ELBs are required for registration")

    if not ec2_secret_key and 'EC2_SECRET_KEY' in os.environ:
        ec2_secret_key = os.environ['EC2_SECRET_KEY']
    if not ec2_access_key and 'EC2_ACCESS_KEY' in os.environ:
        ec2_access_key = os.environ['EC2_ACCESS_KEY']

    if not ec2_region and 'EC2_REGION' in os.environ:
        ec2_region = os.environ['EC2_REGION']

    if not ec2_region:
        response = urllib2.urlopen('http://169.254.169.254/latest/meta-data/placement/availability-zone')
        az = response.read()
        for r in AWS_REGIONS:
            if az.startswith(r):
                ec2_region = r
                break

        if not ec2_region:
            module.fail_json(msg = str("ec2_region not specified and unable to determine region from AWS."))

    instance_id = module.params['instance_id']
    elb_man = ElbManager(module, instance_id, ec2_elbs, ec2_access_key,
                         ec2_secret_key, ec2_region=ec2_region)

    for elb in [ ec2_elbs ]:
        if not elb_man.exists(elb):
            str="ELB %s does not exist" % elb
            module.fail_json(msg=str)

    if module.params['state'] == 'present':
        elb_man.register()
    elif module.params['state'] == 'absent':
        elb_man.deregister()

    ansible_facts = {'ec2_elbs': [lb.name for lb in elb_man.lbs]}
    ec2_facts_result = dict(changed=elb_man.changed, ansible_facts=ansible_facts)

    module.exit_json(**ec2_facts_result)

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

main()