root/trunk/flowmon-collector.py

#!/usr/bin/python
#
# 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 os
import sys
import SocketServer
import FlowHandler
import optparse
from signal import signal, SIGTERM
import atexit
from logging import *
from common import *
import pprint
import traceback
import gamin

def cleanup():
    info('Collector shutdown')


class handler(SocketServer.DatagramRequestHandler):
    '''Connection handler for the UDP Server'''
    def handle(self):
        flowsrc = self.client_address[0]
        # Only accept data from configured probes
        if flowsrc not in routers:
            # raise UnknownRouterException(flowsrc)
            debug('Package from non-configured system %s received' %flowsrc)
            return
        
        flowpkg = self.request[0]
        x = FlowHandler.FlowHandler(flowsrc, flowpkg)

def main():

    # CLI Options
    usage = "usage: %prog [options] arg"
    parser = optparse.OptionParser(usage)
    parser.add_option("-c", "--config", dest="configfile", default="config.py",
                      help="Use FILENAME to read initial configuration from")
    parser.add_option("-D", "--daemon",
                      action="store_true", dest="daemon",
                      help="Run as daemon")
    parser.add_option("-d", "--loglevel=debug",
                      action="store_const", const="DEBUG", dest="loglevel",
                      help="DEBUG loglevel", default="INFO")
    (options, args) = parser.parse_args()


    # Read in configuration
    execfile(options.configfile, globals())

    # Set up logging
    basicConfig(level=eval(options.loglevel),
         format='%(asctime)s %(levelname)s %(module)s %(message)s',
         filename=LOGFILE_COLLECTOR, filemode='a+')

    # Dynamically load the Parser Modules in the FlowParser sub-directory
    # This will import all files with the .py extension
    for file in os.listdir('FlowParser'):
        if file.split('.')[-1] == 'py':
            __import__('FlowParser/' + file[:-3])

    if options.daemon:
        daemonize()

    try:
        atexit.register(cleanup)
        signal(SIGTERM, lambda signum, stack_frame: exit(1))
        s = SocketServer.ThreadingUDPServer(('',PORTNO), handler)
        info('Collector startup success')
        info('UDPServer startup successful. Listening on ' + str(PORTNO) + '...')
        s.serve_forever()
    except (KeyboardInterrupt, SystemExit):
        pass
    except:
        raise
    
if __name__ == "__main__":
    main()