diff --git a/examples/waf_rate_limit.rb b/examples/waf_rate_limit.rb new file mode 100644 index 0000000..d966f44 --- /dev/null +++ b/examples/waf_rate_limit.rb @@ -0,0 +1,36 @@ +# CFNDSL + +Resource('RateLimitRule') { + Type 'Custom::WAFRateLimit' + Property('ServiceToken', FnGetAtt('WAFRateLimitFunction', 'Arn')) + Property('EnvironmentName', Ref('EnvironmentName')) + Property('Region', Ref("AWS::Region")) + Property('Rate', 5000) + Property('Negated', true) + Property('Action', 'BLOCK') + Property('IPSet', waf_ip_set(ip_blocks, ['rate_limited'])) + Property('WebACLId', Ref('WebACL')) + Property('Priority', 2) +} + +Resource('WAFRateLimitFunction') { + Type 'AWS::Lambda::Function' + Property('Code', './waf_rate_limit/') + Property('Handler', 'handler.lambda_handler') + Property('Runtime', 'python3.6') + Property('Timeout', 60) + Property('Role', FnGetAtt('WAFRole', 'Arn')) +} + +Resource("WAFRole") { + Type 'AWS::IAM::Role' + Property('AssumeRolePolicyDocument', { + Statement: [ + Effect: 'Allow', + Principal: { Service: [ 'lambda.amazonaws.com' ] }, + Action: [ 'sts:AssumeRole' ] + ] + }) + Property('Path','/') + Property('Policies', Policies.new.get_policies('waf')) +} diff --git a/ssm-secure-parameter/handler.py b/ssm-secure-parameter/handler.py index ba21b42..81e202b 100644 --- a/ssm-secure-parameter/handler.py +++ b/ssm-secure-parameter/handler.py @@ -21,6 +21,8 @@ def lambda_handler(event, context): lambda_response.respond_error(f"{key} property missing") return + replace = cr_params.get('Update', True) + try: parameter = logic.SSMSecureParameterLogic(cr_params['Path']) length = 16 or cr_params['Length'] @@ -40,7 +42,7 @@ def lambda_handler(event, context): elif event['RequestType'] == 'Update': password, version = parameter.create( length=length, - update=True + update=replace ) event['PhysicalResourceId'] = cr_params['Path'] diff --git a/waf_rate_limit/__init__.py b/waf_rate_limit/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/waf_rate_limit/cr_response.py b/waf_rate_limit/cr_response.py new file mode 100644 index 0000000..a1aa830 --- /dev/null +++ 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 new file mode 100644 index 0000000..1bae318 --- /dev/null +++ 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 new file mode 100644 index 0000000..9e1dac3 --- /dev/null +++ b/waf_rate_limit/logic.py @@ -0,0 +1,228 @@ +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.regional = resource_properties.get('Regional', 'false') + self.web_acl_id = resource_properties['WebACLId'] + self.priority = int(resource_properties['Priority']) + + if 'EnvironmentName' in resource_properties: + self.rule_name = f"{resource_properties['EnvironmentName']}-rate-limit" + self.ip_set_name = f"{resource_properties['EnvironmentName']}-rate-limit-ip-set" + else: + self.rule_name = resource_properties['RuleName'] + self.ip_set_name = resource_properties['IpSetName'] + + self.metric_name = self.rule_name.replace('-', '') + + if to_bool(self.regional): + self.client = boto3.client('waf-regional', region_name=self.region) + else: + 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() + + if len(self.ip_set): + 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 new file mode 100644 index 0000000..67399b0 --- /dev/null +++ b/waf_rate_limit/sample-event.json @@ -0,0 +1,26 @@ +{ + "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" + } + ], + "Regional": "true" + }, + "RequestType": "Delete", + "ResourceType": "Custom::WAFRateLimitFunction", + "RequestId": "unique id for this create request", + "LogicalResourceId": "WAFRateLimitFunction", + "PhysicalResourceId": "d55e2c88-3eb7-40d2-8135-cf65d3624b35" +} diff --git a/waf_regex/logic.py b/waf_regex/logic.py index 1a97266..8a2453d 100644 --- a/waf_regex/logic.py +++ b/waf_regex/logic.py @@ -14,7 +14,7 @@ class WafRegexLogic: def __init__(self, resource_properties): self.regex_patterns = resource_properties['RegexPatterns'] self.match_type = resource_properties['Type'] - self.match_data = resource_properties['Data'] + self.match_data = resource_properties.get('Data', '') self.transform = resource_properties['Transform'] self.match_name = resource_properties['Name'] self.pattern_name = f"{resource_properties['Name']}-pattern" @@ -99,22 +99,26 @@ def new_match_set(self): return response_create_match_set['RegexMatchSet']['RegexMatchSetId'] def insert_match_set(self, match_set_id, pattern_set_id): + update = { + 'Action': 'INSERT', + 'RegexMatchTuple': { + 'FieldToMatch': { + 'Type': self.match_type + }, + 'TextTransformation': self.transform, + 'RegexPatternSetId': pattern_set_id + } + } + + # This applies when `match_type` is 'HEADER' or 'SINGLE_QUERY_ARG' + if (self.match_data != '') { + update['RegexMatchTuple']['FieldToMatch']['Data'] = self.match_data + } + changeToken = self.client.get_change_token() update_regex_matchset = self.client.update_regex_match_set( RegexMatchSetId=match_set_id, - Updates=[ - { - 'Action': 'INSERT', - 'RegexMatchTuple': { - 'FieldToMatch': { - 'Type': self.match_type, - 'Data': self.match_data - }, - 'TextTransformation': self.transform, - 'RegexPatternSetId': pattern_set_id - } - }, - ], + Updates=[update], ChangeToken=changeToken['ChangeToken'] )