import socket
from os import environ
from socket import gethostbyname_ex
import boto3

CLUSTER_ENDPOINT = environ.get('CLUSTER_ENDPOINT')
TARGET_GROUP = environ.get('TARGET_GROUP')
client = boto3.client('elbv2')

def lambda_handler(event, context):
    ip_addresses = []
    try:
        resolved_ip_addresses = resolve_dns(CLUSTER_ENDPOINT)
        configured_ip_addresses = get_registered_target_ips(TARGET_GROUP)

        targets_to_deregister = [x for x in configured_ip_addresses if x not in resolved_ip_addresses]
        if(targets_to_deregister) :
            deregister_targets(targets_to_deregister, TARGET_GROUP)
       
        targets_to_register = [x for x in resolved_ip_addresses if x not in configured_ip_addresses]
        if(targets_to_register):
            register_targets(targets_to_register, TARGET_GROUP)
            
    except Exception as e:
        print(f"Error: {e}")
    
    return ip_addresses

def resolve_dns(dns_name):
    ip_addresses = gethostbyname_ex(dns_name)[-1]
    return ip_addresses

def get_registered_target_ips(target_group):
    registered_ips = []
    response = client.describe_target_health(TargetGroupArn=target_group)
    return [x["Target"]["Id"] for x in response["TargetHealthDescriptions"]]

def deregister_targets(ips, target_group):
    client.deregister_targets(
        TargetGroupArn = target_group,
        Targets = [{'Id': x} for x in ips]
    )
    print("deregistered %s" % ips)

def register_targets(ips, target_group):
    client.register_targets(
        TargetGroupArn = target_group,
        Targets = [{'Id': x} for x in ips]
    )
    print("registered %s" % ips)

if __name__ == "__main__":
    lambda_handler(None, None)