Refactoring of ec2_eip module.

cloud/amazon/ec2_eip.py

@@ -100,114 +100,72 @@ try:
 except ImportError:
     HAS_BOTO = False
 
-wait_timeout = 0
 
+class EIPException(Exception):
+    pass
 
-def associate_ip_and_instance(ec2, address, instance_id, module):
-    if ip_is_associated_with_instance(
-            ec2, address.public_ip, instance_id, module):
-        module.exit_json(changed=False, public_ip=address.public_ip)
 
-    # If we're in check mode, nothing else to do
+def associate_ip_and_instance(ec2, address, instance_id, check_mode):
-    if module.check_mode:
+    if address_is_associated_with_instance(ec2, address, instance_id):
-        module.exit_json(changed=True)
+        return {'changed': False}
 
-    try:
+    # If we're in check mode, nothing else to do
+    if not check_mode:
         if address.domain == 'vpc':
             res = ec2.associate_address(instance_id,
                                         allocation_id=address.allocation_id)
         else:
             res = ec2.associate_address(instance_id,
                                         public_ip=address.public_ip)
-    except boto.exception.EC2ResponseError, e:
+        if not res:
-        module.fail_json(msg=str(e))
+            raise EIPException('association failed')
 
-    if res:
+    return {'changed': True}
-        module.exit_json(changed=True, public_ip=address.public_ip)
-    module.fail_json(msg='association failed')
 
 
-def disassociate_ip_and_instance(ec2, address, instance_id, module):
+def disassociate_ip_and_instance(ec2, address, instance_id, check_mode):
-    if not ip_is_associated_with_instance(
+    if not address_is_associated_with_instance(ec2, address, instance_id):
-            ec2, address.public_ip, instance_id, module):
+        return {'changed': False}
-        module.exit_json(changed=False, public_ip=address.public_ip)
 
     # If we're in check mode, nothing else to do
-    if module.check_mode:
+    if not check_mode:
-        module.exit_json(changed=True)
+        if address.domain == 'vpc':
-
-    try:
-        if address.domain == "vpc":
             res = ec2.disassociate_address(
                 association_id=address.association_id)
         else:
             res = ec2.disassociate_address(public_ip=address.public_ip)
-    except boto.exception.EC2ResponseError, e:
-        module.fail_json(msg=str(e))
 
-    if res:
+        if not res:
-        module.exit_json(changed=True)
+            raise EIPException('disassociation failed')
-    else:
-        module.fail_json(msg="disassociation failed")
 
+    return {'changed': True}
 
-def find_address(ec2, public_ip, module, fail_on_not_found=True):
+
+def find_address(ec2, public_ip, wait_timeout):
     """ Find an existing Elastic IP address """
-    if wait_timeout != 0:
+    deadline = time.time() + wait_timeout
-        timeout = time.time() + wait_timeout
+    while True:
-        while timeout > time.time():
-            try:
-                addresses = ec2.get_all_addresses([public_ip])
-                break
-            except boto.exception.EC2ResponseError, e:
-                if "Address '%s' not found." % public_ip in e.message:
-                    if not fail_on_not_found:
-                        return None
-                else:
-                    module.fail_json(msg=str(e.message))
-            time.sleep(5)
-
-        if timeout <= time.time():
-            module.fail_json(msg="wait for EIPs timeout on %s" %
-                             time.asctime())
-    else:
         try:
-            addresses = ec2.get_all_addresses([public_ip])
+            return ec2.get_all_addresses([public_ip])[0]
-        except boto.exception.EC2ResponseError, e:
+        except boto.exception.EC2ResponseError as e:
-            if "Address '%s' not found." % public_ip in e.message:
+            if "Address '{}' not found.".format(public_ip) not in e.message:
-                if not fail_on_not_found:
+                raise
-                    return None
-            module.fail_json(msg=str(e.message))
 
-    return addresses[0]
+        if time.time() >= deadline:
+            raise EIPException('wait for EIPs timeout on {}'
+                               .format(time.asctime()))
+        time.sleep(5)
 
 
-def ip_is_associated_with_instance(ec2, public_ip, instance_id, module):
+def address_is_associated_with_instance(ec2, address, instance_id):
     """ Check if the elastic IP is currently associated with the instance """
-    address = find_address(ec2, public_ip, module)
     if address:
-        return address.instance_id == instance_id
+        return address and address.instance_id == instance_id
     return False
 
 
-def instance_is_associated(ec2, instance, module):
+def allocate_address(ec2, domain, reuse_existing_ip_allowed):
-    """
-    Check if the given instance object is already associated with an
-    elastic IP
-    """
-    instance_ip = instance.ip_address
-    if not instance_ip:
-        return False
-    eip = find_address(ec2, instance_ip, module, fail_on_not_found=False)
-    return (eip and (eip.public_ip == instance_ip))
-
-
-def allocate_address(ec2, domain, module, reuse_existing_ip_allowed):
     """ Allocate a new elastic IP address (when needed) and return it """
-    # If we're in check mode, nothing else to do
-    if module.check_mode:
-        module.exit_json(change=True)
-
     if reuse_existing_ip_allowed:
         domain_filter = {'domain': domain or 'standard'}
         all_addresses = ec2.get_all_addresses(filters=domain_filter)
@@ -220,45 +178,67 @@ def allocate_address(ec2, domain, module, reuse_existing_ip_allowed):
     return ec2.allocate_address(domain=domain)
 
 
-def release_address(ec2, public_ip, module):
+def release_address(ec2, address, check_mode):
     """ Release a previously allocated elastic IP address """
 
-    address = find_address(ec2, public_ip, module)
-
     # If we're in check mode, nothing else to do
-    if module.check_mode:
+    if not check_mode:
-        module.exit_json(change=True)
+        if not address.release():
+            EIPException('release failed')
 
-    res = address.release()
+    return {'changed': True}
-    if res:
-        module.exit_json(changed=True)
-    else:
-        module.fail_json(msg="release failed")
 
 
-def find_instance(ec2, instance_id, module):
+def find_instance(ec2, instance_id):
     """ Attempt to find the EC2 instance and return it """
 
-    try:
+    reservations = ec2.get_all_reservations(instance_ids=[instance_id])
-        reservations = ec2.get_all_reservations(instance_ids=[instance_id])
-    except boto.exception.EC2ResponseError, e:
-        module.fail_json(msg=str(e))
 
     if len(reservations) == 1:
         instances = reservations[0].instances
         if len(instances) == 1:
             return instances[0]
 
-    module.fail_json(msg="could not find instance" + instance_id)
+    raise EIPException("could not find instance" + instance_id)
+
+
+def ensure_present(ec2, domain, address, instance_id,
+                   reuse_existing_ip_allowed, check_mode):
+    changed = False
 
+    # Return the EIP object since we've been given a public IP
+    if not address:
+        if check_mode:
+            return {'changed': True}
 
-def allocate_eip(ec2, eip_domain, module, reuse_existing_ip_allowed, new_eip_timeout):
+        address = allocate_address(ec2, domain, reuse_existing_ip_allowed)
-    # Allocate a new elastic IP
+        changed = True
-    address = allocate_address(ec2, eip_domain, module, reuse_existing_ip_allowed)
+
-    # overriding the timeout since this is a a newly provisioned ip
+    if instance_id:
-    global wait_timeout
+        # Allocate an IP for instance since no public_ip was provided
-    wait_timeout = new_eip_timeout
+        instance = find_instance(ec2, instance_id)
-    return address
+        if instance.vpc_id:
+            domain = 'vpc'
+
+        # Associate address object (provided or allocated) with instance
+        assoc_result = associate_ip_and_instance(ec2, address, instance_id,
+                                                 check_mode)
+        changed = changed or assoc_result['changed']
+
+    return {'changed': changed, 'public_ip': address.public_ip}
+
+
+def ensure_absent(ec2, domain, address, check_mode):
+    if not address:
+        return {'changed': False}
+
+    # disassociating address from instance
+    if instance_id:
+        return disassociate_ip_and_instance(ec2, address, instance_id,
+                                            check_mode)
+    # releasing address
+    else:
+        return release_address(ec2, address, check_mode)
 
 
 def main():
@@ -290,48 +270,22 @@ def main():
     in_vpc = module.params.get('in_vpc')
     domain = 'vpc' if in_vpc else None
     reuse_existing_ip_allowed = module.params.get('reuse_existing_ip_allowed')
-    new_eip_timeout = int(module.params.get('wait_timeout'))
+    wait_timeout = int(module.params.get('wait_timeout'))
 
-    if state == 'present':
+    try:
-        # If both instance_id and public_ip are not specified, allocate a new
-        # elastic IP, and exit.
-        if not instance_id and not public_ip:
-            address = allocate_eip(ec2, domain, module,
-                                   reuse_existing_ip_allowed, new_eip_timeout)
-            module.exit_json(changed=True, public_ip=address.public_ip)
-
-        # Return the EIP object since we've been given a public IP
         if public_ip:
-            address = find_address(ec2, public_ip, module)
+            address = find_address(ec2, public_ip, wait_timeout)
 
-        # Allocate an IP for instance since no public_ip was provided
+        if state == 'present':
-        if instance_id and not public_ip:
+            result = ensure_present(ec2, domain, address, instance_id,
-            instance = find_instance(ec2, instance_id, module)
+                                    reuse_existing_ip_allowed,
-
+                                    module.check_mode)
-            if instance.vpc_id:
-                domain = "vpc"
-
-            # Do nothing if the instance is already associated with an
-            # elastic IP.
-            if instance_is_associated(ec2, instance, module):
-                module.exit_json(changed=False, public_ip=instance.ip_address)
-
-            # If the instance is not already associated with an elastic IP,
-            # allocate a new one.
-            address = allocate_eip(
-                ec2, domain, module, reuse_existing_ip_allowed, new_eip_timeout)
-
-        # Associate address object (provided or allocated) with instance
-        associate_ip_and_instance(ec2, address, instance_id, module)
-
-    else:
-        # disassociating address from instance
-        if instance_id:
-            address = find_address(ec2, public_ip, module)
-            disassociate_ip_and_instance(ec2, address, instance_id, module)
-        # releasing address
         else:
-            release_address(ec2, public_ip, module)
+            result = ensure_absent(ec2, domain, address, module.check_mode)
+    except (boto.exception.EC2ResponseError, EIPException) as e:
+        module.fail_json(msg=str(e))
+
+    module.exit_json(**result)
 
 
 # import module snippets