Add WAF rate limiting rule.

zac · zac · commit 1a98623eb541 · 2018-12-16T22:03:00.000+11:00
diff --git a/waf_rate_limit/__init__.py b/waf_rate_limit/__init__.py
diff --git a/waf_rate_limit/cr_response.py b/waf_rate_limit/cr_response.py
@@ -0,0 +1,58 @@
+import logging
+from urllib.request import urlopen, Request, HTTPError, URLError
+import json
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+class CustomResourceResponse:
+    def __init__(self, request_payload):
+        self.payload = request_payload
+        self.response = {
+            "StackId": request_payload["StackId"],
+            "RequestId": request_payload["RequestId"],
+            "LogicalResourceId": request_payload["LogicalResourceId"],
+            "Status": 'SUCCESS',
+        }
+
+
+    def respond_error(self, message):
+        self.response['Status'] = 'FAILED'
+        self.response['Reason'] = message
+        self.respond({})
+
+    def respond(self, data):
+        event = self.payload
+        response = self.response
+        ####
+        #### copied from https://github.com/ryansb/cfn-wrapper-python/blob/master/cfn_resource.py
+        ####
+
+        if event.get("PhysicalResourceId", False):
+            response["PhysicalResourceId"] = event["PhysicalResourceId"]
+
+        logger.debug("Received %s request with event: %s" %
+                     (event['RequestType'], json.dumps(event)))
+
+        response["Data"] = data
+
+        serialized = json.dumps(response)
+        
+        logger.info(f"Responding to {event['RequestType']} request with: {serialized}")
+        req_data = serialized.encode('utf-8')
+
+        req = Request(
+            event['ResponseURL'],
+            data=req_data,
+            headers={'Content-Length': len(req_data), 'Content-Type': ''}
+        )
+        req.get_method = lambda: 'PUT'
+
+        try:
+            urlopen(req)
+            logger.debug("Request to CFN API succeeded, nothing to do here")
+        except HTTPError as e:
+            logger.error("Callback to CFN API failed with status %d" % e.code)
+            logger.error("Response: %s" % e.reason)
+        except URLError as e:
+            logger.error("Failed to reach the server - %s" % e.reason)
diff --git a/waf_rate_limit/handler.py b/waf_rate_limit/handler.py
@@ -0,0 +1,44 @@
+import sys
+import os
+
+sys.path.append(f"{os.environ['LAMBDA_TASK_ROOT']}/lib")
+sys.path.append(os.path.dirname(os.path.realpath(__file__)))
+
+import cr_response
+from logic import WafRateLimit
+import json
+
+def lambda_handler(event, context):
+
+    print(f"Received event:{json.dumps(event)}")
+
+    lambda_response = cr_response.CustomResourceResponse(event)
+    cr_params = event['ResourceProperties']
+    waf_logic = WafRateLimit(cr_params)
+    try:
+        # if create request, generate physical id, both for create/update copy files
+        if event['RequestType'] == 'Create':
+            event['PhysicalResourceId'] = waf_logic._create_rate_based_rule()
+            data = {
+                "RuleID" : event['PhysicalResourceId']
+            }
+            lambda_response.respond(data)
+
+        elif event['RequestType'] == 'Update':
+            waf_logic._update_rate_based_rule(event['PhysicalResourceId'])
+            data = {
+                "RuleID" : event['PhysicalResourceId']
+            }
+            lambda_response.respond(data)
+
+        elif event['RequestType'] == 'Delete':
+            print(event['PhysicalResourceId'])
+            waf_logic._delete_rate_based_rule(event['PhysicalResourceId'])
+            data = { }
+            lambda_response.respond(data)
+
+    except Exception as e:
+        message = str(e)
+        lambda_response.respond_error(message)
+
+    return 'OK'
diff --git a/waf_rate_limit/logic.py b/waf_rate_limit/logic.py
@@ -0,0 +1,215 @@
+import boto3
+import os
+import glob
+import logging
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+
+class WafRateLimit:
+
+    def __init__(self, resource_properties):
+        self.rate = resource_properties['Rate']
+        self.action = resource_properties['Action']
+        self.region = resource_properties['Region']
+        self.env = resource_properties['EnvironmentName']
+        self.ip_set = resource_properties['IPSet']
+        self.negated = resource_properties['Negated']
+        self.region = resource_properties['Region']
+        self.web_acl_id = resource_properties['WebACLId']
+        self.priority = int(resource_properties['Priority'])
+        self.rule_name = f"{resource_properties['EnvironmentName']}-rate-limit"
+        self.ip_set_name = f"{resource_properties['EnvironmentName']}-rate-limit-ip-set"
+        self.metric_name = self.rule_name.replace('-', '')
+
+        self.client = boto3.client('waf', region_name=self.region)
+
+    def retry(func):
+        # Reattempt to execute a given function with optional arguments.
+        # This is to avoid the insane error about a token already being expired.
+        def wrapper(self, *args, **kwargs):
+            attempts = 5
+            remaining = attempts
+
+            while remaining:
+                try:
+                    result = func(self, *args, **kwargs)
+                    return result
+                except self.client.exceptions.WAFStaleDataException as e:
+                    logger.info(str(e))
+                    logger.info("(%d/%d) Retrying request with a new change token..." % (remaining + 1, attempts))
+                    remaining -= 1
+
+            logger.info("ERROR - failed to execute request.")
+            exit(1)
+
+        return wrapper
+
+    def _create_rate_based_rule(self):
+        rule_id = self.create_rate_based_rule()
+        ip_set_id = self.create_ip_set()
+        self.update_ip_set('INSERT', ip_set_id, self.ip_set)
+        self.update_rate_based_rule('INSERT', ip_set_id, rule_id)
+        self._add_to_web_acl(rule_id)
+
+        return rule_id
+
+    @retry
+    def create_rate_based_rule(self):
+        change_token = self._get_change_token()
+        logger.info("Creating WAF rule '%s' ..." % self.rule_name)
+
+        rule_id = self.client.create_rate_based_rule(
+            Name=self.rule_name,
+            MetricName=self.metric_name,
+            RateLimit=int(self.rate),
+            RateKey='IP',
+            ChangeToken=change_token
+        )['Rule']['RuleId']
+
+        return rule_id
+
+    @retry
+    def create_ip_set(self):
+        change_token = self._get_change_token()
+        logger.info("Creating IP set '%s' ..." % self.ip_set_name)
+
+        ip_set_id = self.client.create_ip_set(
+            Name=self.ip_set_name,
+            ChangeToken=change_token
+        )['IPSet']['IPSetId']
+
+        return ip_set_id
+
+    @retry
+    def update_ip_set(self, action, ip_set_id, ip_set):
+        change_token = self._get_change_token()
+        logger.info("Updating IP set '%s' (%s) with %d IPs as %s ..." % (self.ip_set_name, ip_set_id, len(self.ip_set), action))
+
+        self.client.update_ip_set(
+            IPSetId=ip_set_id,
+            ChangeToken=change_token,
+            Updates=generate_waf_ip_set(action, ip_set)
+        )
+
+    def _update_rate_based_rule(self, rule_id):
+        self._delete_rate_based_rule(rule_id)
+        return self._create_rate_based_rule()
+
+    @retry
+    def update_rate_based_rule(self, action, ip_set_id, rule_id):
+        change_token = self._get_change_token()
+        logger.info("Updating rule '%s' (%s) with IP set '%s' (%s) as %s ..." % (self.rule_name, rule_id, self.ip_set_name, ip_set_id, action))
+
+        self.client.update_rate_based_rule(
+            RuleId=rule_id,
+            ChangeToken=change_token,
+            Updates=[{
+                'Action': action,
+                'Predicate': {
+                    'Negated': to_bool(self.negated),
+                    'Type': 'IPMatch',
+                    'DataId': ip_set_id
+                }
+            }],
+            RateLimit=int(self.rate)
+        )
+
+    def _delete_rate_based_rule(self, rule_id):
+        logger.info("Getting IP set for rule '%s' (%s) ..." % (self.rule_name, rule_id))
+
+        try:
+            predicates = self.client.get_rate_based_rule(
+                RuleId=rule_id
+            )['Rule']['MatchPredicates']
+        except self.client.exceptions.WAFNonexistentItemException as e:
+            logger.info("%s: rule ID '%s' does not exist. Returning success" % (str(e), rule_id))
+            return
+
+        if len(predicates):
+            ip_set_id = predicates[0]['DataId']
+
+            logger.info("Getting IPs for IP set '%s' ..." % (ip_set_id))
+
+            current_ip_set = self.client.get_ip_set(
+                IPSetId=ip_set_id
+            )['IPSet']['IPSetDescriptors']
+
+            if len(current_ip_set):
+                self.update_ip_set('DELETE', ip_set_id, current_ip_set)
+
+            self.update_rate_based_rule('DELETE', ip_set_id, rule_id)
+            self.delete_ip_set(ip_set_id)
+
+        self._delete_from_web_acl(rule_id)
+        self.delete_rate_based_rule(rule_id)
+
+    @retry
+    def delete_ip_set(self, ip_set_id):
+        change_token = self._get_change_token()
+        logger.info("Deleting IP set '%s' ..." % (ip_set_id))
+
+        self.client.delete_ip_set(
+            IPSetId=ip_set_id,
+            ChangeToken=change_token
+        )
+
+    @retry
+    def delete_rate_based_rule(self, rule_id):
+        change_token = self._get_change_token()
+        logger.info("Deleting rule '%s' (%s) ..." % (self.rule_name, rule_id))
+
+        self.client.delete_rate_based_rule(
+            RuleId=rule_id,
+            ChangeToken=change_token
+        )
+
+    def _get_change_token(self):
+        token = self.client.get_change_token()['ChangeToken']
+        logger.info("Got change token: %s" % token)
+        return token
+
+    def _add_to_web_acl(self, rule_id):
+        self._update_web_acl('INSERT', self.action, self.priority, rule_id)
+
+    def _delete_from_web_acl(self, rule_id):
+        # Get the current rule priority, as it is needed in the update request
+        web_acl_rules = self.client.get_web_acl(
+            WebACLId=self.web_acl_id
+        )['WebACL']['Rules']
+
+        current_rule = list(filter(lambda rule: rule['RuleId'] == rule_id, web_acl_rules))[0]
+        current_action = current_rule['Action']['Type']
+        current_priority = int(current_rule['Priority'])
+
+        self._update_web_acl('DELETE', current_action, current_priority, rule_id)
+
+    @retry
+    def _update_web_acl(self, new_action, current_action, priority, rule_id):
+        """Add a rule ID with a web ACL.
+        """
+        change_token = self._get_change_token()
+        logger.info("%sing rule '%s' (%s) in web ACL ID '%s'" % (new_action, self.rule_name, rule_id, self.web_acl_id))
+
+        self.client.update_web_acl(
+            WebACLId=self.web_acl_id,
+            Updates=[{
+              "Action": new_action,
+              "ActivatedRule": {
+                "Action": {
+                    "Type": current_action
+                },
+                "Priority": priority,
+                "RuleId": rule_id,
+                "Type": "RATE_BASED"
+              }
+            }],
+            ChangeToken=change_token
+        )
+
+def generate_waf_ip_set(action, ips):
+    return [{'Action': action, 'IPSetDescriptor': ip } for ip in ips]
+
+def to_bool(value):
+    return value.lower() == 'true'
diff --git a/waf_rate_limit/sample-event.json b/waf_rate_limit/sample-event.json
@@ -0,0 +1,25 @@
+{
+  "StackId": "arn:aws:cloudformation:us-west-2:EXAMPLE/stack-name/guid",
+  "ResponseURL": "http://pre-signed-S3-url-for-response",
+  "ResourceProperties": {
+    "EnvironmentName": "prod",
+    "Region": "ap-southeast-2",
+    "Rate": "5000",
+    "Priority": "10",
+    "Action": "BLOCK",
+    "Negated": "true",
+    "WebACLId": "98f6fb51-3ad7-4cff-8f68-a2f96e707ac4",
+    "Priority": "2",
+    "IPSet": [
+        {
+            "Type": "IPV4",
+            "Value": "123.22.64.68/32"
+        }
+    ]
+  },
+  "RequestType": "Delete",
+  "ResourceType": "Custom::WAFRateLimitFunction",
+  "RequestId": "unique id for this create request",
+  "LogicalResourceId": "WAFRateLimitFunction",
+  "PhysicalResourceId": "d55e2c88-3eb7-40d2-8135-cf65d3624b35"
+}
