#!/usr/bin/env python

__version__ = "$Revision: 0.4 $"
__author__ = "Emiliano Muzzurru Renzi <emiliano.muzzurru@dada.eu>"
__status__ = "Production"


import psutil, ipaddress, sqlite3, json
from eslock import manage_config, validate_blacklist, set_logger, open_connection
import os, sys, time

geolite = '/usr/share/GeoIP/GeoLiteCity.dat'

default_config = {
    'log_file':'/var/log/eslock/producer_ip.log',
    'log_level_console':'ERROR',
    'log_level_file':'INFO',
    'log_level_main':'DEBUG',
    'log_backupcounter': 3,
    'rmq_user':'eslock',
    'rmq_pwd': 'PWD',
    'rmq_host': 'rabbitmq.it.dadainternal',
    'rmq_port': 5672,
    'rmq_vhost':'/IT',
    'rmq_exchange':'eslock',
    'rmq_heartbeat': 60,
}

def open_channel(connection):
    channel = connection.channel()
    channel.exchange_declare(exchange=config['rmq_exchange'],
                             exchange_type='fanout')
    return channel

def load_custom_wl():

    file_name = '/opt/eslock/lista_wl.txt'
    if os.path.exists(file_name):
        file = open(file_name,"r")
        list_wl = file.readlines()
        file.close()
    else:
        list_wl=[]

    return list_wl

def check_whitelist(ip):

    # 46.20.82.138 --> ETINET
    whitelist = ['93.108.244.99', '89.152.248.117', '127.0.0.0/8', '185.2.4.0/22', '81.88.57.64/27', '81.88.48.0/20', '195.110.96.0/19', '88.36.63.164', '80.172.25.69', '212.84.117.126','212.84.117.150', '212.84.117.26', '195.7.255.123', '195.7.255.124', '195.7.255.125', '213.229.80.64/28', '195.7.254.96/28', '83.71.168.34', '85.233.160.0/24', '85.233.164.0/24', '172.22.131.0/24', '194.244.25.144/28', '198.143.32.0/19', '107.154.0.0/16', '213.151.116.122', '77.226.234.3', '46.20.82.138', '83.211.81.0/24']

    for wl in whitelist:
        if ipaddress.ip_address(unicode(ip)) in ipaddress.ip_network(unicode(wl)):
        #if ipaddress.ip_address(ip) in ipaddress.ip_network(wl):
            return True

    #list_wl=load_custom_wl()
    #for wl in list_wl:
    #    if ipaddress.ip_address(unicode(ip)) in ipaddress.ip_network(unicode(wl)):
    #        return True

    return False

def check_modsecurity(list_ip):

    database='/var/cpanel/modsec/modsec.sqlite'
    conn = sqlite3.connect(database,timeout=30)
    conn.text_factory = str
    c = conn.cursor()
    QUERY = "SELECT ip from hits h WHERE h.timestamp > datetime('now','-1 minute','localtime') and h.http_status = '410' GROUP BY ip  HAVING count(ip) > 1 ORDER BY count(ip)"
    try:
        c.execute(QUERY)
        for ip in c.fetchall():
            if not check_whitelist(ip[0]):
                list_ip.add(ip[0])
    except sqlite3.Error:
        logger.error('Problem with sqlite database')
    finally:
        conn.close()
    
    logger.debug('Modsecurity Block IP :',list_ip)
    return list_ip

def check_network(list_ip):

    ip_address = [] 

    for c in psutil.net_connections(kind='inet4'):
        raddr = ""
        if (c.raddr and ( c.status == "ESTABLISHED" or c.status == "SYN_RECV")):
            ip_address.append(c.raddr[0])
    ip_address = dict((x, ip_address.count(x)) for x in set(ip_address))
    #logger.debug(ip_address)
    ip_address = dict((k, v) for k, v in ip_address.items() if v >= 80)

    for key, value in ip_address.iteritems():
        if not check_whitelist(key):
    #       list_ip.add(key)
	    logger.debug('IP :',key,' Connections: ',value)

    return list_ip

if __name__ == '__main__':
     config = manage_config('/etc/eslock/eslock_ddos.yaml', default_config)
     logger = set_logger(config['log_file'],
                         config['log_level_main'],
                         config['log_level_file'],
                         config['log_level_console'],
                         config['log_backupcounter']) #inizializziamo il logger
     logger.info('START '+__version__)

     while True:
         list_ip = set()
         #list_ip=check_network(list_ip)
         list_ip=check_modsecurity(list_ip)
         list_ip = list(list_ip)

         if list_ip:
             while 1:
                 connection = open_connection(config['rmq_user'],
                                              config['rmq_pwd'],
                                              config['rmq_host'],
                                              config['rmq_port'],
                                              config['rmq_vhost'],
                                              config['rmq_heartbeat'])
                 if connection:
                      break
                 logger.error('Failed connection to rmq, waiting 30 secs')
                 time.sleep(30)

             logger.info('%(rmq)s connected'%{'rmq':config['rmq_host']})
             channel = open_channel(connection)
             logger.info('Rabbitmq channel opened')
             #Build message for rabbitmq
             message = dict()
             message['type'] = 'blacklist'
             message['name'] = 'tmp'
             message['comment'] = 'Modsecurity LHCP'
             message['ip_list'] = list_ip
             message['cluster'] = ['lhcp','opus']
             json_message = json.dumps(message)
             json_message = json.dumps(message, indent=2,
                                       sort_keys=True, separators=(',', ': '))

             channel.basic_publish(exchange=config['rmq_exchange'],
                                   routing_key='',
                                   body=json.dumps(message))
             logger.info('Message sent, len: '+str(len(message)))
             for ip in list_ip:
                  logger.info('IP Block : '+ ip)
             connection.close()

	 time.sleep(30)