root/trunk/FlowHandler.py

#
# Copyright (C) 2007-2008 Red Hat, Inc.
# Author: Andreas Thienemann <athienem@redhat.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Library General Public License as published by
# the Free Software Foundation; version 2 only
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Library General Public License for more details.
#
# You should have received a copy of the GNU Library General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
# Copyright 2004, 2005 Red Hat, Inc.
#
# AUTHOR: Andreas Thienemann <athienem@redhat.com>
#

import struct
import os
import sys
import IPy
import binascii
import pickle
import hashlib
import shutil
import pprint
from logging import *

## Read the config
execfile('config.py')

class AbstractFlowParserMeta(type):
    def __init__(cls, name, bases, attrs):
        if not hasattr(cls, 'plugins'):
            # This branch only executes when processing the mount point itself.
            # So, since this is a new plugin type, not an implementation, this
            # class shouldn't be registered as a plugin. Instead, it sets up a
            # list where plugins can be registered later.
            cls.plugins = {}
        else:
            # This must be a plugin implementation, which should be registered.
            # Simply appending it to the list is all that's needed to keep
            # track of it later.
            cls.plugins[cls.name] = cls


class AbstractFlowParser(object):
    """Plugin registry.
    See http://gulopine.gamemusic.org/2008/jan/10/simple-plugin-framework/"""
    __metaclass__ = AbstractFlowParserMeta

    def __init__(self, flowsrc, pkgdata):
        super(AbstractFlowParser, self).__init__()
        self.flowsrc = flowsrc
        self.pkgdata = pkgdata
        self.flowpkg_decode = []

    def writeout(self, flow_data=''):

        debug('Writeout called') 
        # Sanity check, do we have an empty flowset (only a header)?
        if flow_data == '' and len(self.flowpkg_decode) == 1:
            debug('Writeout canceled') 
            return

        # We're not being passed anything specific to write, get the default self.flowpkg_decode data
        if flow_data == '':
            flow_data = self.flowpkg_decode
            # Delete the committed data
            del(self.flowpkg_decode)

        # Build an unique filename
        filename = hashlib.sha1()
        filename.update(str(flow_data[0]['sourceIP'])+'-'+str(flow_data[0]['unixtime'])+'-'+str(flow_data[0]['sequence']))
        filename = filename.hexdigest() + '.pkl'

        # Pickle this
        fp = open(os.path.join(SPOOLDIR, filename + '.tmp'), "w")
        pickle.dump(flow_data, fp)
        fp.close

        # Move the file to it's final name to guarantee atomicity
        shutil.move(os.path.join(SPOOLDIR, filename + '.tmp'), os.path.join(SPOOLDIR, filename))
        

class FlowHandler(object):

    def __init__(self, flowsrc, pkgdata):
        super(FlowHandler, self).__init__()
       
        self.flowsrc = flowsrc
        self.pkgdata = pkgdata

        klass = AbstractFlowParser.plugins[routers[flowsrc]]
        parser = klass(flowsrc, pkgdata)

        # Execute the parser
        parser.parse()
        # Flush it's cache to disk early
        parser.writeout()