Instantly share code, notes, and snippets.

@BuiDuySon /lambda_function.py Secret
Last active Nov 7, 2018

Embed
What would you like to do?
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment