cloudtools · markpeek · Jun 13, 2021 · May 10, 2021 · May 10, 2021 · May 10, 2021
diff --git a/tests/test_validators.py b/tests/test_validators.py
@@ -27,6 +27,8 @@
     task_type,
     tg_healthcheck_port,
     waf_action_type,
+    wafv2_custom_body_response_content,
+    wafv2_custom_body_response_content_type
 )
 
 
@@ -273,6 +275,20 @@ def test_waf_action_type(self):
             with self.assertRaises(ValueError):
                 waf_action_type(s)
 
+    def test_wafv2_custom_body_response_content(self):
+        for s in ["{'hello': 'world'}", "<!DOCTYPE html><html><head><title>Test</title></head><body><h1>Test</h1><p>Test.</p></body></html>", "Health"]:
+            wafv2_custom_body_response_content(s)
+        for s in ["", "a"*10241]:
+            with self.assertRaises(ValueError):
+                wafv2_custom_body_response_content(s)
+
+    def test_wafv2_custom_body_response_content_type(self):
+        for s in ["APPLICATION_JSON", "TEXT_HTML", "TEXT_PLAIN"]:
+            wafv2_custom_body_response_content_type(s)
+        for s in ["", "APPLICATION", "HTML", "TEXT"]:
+            with self.assertRaises(ValueError):
+                wafv2_custom_body_response_content_type(s)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/troposphere/validators.py b/troposphere/validators.py
@@ -658,3 +658,23 @@ def ecs_efs_encryption_status(status):
             % (", ".join(valid_status))
         )
     return status
+
+
+def wafv2_custom_body_response_content(content):
+    """validate wafv2 custom body response content. Any character between 1 to 10240
+    """
+    if not content:
+        raise ValueError("Content must not be empty")
+    if len(content) > 10240:
+        raise ValueError("Content maximum length must not exceed 10240")
+
+    return content
+
+
+def wafv2_custom_body_response_content_type(content_type):
+    """validate wafv2 custom response content type
+    """
+    valid_types = ["APPLICATION_JSON", "TEXT_HTML", "TEXT_PLAIN"]
+    if content_type not in valid_types:
+        raise ValueError('ContentType must be one of: "%s"' % (", ".join(valid_types)))
+    return content_type
diff --git a/troposphere/wafv2.py b/troposphere/wafv2.py
@@ -4,7 +4,7 @@
 # See LICENSE file for full license.
 
 from . import AWSObject, AWSProperty, Tags
-from .validators import boolean, integer
+from .validators import boolean, integer, wafv2_custom_body_response_content, wafv2_custom_body_response_content_type
 
 VALID_TRANSFORMATION_TYPES = (
     "CMD_LINE",
@@ -68,6 +68,22 @@ def validate_positional_constraint(positional_constraint):
     return positional_constraint
 
 
+def validate_custom_response_bodies(custom_response_bodies):
+    """validate custom response bodies
+    """
+    if not isinstance(custom_response_bodies, dict):
+        raise ValueError("CustomResponseBodies must be dict")
+
+    for k, v in custom_response_bodies.items():
+        if not isinstance(v, CustomResponseBody):
+            raise ValueError(
+                "value of %s must be type of CustomResponseBody"
+                % (k)
+            )
+
+    return custom_response_bodies
+
+
 class ExcludedRule(AWSProperty):
     props = {"Name": (str, False)}
 
@@ -354,6 +370,7 @@ class WebACL(AWSObject):
     resource_type = "AWS::WAFv2::WebACL"
 
     props = {
+        "CustomResponseBodies": (validate_custom_response_bodies, False),
         "DefaultAction": (DefaultAction, False),
         "Description": (str, False),
         "Name": (str, False),
@@ -404,6 +421,7 @@ class RuleGroup(AWSObject):
 
     props = {
         "Capacity": (integer, False),
+        "CustomResponseBodies": (validate_custom_response_bodies, False),
         "Description": (str, False),
         "Name": (str, False),
         "Rules": ([RuleGroupRule], False),
@@ -420,3 +438,10 @@ class WebACLAssociation(AWSObject):
         "ResourceArn": (str, True),
         "WebACLArn": (str, True),
     }
+
+
+class CustomResponseBody(AWSObject):
+    props = {
+        "Content": (wafv2_custom_body_response_content, True),
+        "ContentType": (wafv2_custom_body_response_content_type, True)
+    }