自動アップデートする AWS WAF IP セットを CloudFormation で作成してみる

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

AWSTemplateFormatVersion: 2010-09-09
Description: Creates two regional AWS WAF IP sets that are automatically updated with AWS service's IP ranges.

Parameters:
  IPV4SetNameSuffix: 
    Type: String
    Default: IPv4Set
    Description: The AWS WAF IPv4 set suffix. The prefix will be the stack name. The IP set is initially created with a bogus address
  SERVICES:
    Type: String
    Default: AMAZON
    Description: Enter the name of the AWS services to add, separated by commas and as explained in https://docs.aws.amazon.com/general/latest/gr/aws-ip-ranges.html
  EC2REGIONS:
    Type: String
    Default: 'all'
    Description: For the "EC2" service, specify the AWS regions to add, separated by commas. Use 'all' to add all AWS regions.

Resources:
  IPv4Set:
    Type: AWS::WAFv2::IPSet
    Properties: 
      Addresses: 
        - 0.0.0.0/1
      Description: IPv4 set automatically updated with AWS IP ranges
      IPAddressVersion: IPV4
      Name: !Sub
        - ${AWS::StackName}-${Suffix}
        - {Suffix: !Ref IPV4SetNameSuffix}
      Scope: REGIONAL

  LambdaUpdateWAFIPSet:
    Type: AWS::Lambda::Function
    Properties: 
      Description: This Lambda function, invoked by an incoming SNS message, updates the IPv4 and IPv6 sets with the addresses from the specified services
      Environment: 
        Variables:
          IPV4_SET_NAME:
            !Select 
              - "0"
              - !Split [ "|" , Ref: IPv4Set]
          IPV4_SET_ID:
            Fn::GetAtt: [ IPv4Set, Id ]
          SERVICES:
            Ref: SERVICES
          EC2_REGIONS:
            Ref: EC2REGIONS
          INFO_LOGGING: "false"

      FunctionName: !Sub '${AWS::StackName}-UpdateWAFIPSets'
      Handler: index.lambda_handler
      MemorySize: 128
      Role:
        Fn::GetAtt: [ LambdaUpdateWAFIPSetIamRole, Arn ]
      Runtime: python3.8
      Timeout: 10
      Code:
        ZipFile: !Sub |
          '''
          Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
          SPDX-License-Identifier: MIT-0
          '''

          import boto3
          import hashlib
          import json
          import logging
          import os
          from urllib import request
          import cfnresponse


          ####### Get values from environment variables  ######
          IPV4_SET_NAME=os.environ['IPV4_SET_NAME'].strip()
          IPV4_SET_ID=os.environ['IPV4_SET_ID'].strip()
          # IPV6_SET_NAME=os.environ['IPV6_SET_NAME'].strip()
          # IPV6_SET_ID=os.environ['IPV6_SET_ID'].strip()

          # Set default services if env. variable does not exist or its an empty string
          SERVICES = os.getenv( 'SERVICES', 'ROUTE53_HEALTHCHECKS,CLOUDFRONT').split(',')
          if SERVICES == ['']: SERVICES = ['ROUTE53_HEALTHCHECKS','CLOUDFRONT']

          # Set EC2 region to 'all' if env. variable does not exist or its an empty string
          EC2_REGIONS = os.getenv('EC2_REGIONS','all').split(',')
          if EC2_REGIONS == ['']: EC2_REGIONS = ['all']

          # Set logging level from environment variable
          INFO_LOGGING = os.getenv('INFO_LOGGING','false')
          if INFO_LOGGING == ['']: INFO_LOGGING = 'false'

          #######

          def lambda_handler(event, context):
              # Set up logging. Set the level if the handler is already configured.
              if len(logging.getLogger().handlers) > 0:
                  logging.getLogger().setLevel(logging.ERROR)
              else:
                  logging.basicConfig(level=logging.ERROR)
              # Set the environment variable DEBUG to 'true' if you want verbose debug details in CloudWatch Logs.
              if INFO_LOGGING == 'true':
                  logging.getLogger().setLevel(logging.INFO)
              
              if event.get('RequestType'):
                ip_ranges = json.loads(get_ip_groups_json('https://ip-ranges.amazonaws.com/ip-ranges.json', 'test-hash'))
                # Extract the service ranges
                ranges = get_ranges_for_service(ip_ranges,SERVICES,EC2_REGIONS)
                # Update the AWS WAF IP sets
                update_waf_ipset(IPV4_SET_NAME,IPV4_SET_ID,ranges['ipv4'])
                cfnresponse.send(event, context, cfnresponse.SUCCESS,
                                    {'Response': 'Success'})
              else:
                # If you want different services, set the SERVICES environment variable
                # It defaults to ROUTE53_HEALTHCHECKS and CLOUDFRONT. Using 'jq' and 'curl' get the list of possible
                # services like this:
                # curl -s 'https://ip-ranges.amazonaws.com/ip-ranges.json' | jq -r '.prefixes[] | .service' ip-ranges.json | sort -u 
              
                message = json.loads(event['Records'][0]['Sns']['Message'])

                # Load the ip ranges from the url
                ip_ranges = json.loads(get_ip_groups_json(message['url'], message['md5']))

                # Extract the service ranges
                ranges = get_ranges_for_service(ip_ranges,SERVICES,EC2_REGIONS)

                # Update the AWS WAF IP sets
                update_waf_ipset(IPV4_SET_NAME,IPV4_SET_ID,ranges['ipv4'])
                # update_waf_ipset(IPV6_SET_NAME,IPV6_SET_ID,ranges['ipv6'])

                return ranges
              
          def get_ip_groups_json(url, expected_hash):

              logging.debug("Updating from " + url)

              response = request.urlopen(url)
              ip_json = response.read()

              m = hashlib.md5()
              m.update(ip_json)
              hash = m.hexdigest()

              # If the hash provided is 'test-hash', returns the JSON without checking the hash
              if expected_hash == 'test-hash':
                  print('Running in test mode')
                  return ip_json

              if hash != expected_hash:
                  raise Exception('MD5 Mismatch: got ' + hash + ' expected ' + expected_hash)

              return ip_json

          def get_ranges_for_service(ranges, services,ec2_regions):
              """Gets IPv4 and IPv6 prefixes from the matching services"""
              service_ranges = {'ipv6':[],'ipv4':[]}
              ec2_regions = strip_list(ec2_regions)
              services = strip_list(services)

              # Loop over the IPv4 prefixes and appends the matching services
              print(f'Searching for {services} IPv4 prefixes')
              for prefix in ranges['prefixes']:

                  if prefix['service'] in services and \
                      (
                          (prefix['service'] != 'EC2') \
                          or \
                          (prefix['service']=='EC2' and ec2_regions != ['all'] and prefix['region'] in ec2_regions) \
                          or \
                          (prefix['service']=='EC2' and ec2_regions == ['all'])
                      ):

                      logging.info((f"Found {prefix['service']} region: {prefix['region']} range: {prefix['ip_prefix']}"))
                      service_ranges['ipv4'].append(prefix['ip_prefix'])

              # Loop over the IPv6 prefixes and appends the matching services
              print(f'Searching for {services} IPv6 prefixes')
              for ipv6_prefix in ranges['ipv6_prefixes']:

                  if ipv6_prefix['service'] in services and \
                      (
                          (ipv6_prefix['service'] != 'EC2') \
                          or \
                          (ipv6_prefix['service']=='EC2' and ec2_regions != ['all'] and ipv6_prefix['region'] in ec2_regions) \
                          or \
                          (ipv6_prefix['service']=='EC2' and ec2_regions == ['all'])
                      ):

                      logging.info((f"Found {ipv6_prefix['service']} region: {ipv6_prefix['region']} ipv6 range: {ipv6_prefix['ipv6_prefix']}"))
                      service_ranges['ipv6'].append(ipv6_prefix['ipv6_prefix'])

              return service_ranges

          def update_waf_ipset(ipset_name,ipset_id,address_list):
              """Updates the AWS WAF IP set"""
              waf_client = boto3.client('wafv2')

              lock_token = get_ipset_lock_token(waf_client,ipset_name,ipset_id)

              logging.info(f'Got LockToken for AWS WAF IP Set "{ipset_name}": {lock_token}')

              waf_client.update_ip_set(
                  Name=ipset_name,
                  Scope='REGIONAL',
                  Id=ipset_id,
                  Addresses=address_list,
                  LockToken=lock_token
              )

              print(f'Updated IPSet "{ipset_name}" with {len(address_list)} CIDRs')

          def get_ipset_lock_token(client,ipset_name,ipset_id):
              """Returns the AWS WAF IP set lock token"""
              ip_set = client.get_ip_set(
                  Name=ipset_name,
                  Scope='REGIONAL',
                  Id=ipset_id)
              
              return ip_set['LockToken']

          def strip_list(list):
              """Strips individual elements of the strings"""
              return [item.strip() for item in list]

  FirstSetIP:
      Type: Custom::SetupLambda
      Properties:
        ServiceToken: 
          Fn::GetAtt: 
            - LambdaUpdateWAFIPSet
            - Arn

  LambdaUpdateWAFIPSetIamRole:
    Type: AWS::IAM::Role
    Properties: 
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement: 
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      Description: Lambda execution role
      Path: /service-role/

  LambdaUpdateWAFIPSetIamPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyName: !Sub '${AWS::StackName}-LambdaUpdateWAFIPSetIamPolicy'
      Roles:
        - Ref: LambdaUpdateWAFIPSetIamRole
      PolicyDocument: |
        {
            "Version": "2012-10-17",
            "Statement": [{
                    "Sid": "CloudWatchLogsPermissions",
                    "Effect": "Allow",
                    "Action": [
                        "logs:CreateLogGroup",
                        "logs:CreateLogStream",
                        "logs:PutLogEvents"
                    ],
                    "Resource": "arn:aws:logs:*:*:*"
                }, {
                        "Sid": "WAFPermissions",
                        "Effect": "Allow",
                        "Action": [
                            "wafv2:UpdateIPSet",
                            "wafv2:GetIPSet"
                        ],
                        "Resource": "*"
                    }
            ]
        }

  LambdaPermission:
    Type: AWS::Lambda::Permission
    Properties: 
      Action: lambda:InvokeFunction
      FunctionName:
        Fn::GetAtt: [ LambdaUpdateWAFIPSet, Arn ]
      Principal: sns.amazonaws.com
      SourceArn: arn:aws:sns:us-east-1:806199016981:AmazonIpSpaceChanged

  LambdaSNSSubscription:
    Type: AWS::SNS::Subscription
    Properties: 
      Endpoint:
        Fn::GetAtt: [ LambdaUpdateWAFIPSet, Arn ]
      Protocol: lambda
      Region: us-east-1
      TopicArn: arn:aws:sns:us-east-1:806199016981:AmazonIpSpaceChanged
