Mentat

#!/usr/bin/env python3

"""Mentat IP aggregation correctness tester"""

import timeit

import random

import ipaddress

import sys

import psycopg2

from psycopg2 import sql

DBNAME = 'mentat_events'

BITS = 32

def main():

    start = timeit.default_timer()

    count = 0

    con = psycopg2.connect(dbname=DBNAME)

    cursor = con.cursor()

    raw_query = "SELECT id FROM (SELECT id FROM events WHERE %s && ANY({srctgt}) UNION ALL SELECT id FROM events WHERE ({aggr} && %s AND %s && ANY({srctgt}))) AS sub GROUP BY id HAVING count(*) = 1"

    try:

        while True:

            count = count + 1

            ip = str(ipaddress.ip_address(random.getrandbits(BITS)))

            for srctgt, aggr in [('source_ip', 'source_ip_aggr_ip4'), ('target_ip', 'target_ip_aggr_ip4')]:

                query = sql.SQL(raw_query).format(srctgt=sql.Identifier(srctgt), aggr=sql.Identifier(aggr))

                cursor.execute(query, (ip, ip, ip))

                ids = cursor.fetchall()

                con.rollback()

                if ids:

                    print('\nIP: %s (%s) mismatch in %s' % (ip, srctgt, [eid for eidt in ids for eid in eidt]))

                    raise EnvironmentError

            print('.', end='')

            sys.stdout.flush()

    except (KeyboardInterrupt, EnvironmentError):

        pass

    finally:

        print('\nTotal tested %d addresses in %fs' % (count, timeit.default_timer() - start))

if __name__ == '__main__':

    main()