Lambda function to update whitelist with newrelic IPs

update_newrelic_whitelist.py

import os
import json
import boto3
import urllib3

# Set the following in the Lambda Environment variables
IP_SET_ID = os.environ.get('IP_SET_ID')
Special_IPs = os.environ.get('SPECIAL_IPS').split(',')


def get_change_token():
    client = boto3.client('waf-regional')
    return client.get_change_token()['ChangeToken']


def get_latest_minion_ips():
    http = urllib3.PoolManager()
    response = http.request(
        'GET', 'http://s3.amazonaws.com/nr-synthetics-assets/nat-ip-dnsname/production/ip.json')
    if response.status == 200:
        nr_ips = json.loads(response.data)
    all_ips = [nr_ips[location] for location in nr_ips]
    return sum(all_ips, [])


def get_ip_set_ips():
    client = boto3.client('waf-regional')
    response = client.get_ip_set(IPSetId=IP_SET_ID)
    if 'IPSet' in response:
        set_ips = [ip['Value'].rstrip(
            '/32') for ip in response['IPSet']['IPSetDescriptors'] if ip['Value'] not in Special_IPs]
    return set_ips


def update_waf_rule():
    client = boto3.client('waf-regional')
    nr_ips = get_latest_minion_ips()
    waf_ips = get_ip_set_ips()

    to_delete = [{'Action': 'DELETE', 'IPSetDescriptor': {
        'Type': 'IPV4', 'Value': '{}/32'.format(v)}} for v in waf_ips if v not in nr_ips]
    to_add = [{'Action': 'INSERT', 'IPSetDescriptor': {'Type': 'IPV4',
                                                       'Value': '{}/32'.format(v)}} for v in nr_ips if v not in waf_ips]
    updates = to_add + to_delete

    response = client.update_ip_set(
        IPSetId=IP_SET_ID,
        ChangeToken=get_change_token(),
        Updates=updates
    )


def lambda_handler(event, context):
    update_waf_rule()


if __name__ == "__main__":
    update_waf_rule()