efs_modify_security_groups.lambda.py

import cfnresponse
import boto3

CLIENT = boto3.client('efs')


def execute(client, file_system_id, action, security_group_ids):
    if client is None or not client:
        raise TypeError('client must be present')
    if not isinstance(file_system_id, basestring):
        raise TypeError('file_system_id must be a valid string - {} is {}'.format(file_system_id, type(file_system_id)))
    if not isinstance(action, basestring):
        raise TypeError('action must be a valid string - {} is {}'.format(action, type(action)))
    if not isinstance(security_group_ids, list):
        raise TypeError('security_group_ids must be a list - {} is {}'.format(security_group_ids, type(security_group_ids)))
    if len(security_group_ids) not in range(1, 6):
        raise ValueError(
            'security_group_ids must have a length between 1 to 5 - {} ids were provided'.format(len(security_group_ids)))

    mount_targets = client.describe_mount_targets(
        FileSystemId=file_system_id
    )

    mount_target_sgs = dict(map(
        lambda x: (
            x['MountTargetId'], client.describe_mount_target_security_groups(
                MountTargetId=x['MountTargetId'])['SecurityGroups']
        ),
        mount_targets['MountTargets']
    ))

    if action == 'replace':
        print '[REPLACE]: Replacing security groups...'

        for mt_id, sg_ids in mount_target_sgs.iteritems():
            current_sg_ids = set(sg_ids)
            new_sg_ids = set(security_group_ids)
            delta_removed_sg_ids = list(current_sg_ids - new_sg_ids)
            delta_added_sg_ids = list(new_sg_ids - current_sg_ids)
            delta_sg_ids = delta_removed_sg_ids + delta_added_sg_ids

            if len(delta_removed_sg_ids):
                print '[REPLACE]: Delta found. Removing security groups ({}) for the mount target ({}).'.format(", ".join(delta_removed_sg_ids), mt_id)

            if len(delta_added_sg_ids):
                print '[REPLACE]: Delta found. Adding security groups ({}) for the mount target ({}).'.format(", ".join(delta_added_sg_ids), mt_id)

            if len(delta_sg_ids):
                client.modify_mount_target_security_groups(
                    MountTargetId=mt_id,
                    SecurityGroups=security_group_ids
                )
            else:
                print '[REPLACE]: Delta not found. Skipping replace action for the mount target ({}).'.format(mt_id)

        return security_group_ids
    else:
        raise ValueError('Invalid action: {}.'.format(action))


def handler(event, context):
    resource_id = event['LogicalResourceId']
    event_request_type = event['RequestType']
    event_data = event['ResourceProperties']
    response_data = {}
    response_data['LogStreamName'] = context.log_stream_name

    print '[CFN]: {} triggered... (md5: {})'.format(event_request_type, event_data['MD5'])

    try:
        if event_request_type == 'Create' or event_request_type == 'Update':
            response_data['SecurityGroupIds'] = execute(CLIENT, event_data['FileSystemId'], 'replace', filter(
                bool, event_data['AddSecurityGroupIds'].split(",")))
        elif event_request_type == 'Delete':
            response_data['SecurityGroupIds'] = execute(CLIENT, event_data['FileSystemId'], 'replace', filter(
                bool, event_data['RemoveSecurityGroupIds'].split(",")))
        else:
            raise ValueError(
                'Invalid request type for {}.'.format(event_request_type))

        cfnresponse.send(event, context, cfnresponse.SUCCESS,
                         response_data, resource_id)
    except Exception as err:
        message = "[FAILURE]: {}".format(err)

        print message

        response_data['Reason'] = message

        if event_request_type == 'Delete':
            cfnresponse.send(event, context, cfnresponse.SUCCESS,
                          response_data, resource_id)
        else:
            cfnresponse.send(event, context, cfnresponse.FAILED,
                         response_data, resource_id)