full.py

#!/usr/bin/env python

from sys import exit
import json
import datetime
import requests
from uuid import uuid4
import psycopg2
import manicthief

f = open('analytics.json','r')
data = f.read()
f.close()
analytics = json.loads(data)

cache_conn = psycopg2.connect("dbname=manicthief user=manicthief password=test host=127.0.0.1")
cache_cursor = cache_conn.cursor()

tag_conn = psycopg2.connect("dbname=manicthief user=manicthief password=test host=127.0.0.1")
tag_cursor = tag_conn.cursor()

insert_tag_buffer = []
update_tag_buffer = []

s = manicthief.Utils.stat

# Helper function for issuing queries
def get_new_records(query):
    
    statement = query
    cache_cursor.execute(statement)
    results = cache_cursor.fetchall()
    return results

def generate_insert_statement(tag_name, ip, category, confidence, intention, tag_id):

    insert_statement = tag_cursor.mogrify('''INSERT INTO analytics ( ip, tag_name, category, confidence, intention, tag_id, created, updated, tag_uuid, negated ) VALUES ( %s, %s, %s, %s, %s, %s, %s, %s, %s, %s )''',
            ( ip, tag_name, category, confidence, intention, tag_id, 
                datetime.datetime.utcnow(), 
                datetime.datetime.utcnow(),
                str(uuid4()),
                False))
    #s("analyticsd.generate.insert")
    return insert_statement+';'

def get_previous_records(tag_name):
    tag_cursor.execute('''SELECT * FROM analytics WHERE tag_name = '{TAG_NAME}' AND updated >= now() - interval '3 days'
                            AND negated IS false ORDER BY updated DESC'''.format(TAG_NAME=tag_name))
    s("analyticsd.get.record")
    return tag_cursor.fetchall()

def convert_to_sql_list(l):
    x = '('
    for i in l:
        x = x+"'"+i+"',"
    x = x[:-1]
    x = x+')'
    return x

def hack_convert(records):
    template = "select * from analytics where tag in ({CODE})"
    root = "(VALUES "
    for rec in records:
        root = root + "('"+rec+"'::uuid),"

    root = root[:-1]

    root = root + ")"
    return root

TEMPLATE = """UPDATE analytics SET updated = now() WHERE tag_uuid IN {UUIDS}"""
def update_records(tag_id, uuids):
    "running against", len(uuids)
    #structured_uuids = convert_to_sql_list(uuids)
    structured_uuids = hack_convert(uuids)
    print structured_uuids
    query = TEMPLATE.format(UUIDS=structured_uuids, IPS=structured_uuids)
    print 'executing function - psql call execute'
    #print tag_cursor.mogrify(query)
    #exit("manual exit")
    tag_cursor.execute(query)
    print 'finish running exec'
    #s("analyticsd.update.record")
    return

def bulk_insert(statements):
    insert_blob = ' '.join(statements)
    tag_cursor.execute(insert_blob)
    s("analyticsd.insert.record")
    return

def main():
  
    #while True:
        for analytic in analytics:

            # Print the analytic we're running
            print('[+] %s' % analytic['name'])

            # grab all the records of that type that were not negated
            # and updated sooner than 3 days ago
            previous_records = get_previous_records(analytic['name'])
            pr_ips = []
            for i in previous_records:
                pr_ips.append(i[0])

            # grab all the new IPs yielded by the analytics
            new_records = get_new_records(analytic['query'])
            nr_ips = []
            for i in new_records:
                nr_ips.append(i[0])

            # which records match existing ones?
            ips_to_update = set(pr_ips).intersection(nr_ips)

            print('[+] %s: Updating %d IPs' % (analytic['name'], len(ips_to_update)))
            
            # get the uuids associated with the records that need to be updated
            uuids_to_update = []
            for i in previous_records:
                s("analyticsd.prepare.update")
                uuids_to_update.append(i[8])
            
            if len(ips_to_update) > 0:
                
                while len(uuids_to_update) > 0:
                    print 'updating', len(uuids_to_update[:1000])
                    update_records(analytic['id'], uuids_to_update[:1000])
                    del uuids_to_update[:1000]
                tag_conn.commit()
            
            ips_to_add = []
            for i in nr_ips:
                if i not in pr_ips:
                    ips_to_add.append(i)

            print('[+] %s: Adding %d IPs' % (analytic['name'], len(ips_to_add)))
            insert_buffer = []
            if ips_to_add > 0:
                for ip in ips_to_add:
                    s("analyticsd.prepare.add")
                    insert_buffer.append(generate_insert_statement(analytic['name'], ip, analytic['category'], analytic['confidence'], analytic['intention'], analytic['id']))
                   
            while len(insert_buffer) > 0:
                bulk_insert(insert_buffer[:100])
                del insert_buffer[:100]
            tag_conn.commit()

            s("analyticsd.analytic.completed")

if __name__ == '__main__':
    main()