gld

gld.py

# -*- coding: utf-8 -*-
#
# Dependencies
#   apt-get install libmysqlclient-dev python python-dev python-mysqldb
# 
# Example
#   python gld.py --help
#

import argparse
import MySQLdb
from datetime import datetime

# greylist treshold
N = 5
# default comment
DEFAULT_COMMENT = 'auto-whitelist %s' % datetime.utcnow().isoformat(' ')

def main(host, port, user, password, database):
    """
    :param host: {str} database hostname
    :param port: {int} database port
    :param user: {str} database user name
    :param password: {str} database password
    :param database: {str} database name
    """
    # connecting to database
    db = MySQLdb.connect(host=host, 
                         port=port, 
                         user=user, 
                         passwd=password, 
                         db=database)
    c = db.cursor()

    # selecting all IPs from greylist what has N bigger then N
    c.execute('SELECT ip FROM greylist WHERE n >= %d' % N)
    ips = [ip for ip, in c.fetchall()]

    if len(ips) > 0:
        # find out which IPs has not bee added to whitelist yet
        ins = ','.join(['%s'] * len(ips))
        c.execute('SELECT mail FROM whitelist WHERE mail IN (%s)' % ins,
                  tuple(ips))
        ips_to_insert = list(set(ips) - set([ip for ip, in c.fetchall()]))

        if len(ips_to_insert) > 0:
            # insert missing IPs to whitelist database  
            c.executemany('INSERT INTO whitelist (mail, comment) VALUES (%s, %s)',
                          [(ip, DEFAULT_COMMENT) for ip in ips_to_insert])

    # close connections
    c.close()
    db.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--host', dest='host', help='Database hostname')
    parser.add_argument('--port', dest='port', default=3306,
                        help='Database port')
    parser.add_argument('--user', dest='user', help='Database user name')
    parser.add_argument('--password', dest='password', help='Database password')
    parser.add_argument('--database', dest='database', help='Database name')

    main(**vars(parser.parse_args()))

    print "Done - %s" % datetime.utcnow().isoformat(' ')