ec2_spots_eip_handler.py

ec2_spots_eip_handler.py

import boto3
import logging

eip_tag_key = "EIP-fleet"
eip_tag_value = "true"
dynamodb_table = "EC2-Spots-EIP"

logger = logging.getLogger()
logger.setLevel("INFO")

def reserve_ip(instance_id):
    dynamodb = boto3.resource('dynamodb')
    table = dynamodb.Table(dynamodb_table)
    logger.info(f'Checking if an instance already has associated EIP (in reboot case for example)')

    response = table.scan(
        FilterExpression="NOT attribute_exists(InstanceID) OR #instance_id = :instance_id",
        ExpressionAttributeNames={"#instance_id": "InstanceID"},
        ExpressionAttributeValues={":instance_id": instance_id}
    )

    if response.get('Items'):
        item = response['Items'][0]
        eip = item['EIP']

        if 'InstanceID' not in item or not item['InstanceID']:
            table.update_item(
                Key={'EIP': eip},
                UpdateExpression='SET InstanceID = :instance_id',
                ExpressionAttributeValues={':instance_id': instance_id}
            )
            logger.info(f'Looking for available EIP for instance {instance_id}')
            logger.info(f'Reserving EIP: {eip} with instance')
            get_allocation_id_by_ip(instance_id, eip)
        else:
            logger.warning(f'Instance {instance_id} already has an associated EIP: {eip}')       
    else:
        logger.warning('No available EIP found for the specified instance')

def release_ip(instance_id):
    dynamodb = boto3.resource('dynamodb')
    table = dynamodb.Table(dynamodb_table)
    logger.info(f'Looking for EIP attached to instance {instance_id}')

    response = table.scan(
        FilterExpression="attribute_exists(InstanceID) AND InstanceID = :instance_id",
        ExpressionAttributeValues={":instance_id": instance_id}
    )

    if response.get('Items'):
        item = response['Items'][0]
        eip = item['EIP']

        table.update_item(
            Key={'EIP': eip},
            UpdateExpression='REMOVE InstanceID'
        )
        logger.info(f'Cleared InstanceID for EIP: {eip}')
    else:
        logger.warning('No matching record found')

def get_allocation_id_by_ip(instance_id, eip):
    logger.info(f'Get allocation ID for EIP: {eip}')
    ec2 = boto3.client('ec2')

    try:
        response = ec2.describe_addresses(Filters=[{'Name': 'public-ip', 'Values': [eip]}])
        if 'Addresses' in response and len(response['Addresses']) > 0:
            allocation_id = response['Addresses'][0]['AllocationId']
            logger.info(f'Allocation ID: {allocation_id}')
            attach_eip(instance_id, allocation_id)
        else:
            logger.error(f'No allocation ID found for IP address: {eip}')
    except Exception as e:
        logger.error(f'Error getting AllocationId for IP address {eip}: {str(e)}')

def attach_eip(instance_id, allocation_id):
    logger.info(f'Attaching EIP to instance: {instance_id}')
    ec2_client = boto3.client('ec2')

    try:
        ec2_client.associate_address(
            InstanceId=instance_id,
            AllocationId=allocation_id
        )
        logger.info(f'Elastic IP attached to instance: {instance_id}')
    except Exception as e:
        logger.error(f'Error attaching Elastic IP to instance {instance_id}: {str(e)}')
    
def eip_handler(instance_id, instance_state):
    if instance_state == "running":
        logger.info(f'EIP reservation for instance {instance_id}')
        reserve_ip(instance_id)
    elif instance_state == "terminated":
        logger.info(f'Release EIP for instance {instance_id}')
        release_ip(instance_id)
    else:
        logger.warning(f'Unsupported state of instance {instance_state}')

def lambda_handler(event, context):
    instanceId = event['detail']['instance-id']
    instanceState = event['detail']['state']
    ec2 = boto3.client('ec2')

    try:
        response = ec2.describe_instances(
            Filters=[{'Name': 'instance-lifecycle', 'Values': ['spot']}],
            InstanceIds=[instanceId]
        )
    except:
        logger.error("Unable to describe instance ID:", instanceId)

    if response['Reservations']:
        logger.info(f'Spot Instance found: {instanceId}, State: {instanceState}')
        logger.info("Checking tag which indicates manage EIP and its value")
        if response["Reservations"][0]["Instances"][0]["Tags"] is not None:
            tags = response["Reservations"][0]["Instances"][0]["Tags"]
            for tag in tags:
                if tag["Key"] == eip_tag_key and tag["Value"] == eip_tag_value:
                    logger.info(f'Tag "{eip_tag_key}" found.')
                    logger.info(f'Value is "{eip_tag_value}"')
                    eip_handler(instanceId, instanceState)
    else:
        logger.info(f'Spot Instance NOT found - skipping {instanceId} State: {instanceState}')

    return