import json
import logging
import boto3
import os

logger = logging.getLogger()
logger.setLevel(logging.INFO)
logging.getLogger("boto3").setLevel(logging.WARNING)
logging.getLogger("botocore").setLevel(logging.WARNING)

def lambda_handler(event, context):

    ec2 = boto3.resource('ec2')
    
    new_sgs = []  
    
    # Get new security group ids from environment variables
    if "Add_SG_1" in os.environ and os.getenv("Add_SG_1") != "":
        new_sgs.append(os.environ['Add_SG_1'])
        logger.info("FOUND ENVAR SG1: %s.", os.environ['Add_SG_1'] )
    if "Add_SG_2" in os.environ and os.getenv("Add_SG_2") != "":
        new_sgs.append(os.environ['Add_SG_2'])
        logger.info("FOUND ENVAR SG2: %s.", os.environ['Add_SG_2'] )
    if "Add_SG_3" in os.environ and os.getenv("Add_SG_3") != "":
        new_sgs.append(os.environ['Add_SG_3'])
        logger.info("FOUND ENVAR SG3: %s.", os.environ['Add_SG_3'] )
    if "Add_SG_4" in os.environ and os.getenv("Add_SG_4") != "":
        new_sgs.append(os.environ['Add_SG_4'])
        logger.info("FOUND ENVAR SG4: %s.", os.environ['Add_SG_4'] )

    # Retrieve network interface id from event data
    logger.info("Querying for EUC ENI id in event data.")
    try :
        eni_id = event['detail']['responseElements']['networkInterface']['networkInterfaceId']
        logger.info("ENI id found in event data: %s.", eni_id)
        
        # Get list of security groups already attached to ENI
        logger.info("Generating list of existing security groups on EUC ENI.")
        euc_eni = ec2.NetworkInterface(eni_id)
        euc_eni_groups = euc_eni.groups
        euc_eni_sg_ids = [ eni.get('GroupId') for eni in euc_eni_groups ]
        logger.info("Found %s existing security groups on EUC ENI: %s.", len(euc_eni_groups), euc_eni_groups )

        # Combine list of new and existing security groups to add to ENI
        logger.info("Generating updated list of security groups to add to EUC ENI.")
        for new_sg in new_sgs:
            if new_sg not in euc_eni_sg_ids :
                logger.info("Adding %s to ENI attach list.", new_sg)
                euc_eni_sg_ids.append(new_sg)
            else :
                logger.info("%s already attached to ENI, will not add again.", new_sg)

        # Check that no more than 5 SGs are added to ENI.
        if len(euc_eni_sg_ids) < 6 :
            logger.info("Attaching %s security groups to ENI: %s", len(euc_eni_sg_ids), euc_eni_sg_ids)
            euc_eni.modify_attribute(Groups=euc_eni_sg_ids)
            logger.info("Completed attachment of security groups to ENI")
        else :
            logger.error("Attempting to attach more that 5 security groups, aborting.")
        
    except Exception as e :
        logger.error(e)
        logger.error("Unable to successfully update security groups on ENI.")
    

    return {
        'statusCode': 200
    }