used to distributing keys to different systems uniquely.. takes on standard input and distributes it uniquely to all clients hitting the /k endpoint

key_dist.py

#!/usr/bin/env python

import atexit
import functools
import gzip
import itertools
import logging
import os
import socket
import sys
import time
import threading

from cStringIO import StringIO as IO
from logging.handlers import RotatingFileHandler
from optparse import OptionParser
from Queue import Queue, Empty

from flask import after_this_request, request, Flask

Q_SZ    = 2**26
LOG_FMT = '%(asctime)s %(module)8s:%(lineno)-5d [%(levelname)-8s] %(message)s'
LOG_DIR = '/var/tmp'

# from http://www.flaskapi.org/api-guide/status-codes/
STATUS_KEYS_INCOMPLETE = 206    # HTTP_206_PARTIAL_CONTENT
STATUS_KEYS_FINISHED = 410      # HTTP_410_GONE
STATUS_KEY_AMOUNT_NEEDED = 411  # HTTP_411_LENGTH_REQUIRED
STATUS_KEY_FOUND = 412          # HTTP_412_PRECONDITION_FAILED

shared_state = {}
shared_state['key_found'] = False
shared_state['keys_finished'] = False

queue  = Queue(maxsize=Q_SZ)
q_thr  = threading.Thread()
lock   = threading.Lock()
lock2  = threading.Lock()

class InputThread(threading.Thread):
    def __init__(self, q, l):
        threading.Thread.__init__(self)
        self.queue = q
        self.logger = l
        self.should_exit = False

        global shared_state
        global lock

        self.shared_state = shared_state
        self.lock = lock
        self.count = 0

    def run(self):
        for line in iter(sys.stdin.readline, None):
            if self.should_exit:
                return

            if line:
                self.queue.put(line)
                self.count += 1
            else:
                with self.lock:
                    self.logger.critical('all {0} keys have been loaded into the queue'.format(self.count))
                    self.shared_state['keys_finished'] = True
                    self.should_exit = True

def create_app():
    app = Flask(__name__)

    # based on http://stackoverflow.com/questions/14384739/how-can-i-add-a-background-thread-to-flask
    def intr_q_thr():
        global q_thr
        q_thr.should_exit = True

    def start_q_thr():
        # Do initialisation stuff here
        global q_thr
        # Create your thread
        q_thr = InputThread(q=queue, l=app.logger)
        q_thr.start()

    # Initiate
    start_q_thr()

    # When you kill Flask (SIGTERM), clear the trigger for the next thread
    atexit.register(intr_q_thr)
    return app


# from http://flask.pocoo.org/snippets/122/
def gzipped(f):
    @functools.wraps(f)
    def view_func(*args, **kwargs):
        @after_this_request
        def zipper(response):
            accept_encoding = request.headers.get('Accept-Encoding', '')

            if 'gzip' not in accept_encoding.lower():
                return response

            response.direct_passthrough = False

            if (response.status_code < 200 or
                response.status_code >= 300 or
                'Content-Encoding' in response.headers):
                return response
            gzip_buffer = IO()
            gzip_file = gzip.GzipFile(mode='wb', 
                                      fileobj=gzip_buffer)
            gzip_file.write(response.data)
            gzip_file.close()

            response.data = gzip_buffer.getvalue()
            response.headers['Content-Encoding'] = 'gzip'
            response.headers['Vary'] = 'Accept-Encoding'
            response.headers['Content-Length'] = len(response.data)

            return response

        return f(*args, **kwargs)

    return view_func


app   = create_app()


@app.route('/')
def cap():
    requesting_system = request.remote_addr
    hostname = socket.gethostbyaddr(requesting_system)
    if hostname:
        requesting_system = hostname[0]

    app.logger.info('sending capture file ({0} bytes) to {1}'.format(len(app.config['FILE']), requesting_system))
    return app.config['FILE']


@app.route('/k')
@gzipped
def next_keys():
    global shared_state
    global lock2

    requesting_system = request.remote_addr
    hostname = socket.gethostbyaddr(requesting_system)
    if hostname:
        requesting_system = hostname[0]

    with lock2:
        if shared_state['key_found']:
            app.logger.warning('system {0} tried to request more keys but key already has been found'.format(requesting_system))
            return '', STATUS_KEYS_FINISHED

    num = int(request.args.get('n', 0))
    if not num:
        app.logger.error("system {0} tried to request keys but didn't say how many".format(requesting_system))
        return 'Please specify number of keys requested with ?n=', STATUS_KEY_AMOUNT_NEEDED

    keys = []
    with lock2:
        if not shared_state['keys_finished']:
            keys = [queue.get() for _ in xrange(num)]
        else: # keys have all been read, so we want to make sure we don't block if we empty queue
            try:
                keys = [queue.get(block=False)]
                while len(keys) < num:
                    keys.append(queue.get(block=False))
            except Empty:
                pass

    # we didn't get any keys
    if not keys:
        app.logger.error('have no more keys to send to {0}'.format(requesting_system))
        return '', STATUS_KEYS_FINISHED

    # we got some of the keys but not all
    if len(keys) < num:
        app.logger.warning('system {0} wanted {1} keys but only could send {2}: {3}-{4}'.format(requesting_system, num, len(keys), keys[0].strip(), keys[-1].strip()))
        return ''.join(keys), STATUS_KEYS_INCOMPLETE

    # seems we got all the keys we asked for
    app.logger.info('sent {0} keys to system {1}: {2}-{3}'.format(num, requesting_system, keys[0].strip(), keys[-1].strip()))
    return ''.join(keys)


@app.route('/x')
def cracked():
    global shared_state

    with lock2:
        shared_state['key_found'] = True

    requesting_system = request.remote_addr
    hostname = socket.gethostbyaddr(requesting_system)
    if hostname:
        requesting_system = hostname[0]

    key = request.args.get('k', '')
    if key:
        app.logger.critical('key cracked by {0}: {1}'.format(requesting_system, key))
    else:
        app.logger.critical('key cracked by {0}'.format(requesting_system))

    return ''


if __name__ == "__main__":
    usage = 'usage: %prog -f file.cap'
    parser = OptionParser(usage=usage)
    parser.add_option('-f', '--file', dest='filename',
                      help='select capture file to send to clients')
    (options, args) = parser.parse_args()

    if not options.filename:
        parser.print_help()
        sys.exit(1)
    else:
        try:
            with open(options.filename) as fh:
                app.config['FILE'] = fh.read()
        except IOError:
            print('ERROR: cannot open: {0}'.format(options.filename))
            sys.exit(1)

    handler = RotatingFileHandler(os.path.join(LOG_DIR, os.path.basename(__file__)[:-3]+'.log'), maxBytes=10000000, backupCount=4)
    handler.setFormatter(logging.Formatter(LOG_FMT))

    app.logger.addHandler(handler)
    app.logger.setLevel(logging.DEBUG)

    app.logger.critical('will send the following file to clients: {0}'.format(options.filename))

    app.run(threaded=True, port=18888)