mikechau · November 9, 2017 19:39 · Nov 9, 2017
diff --git a/efs_modify_security_groups.lambda.py b/efs_modify_security_groups.lambda.py
@@ -0,0 +1,95 @@
+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)