append primary public IP to VPC flow logs entries via lambda

lambda_function.py

from __future__ import print_function

import json
import base64
import gzip
import re
import sys
import boto3
from StringIO import StringIO
from pprint import pprint

print('Loading function')


def lambda_handler(event, context):
    #get the log messages
    decoded_data = event['awslogs']['data'].decode('base64')
    log_events = json.loads(gzip.GzipFile(fileobj=StringIO(decoded_data)).read())
    match = re.compile('eni-[^-]*')
    eni_id = match.findall(log_events['logStream'])[0]
    # find ENI public IP
    client = boto3.client('ec2', region_name='us-west-2')
    response = client.describe_network_interfaces(NetworkInterfaceIds=[eni_id])
    if 'PublicIp' in response['NetworkInterfaces'][0]['PrivateIpAddresses'][0]['Association']:
        primary_public_ip = response['NetworkInterfaces'][0]['PrivateIpAddresses'][0]['Association']['PublicIp']
    else:
        primary_public_ip = None
    print('public IP: {}'.format(primary_public_ip))
    print('ENI ID from regex: {}'.format(eni_id))
    for log_entry in log_events['logEvents']:
        log_entry['extractedFields']['public_ip'] = primary_public_ip
        # do some more logic here with your newly defined data structure