Add get_rules function to WAF builder

Author: MatthewCane

ca_cdk_constructs/edge_services/waf_rule_templates.py

@@ -4,6 +4,11 @@
 # A COLLECTION OF HELPER FUNCTIONS FOR GENERATING AWS WAFV2 RULES #
 ###################################################################
 
+# Note for developers: These rule builder functions are used within
+# the WAF builder class. If you change any interfaces or make any
+# major changes to these functions, please update the WAF builder
+# class to match, including the docstrings.
+
 
 def managed_rule_group_property(
     name: str,
@@ -126,10 +131,8 @@ def ip_rule_property(
 
     # keys of addresses dict can only be "IPV4" or "IPV6" - easiest way to
     # check is by using set differences
-    if set(addresses.keys()) - set(list(['IPV4', 'IPV6'])) != set(list([])):
-        raise AttributeError(
-            "keys for addresses dict must only be 'IPV4' or 'IPV6'!"
-        )
+    if set(addresses.keys()) - set(list(["IPV4", "IPV6"])) != set(list([])):
+        raise AttributeError("keys for addresses dict must only be 'IPV4' or 'IPV6'!")
 
     # Need IPv4 and IPv6 IP sets
     ipv4_arn = waf.CfnIPSet(
@@ -171,8 +174,8 @@ def ip_rule_property(
                         ip_set_reference_statement=waf.CfnWebACL.IPSetReferenceStatementProperty(
                             arn=ipv6_arn
                         ),
-                    )
+                    ),
                 ]
             )
-        )
+        ),
     )

ca_cdk_constructs/edge_services/waf_v2_builder.py

@@ -22,6 +22,8 @@ class WafV2Builder:
     Functions:
     add_custom_rule: Adds a custom rule to the WAFv2 WebACL.
     add_managed_rule: Adds a managed rule to the WAFv2 WebACL.
+    add_ip_rule: Adds an IP rule to the WAFv2 WebACL.
+    get_rules: Returns the list of rules added to the WAFv2 WebACL.
     build: Builds the WAFv2 WebACL.
 
     Example:
@@ -84,6 +86,11 @@ def __init__(
         self.tags = tags | {"Component": "WAF"}
 
     def add_custom_rule(self, rule: waf.CfnWebACL.RuleProperty) -> None:
+        """
+        Adds a custom rule to the WAFv2 WebACL.
+
+        :param rule: The rule to add.
+        """
         self.rules.append(rule)
 
     def add_managed_rule(
@@ -150,6 +157,12 @@ def add_ip_rule(
             )
         )
 
+    def get_rules(self) -> list[waf.CfnWebACL.RuleProperty]:
+        """
+        Returns the list of rules that have been added to the builder.
+        """
+        return self.rules
+
     def build(self) -> waf.CfnWebACL:
         """
         Builds the WAFv2 WebACL.

tests/test_waf_v2_builder.py

@@ -13,7 +13,6 @@ def waf_builder():
         description="A dummy WAF for testing",
         tags={"Foo": "Bar"},
     )
-
     yield waf_builder
 
 
@@ -78,3 +77,16 @@ def test_waf_v2_logging_enabled():
     )
     waf = waf_builder.build()
     assert waf.visibility_config.cloud_watch_metrics_enabled == True
+
+
+def test_waf_v2_get_rules(waf_builder):
+    waf_builder.add_managed_rule(
+        name="TestManaged",
+        priority=1,
+        managed_rule_name="AWSManagedRulesCommonRuleSet",
+        managed_rule_vendor="AWS",
+        count_only=True,
+    )
+    rules = waf_builder.get_rules()
+    assert len(rules) == 1
+    assert type(rules[0]) == aws_wafv2.CfnWebACL.RuleProperty