niku-server/app.py

import chopy.search
import chopy.db
import traceback
import config
import json
import urllib
import os

def application(environ, start_response):
    path = environ['PATH_INFO']
    qs = environ['QUERY_STRING']

    try:
        if path == '/':
            start_response('200 OK', [('Content-type', 'text/plain')])
            yield 'BEEP BOOP WELCOME TO THE API'

        elif path == '/api/search':
            start_response('200 OK', [('Content-type', 'text/plain')])
            for cid in do_search(environ['QUERY_STRING']):
                yield '%d\r\n' % cid

        elif path == '/api/sort':
            start_response('200 OK', [('Content-type', 'text/plain')])
            for cid in do_sort(environ['QUERY_STRING']):
                yield '%d\r\n' % cid

        elif path == '/api/metadata':
            start_response('200 OK', [('Content-type', 'text/plain')])
            for conn in do_lookup(environ['QUERY_STRING']):
                yield '%s\r\n' % json.dumps(conn)

        elif path in extra_paths:
            start_response('200 OK', [('Content-type', 'text/plain')])
            for x in extra_paths[path](qs):
                yield x

        elif path.startswith('/pcap'):
            start_response('200 OK', [('Content-type', 'application/x-octet-stream')])
            pathkeys = path.split('/')
            pathkeys[0:2] = [config.split_dir]
            filepath = os.path.join(*pathkeys)
            yield open(filepath).read()

        else:
            start_response('404 NOT FOUND', [('Content-type', 'text/plain')])
            yield 'The URL you requested could not be serviced by the application.\r\n'
    except: # pylint: disable=bare-except
        tb = traceback.format_exc()
        start_response('500 INTERNAL SERVER ERROR', [('Content-type', 'text/plain')])
        yield 'The server encountered an error during execution. Here is some debug information.\r\n\r\n'
        yield 'Environment:\r\n'
        for key in environ:
            yield '  %s=%s\r\n' % (key, environ[key])
        yield '\r\n'
        yield tb.replace('\n', '\r\n')
        print tb

def nytes_to_bit_string(n):
    bin_str = "".join(bin(ord(c))[2:].zfill(8) for c in n)
    #num_bits = (len(n) * 8) % 9
    #return bin_str[:len(bin_str) - num_bits]
    return bin_str

def bytes_to_nytes(b):
    bin_str = "".join(bin(ord(c))[2:].zfill(9) for c in b)
    bin_str = bin_str.zfill(len(bin_str) + ((8 - (len(bin_str) % 8)) % 8))
    return "".join(chr(int(bin_str[i:i+8], 2)) for i in xrange(0, len(bin_str), 8))

def do_search(query_string):
    _, session = chopy.db.connect(config.database)
    data = json.loads(urllib.unquote(query_string))
    if 'search_regex' in data and data['search_regex']:
        # OH BOY
        bstr = str(data['search_regex'])
        bitstr = ''.join(bin(ord(c))[2:].zfill(9) for c in bstr)
        possible_values = []
        for i in xrange(9):
            trunc_bitstr = bitstr[i:]
            possible_values.append(''.join(chr(int(trunc_bitstr[j:j+8], 2)) for j in xrange(0, len(trunc_bitstr)-7, 8)))

        data['search_regex'] = '|'.join(''.join('\\x%02x' % ord(c) for c in x) for x in possible_values)

    if type(data) is not dict:
        raise ValueError("Query string for /api/search must be a json dictionary")
    data['index_dir'] = config.index_dir
    for cid in chopy.search.search(session, **data):
        yield cid

def do_sort(query_string):
    _, session = chopy.db.connect(config.database)
    data = json.loads(urllib.unquote(query_string))
    if type(data) is not dict:
        raise ValueError("Query string for /api/sort must be a json dictionary")
    cids = data['ids']
    key = data['order_by']

    q = session.query(chopy.db.Connection.id) \
        .filter(chopy.db.Connection.id.in_(cids)) \
        .order_by(getattr(chopy.db.Connection, key))
    for conn in q:
        yield str(conn.id)

def do_lookup(query_string):
    _, session = chopy.db.connect(config.database)
    data = json.loads(urllib.unquote(query_string))
    if type(data) is not list:
        raise ValueError("Query string for /api/metadata must be a json list")
    q = session.query(chopy.db.Connection) \
        .filter(chopy.db.Connection.id.in_(data)) \
        .options(chopy.db.joinedload('tags'))
    for conn in q:
        out = conn.dict()
        out['tags'] = [x.text for x in out['tags']]
        #out['tags'] = conn.tags
        yield out

# this is... the least readable code I've ever written

def make_setter(cls, args, updatable_args):
    def set_something(query_string):
        _, session = chopy.db.connect(config.database)
        data = json.loads(urllib.unquote(query_string))
        if type(data) is not dict:
            raise ValueError("Query string for set API must be a json dictionary")
        if set(args) != set(data):
            raise ValueError("Expected arguments: %r" % args)

        obj = None
        if updatable_args:
            q = session.query(cls)
            for arg in args:
                if arg not in updatable_args:
                    q = q.filter(getattr(cls, arg) == data[arg])
            obj = q.first()
        if obj is not None:
            for arg in updatable_args:
                setattr(obj, arg, data[arg])
        else:
            obj = cls(**data)
            session.add(obj)

        session.commit()
        yield str(data)
    return set_something

def make_getter(cls, result_args):
    def get_something(query_string): # pylint: disable=unused-argument
        _, session = chopy.db.connect(config.database)
        q = session.query(*[getattr(cls, arg) for arg in result_args]).distinct()
        for r in q:
            yield '%s\r\n' % json.dumps({arg: getattr(r, arg) for arg in result_args})
    return get_something

def make_deleter(cls, args):
    def del_something(query_string):
        _, session = chopy.db.connect(config.database)
        data = json.loads(urllib.unquote(query_string))
        if type(data) is not dict:
            raise ValueError("Query string for delete API must be a json dictionary")
        if set(args) != set(data):
            raise ValueError("Expected arguments: %r" % args)
        q = session.query(cls)
        for arg in args:
            q = q.filter(getattr(cls, arg) == data[arg])
        q.delete()
        session.commit()
        yield '%s\r\n' % data
    return del_something


tags_setter = make_setter(chopy.db.Tag, ['connection', 'text'], [])
tags_getter = make_getter(chopy.db.Tag, ['text'])
tags_deleter = make_deleter(chopy.db.Tag, ['connection', 'text'])

services_setter = make_setter(chopy.db.ServiceName, ['protocol', 'host', 'port', 'name'], ['name'])
services_getter = make_getter(chopy.db.ServiceName, ['protocol', 'host', 'port', 'name'])
services_deleter = make_deleter(chopy.db.ServiceName, ['protocol', 'host', 'port', 'name'])

hosts_setter = make_setter(chopy.db.HostName, ['boot_time', 'name'], ['name'])
hosts_getter = make_getter(chopy.db.HostName, ['boot_time', 'name'])
hosts_deleter = make_deleter(chopy.db.HostName, ['boot_time', 'name'])

extra_paths = {
        '/api/tags/get': tags_getter, '/api/tags/set': tags_setter, '/api/tags/del': tags_deleter,
        '/api/services/get': services_getter, '/api/services/set': services_setter, '/api/services/del': services_deleter,
        '/api/hosts/get': hosts_getter, '/api/hosts/set': hosts_setter, '/api/hosts/del': hosts_deleter
}