lambda_function.py

from __future__ import print_function

import os
import re
import json
import urllib
import boto3
import commands
import gzip
import ast
import datetime
import hashlib
from collections import Counter

## Environmental settings
waf_ip_set_id = os.environ['WAF_IP_SET_ID']
black_list_threshold_count = int(os.environ['BLACK_LIST_THRESHOLD_COUNT'])
slack_incomming_hook_url = os.environ['SLACK_INCOMMING_HOOK_URL']

## codes
s3 = boto3.client('s3')
waf = boto3.client('waf-regional', region_name='ap-northeast-1')

def lambda_handler(event, context):
    # Get the object from the event and show its content type
    bucket = event['Records'][0]['s3']['bucket']['name']
    key = urllib.unquote_plus(event['Records'][0]['s3']['object']['key'].encode('utf8'))
    try:
        # open file, unzip
        s3.download_file(bucket, key, '/tmp/file.gz')
        f = gzip.open('/tmp/file.gz', 'rb')
        content = f.read().strip()
        f.close

        # create black list
        ip_address_list = to_ip_address_list(content.decode('utf-8'))
        black_list = to_ip_black_list(ip_address_list)

        if len(black_list):
            print(u'block threshold: {threshold} '.format(threshold = black_list_threshold_count))

            # create ip set for waf
            update_ip_set=map(to_ip_set_for_waf, black_list)

            # waf setting (side-effect function)
            waf_block(update_ip_set)

            # notification (side-effect functions)
            slack_notification(bucket, black_list)
        else:
            print(u'nothing to do.')
            pass
        return black_list
    except Exception as e:
        print(e)
        raise e

# from: String log file -> to: IP Address List.
def to_ip_address_list(file_string):
    result_list=[]
    for r in file_string.split('\n'):
        try:
            rlogCompile = re.compile('^(.*?) \"(.*?)\" \"(.*?)\".*$')
            rlogs = rlogCompile.search(r)
            items = rlogs.group(1).split()
            ip = re.sub(r':.*$', "", items[3])
            result_list.append(ip)
        except Exception as e:
            print(u'ignore: ',r ,e)
    return result_list

# from: IP address list -> to: black list
def to_ip_black_list(ip_list):
    black_list={}
    counter=Counter(ip_list)
    for ip, cnt in counter.most_common():
        if cnt > black_list_threshold_count:
            black_list[ip] = cnt
        else:
            pass
    return black_list

# from: IP address list -> to: WAF update request body.
def to_ip_set_for_waf(blocked_ip):
    ip_set = {'Action': 'INSERT'}
    descripter={'Type': 'IPV4'}
    descripter['Value']= str(blocked_ip) + '/32'
    ip_set['IPSetDescriptor']=descripter
    return ip_set

def waf_block(update_ip_set):
    token = waf.get_change_token()
    waf.update_ip_set(
        IPSetId=waf_ip_set_id,
        ChangeToken=token['ChangeToken'],
        Updates=update_ip_set
    )
    print(u'WAF:', waf.get_ip_set(IPSetId=waf_ip_set_id))
    return

def slack_notification(bucket, blocked_ips):
    url = slack_incomming_hook_url
    ips = ''
    for ip, cnt in blocked_ips.items():
        ips = ips + "IP: " + str(ip) + " / " + str(cnt) + " request/5min\n"
    body = "[S3_Bucket_Name]: {bucket}\n[Blocked IPs]\n {ips}".format(bucket = bucket, ips = ips)
    params = urllib.urlencode({"payload": {"text": body, "icon_emoji": ":hankey:", "username": "blockbot"}})
    f = urllib.urlopen(url, params)
    print(f.read())
    return

def to_ip_address_list(file_string):
    result_list=[]
    for r in file_string.split('\n'):
        try:
            rlogCompile = re.compile('^(.*?) \"(.*?)\" \"(.*?)\".*$')
            rlogs = rlogCompile.search(r)
            items = rlogs.group(1).split()
            ip = re.sub(r':.*$', "", items[3])
            result_list.append(ip)
        except Exception as e:
            print(u'ignore: ',r ,e)
    return result_list