Sample Lambda Python function for VPC Flow Logs

lambda_kinesis_flow_logs.py

from __future__ import division
from botocore.vendored import requests
from base64 import b64decode
from collections import defaultdict
from json import dumps, loads
from zlib import decompress, MAX_WBITS

ENDPOINT = 'https://example.com/connection_sets/{bucket_time:}'
BUCKET_SECONDS = 600


def lambda_handler(event, context):
    connection_data = defaultdict(lambda: defaultdict(set))

    records = event.get('Records', [])
    for record in records:
        compressed_json = b64decode(record['kinesis']['data'])
        uncompressed_json = decompress(compressed_json, 16 + MAX_WBITS)
        input_data = loads(uncompressed_json)

        log_events = input_data.get('logEvents', [])
        for log_event in log_events:
            message = log_event['message'].split()
            srcaddr = message[3]
            dstaddr = message[4]
            start_time = int(message[10])

            bucket_time = (start_time // BUCKET_SECONDS) * BUCKET_SECONDS
            connection_data[bucket_time][srcaddr].add(dstaddr)
            connection_data[bucket_time][dstaddr].add(srcaddr)

    for bucket_time in sorted(connection_data.iterkeys()):
        output_url = ENDPOINT.format(bucket_time=bucket_time)
        output_data = {k: len(v) for k, v in connection_data[bucket_time]}
        requests.post(output_url, data=dumps(output_data))