Outputs:
  AWSIPSetARN:
    Description: AWS IPSet ARN
    Value: !GetAtt IPv4Set.Arn
    Export:
      Name: AWSIPSetARN

AWSTemplateFormatVersion: '2010-09-09'


Parameters:
  Prefix:
    Type: String
    Default: sample
    Description: "Fill in the name of the system name."
  Env:
    Type: String
    Default: dev
    Description: "Fill in the name of the environment."
  Scope:
    Type: String
    Default: REGIONAL
    AllowedValues: ["REGIONAL", "CLOUDFRONT"]
    Description: "Fill in the scope of waf"
  WebAclAssociationResourceArn:
    Type: String
    Default: "arn:aws:elasticloadbalancing:ap-northeast-1:XXXXXXXXXXXX:loadbalancer/app/XXXXXXXXXXXX"
    Description: Enter RegionalResource(ALB,APIGateway,AppSync) ARN or CloudFront ARN to associate with WEBACL.
  MaintenanceMode:
    Type: String
    AllowedValues: ["on", "off"]

Conditions:
  Maintenance: 
    !Equals ["on", !Ref MaintenanceMode]

Resources:
# ------------------------------------------------------------#
# WAF v2
# ------------------------------------------------------------#
  WebAcl:
    Type: AWS::WAFv2::WebACL
    Properties: 
      Name: !Sub ${Env}-${Prefix}-web-acl
      Scope: !Ref Scope
      DefaultAction:
        Allow: {}
      CustomResponseBodies:
        CustomResponseBody: 
          Content: '<h1>Blocked!!</h1>'
          ContentType: "TEXT_HTML"
      VisibilityConfig:
        CloudWatchMetricsEnabled: true
        SampledRequestsEnabled: true
        MetricName: !Sub ${Env}-${Prefix}-web-acl
      Rules:
        # ------------------------------------------------------------#
        # MaintenanceMode ON Rule
        # ------------------------------------------------------------#
        - !If 
          - Maintenance
          - Name: Whitelist-Rule
            Action:
                Block:
                  CustomResponse:
                    ResponseCode: 403
                    CustomResponseBodyKey: CustomResponseBody
            Priority: 0
            Statement:
              NotStatement:
                Statement:
                  OrStatement:
                    Statements:
                      - IPSetReferenceStatement:
                          Arn: !GetAtt WAFv2WhiteIPSet.Arn
                      - IPSetReferenceStatement: 
                          Arn:
                            Fn::ImportValue: AWSIPSetARN
            VisibilityConfig:
              CloudWatchMetricsEnabled: false
              MetricName: !Sub ${Env}-${Prefix}-Whitelist
              SampledRequestsEnabled: false
          - Ref: AWS::NoValue
          
  WebACLAssociation:
    Type: AWS::WAFv2::WebACLAssociation
    Properties:
      ResourceArn: !Ref WebAclAssociationResourceArn
      WebACLArn: !GetAtt WebAcl.Arn

  WAFv2WhiteIPSet:
    Type: "AWS::WAFv2::IPSet"
    Properties:
      Addresses:
        # White IPs
        - 0.0.0.0/1
      IPAddressVersion: IPV4
      Name: !Sub ${Env}-${Prefix}-whitelist-ips
      Scope: !Ref Scope