np_grabrec.py

该软件根据网络数据生成NetFlow记录。NetFlow可用于网络规划、负载均衡、安全监控等

#! /usr/bin/env python
###############################################################################
#                                                                             #
#   Copyright 2005 University of Cambridge Computer Laboratory.               #
#                                                                             #
#   This file is part of Nprobe.                                              #
#                                                                             #
#   Nprobe is free software; you can redistribute it and/or modify            #
#   it under the terms of the GNU General Public License as published by      #
#   the Free Software Foundation; either version 2 of the License, or         #
#   (at your option) any later version.                                       #
#                                                                             #
#   Nprobe 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 General Public License for more details.                              #
#                                                                             #
#   You should have received a copy of the GNU General Public License         #
#   along with Nprobe; if not, write to the Free Software                     #
#   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA #
#                                                                             #
###############################################################################


##############################################################################
## 
##
## 
## 
##
## 

############################################################################ 

import string
import os
import sys
#import nprobe
from nprobe import np_file, intoa_string
from np_filerec import FileRec
from Tkinter import *
from np_notesel import Sellist
#from np_TCPDisplay import *
#import np_obnode
import np_WebHost
import np_TCPConn
import np_HTTPTrans
import np_Obtree
import np_TCPDisplay
import np_TScreen
from np_TCP import TCP_Machine, TCPNoModel
from TCP_Imp import TCP_Imps
#import np_filerec
from sys import argv
import getopt
from signal import *
import tkSimpleDialog, tkFileDialog
from cPickle import dump, load
import types
        
        

#############################################################################
        
DEF_GRAB_FILE = '/usr/groups/nprobe/jch1003/grabs/grabs1'

#global - initially none to avoid file dialogue if no grabbing
grabobj = None

# control pickle format
BINARY_DUMP = 1

#############################################################################


def grab(rec):

    #print 'grab'
    global grabobj

    if not grabobj:
        grabobj = GrabRec('append')
        
        
    grabobj.grab(rec)
    
        
        

#############################################################################

class GrabRec:

    def __init__(self, mode='append', fnm=None):
        
        if fnm:
            self.fnm = fnm
        else:
            self.fnm = DEF_GRAB_FILE
        self.gfile = None
        self.grabs = []

        self.getfile(mode)


    def getfile(self, mode='append'):

        if self.fnm != DEF_GRAB_FILE:
            #c.l. grab file given
            try:
                self.gfile = open(self.fnm, mode[0])
                self.root = None
            except IOError, s:
                print 'Can\'t open file', self.fnm, ':', s
                self.gfile = None

            return self.gfile

        # otherwise use dialogue
        base = os.path.split(self.fnm)[0]
        while 1:
            try:
                prompt = 'Open grab file'
                fnm = os.path.split(self.fnm)[1]
                if mode == 'append':
                    #fnm = tkSimpleDialog.askstring('Grab', prompt,
                                                   #initialvalue=fnm)
                    fnm =  tkFileDialog.asksaveasfilename( \
                        defaultextension='.grabs', title='Save grabs to?',
                        initialdir=base,
                        filetypes=[('np_Grab files', '.grabs'),
                                   ('All files', '.*')])
                    if not fnm:
                        return None
                else:
                    #fnm = raw_input('Load grab file?')
                    self.root = Tk()
                    self.root.winfo_toplevel().title('Grabs:')
                    fnm =  tkFileDialog.askopenfilename( \
                        title='Grab', initialdir=base,     
                        filetypes=[('np_Grab files', '.grabs'),
                                   ('All files', '.*')])
                    if not len(fnm):
                        #fnm = DEF_GRAB_FILE
                        sys.exit(0)
                fnm = os.path.join(base, fnm)
                self.gfile = open(fnm, mode[0])
                self.fnm = fnm
                break
            except IOError, s:
                print 'Can\'t open file', fnm, ':', s
                self.gfile = None

        return self.gfile

            


    def grab(self, ob):

        #print 'grab got',
        #print rec

        if not self.gfile:
            f = self.getfile('append')
            if not f:
                return

        prompt = 'Grab %s object - comment?' % (ob.Class)
        comment = tkSimpleDialog.askstring('Grab', prompt,
                                                    initialvalue='')
        if comment == None: # cancel
            return

        rec = ob.FileRec

        # make grabrec list of lists for compat with 'WebClient' class
        if rec.type == 'TCPConn':
            rec.recsets = [rec.recsets]
            rooturl = None
        elif rec.type.find('Web') == 0:
            rooturl = ob.rooturl
            
        newrec = [rec.type, comment, rooturl, []]
        # save filerec data rather than filerec
        for r in rec.recsets:
            set = []
            newrec[3].append(set)
            for rr in r:
                set.append((rr[0], rr[1], rr[2].fnm, rr[3]))


        #self.grabs.append(newrec)

        # dump immediately in case we crash
        dump([newrec], self.gfile, BINARY_DUMP)
        self.gfile.flush()
        
    def close(self):

##         # print 'grabs close'
##         # if len(self.grabs):
##             # dump(self.grabs, self.gfile, BINARY_DUMP)
##             # self.gfile.close()

        self.gfile.close() 

    def write_log(self, txt):

        self.log.append(txt)

    def load(self):

        files = {}
        #self.getfile('read')

        while 1:
            try:
                read = load(self.gfile)
                #print 'load'
                for recs in read:
                    if len(recs) == 3: # old type - no rooturl
                        sets = recs[2]
                        rooturl = None
                    elif len(recs) == 4: # new type with rooturl
                        sets = recs[3]
                        rooturl = recs[2]
                    else:
                        print 'Recs wrong len %d' % (len(recs))
                        sys.exit(1)
                    FR = FileRec(recs[0], [])
                    self.grabs.append((FR, recs[1], rooturl))
                    for r in sets:
                        set = []
                        for rr in r:
                            try:
                                file = files[rr[2]] # file already open?
                            except KeyError:
                                try:
                                    file = np_file(rr[2], 0)
                                except IOError, s:
                                    print 'Can\'t open %s: %s' % (rr[2], s)
                                    sys.exit(1)
                                files[rr[2]] = file
                            set.append((rr[0], rr[1], file, rr[3]))
                        FR.add(set)

                    # reduce to single level list
                    if recs[0] == 'TCPConn':
                        FR.recsets = FR.recsets[0]
                                    
            except EOFError:
                break

    def show_files(self):

        files = {}
        #self.getfile('read')
        
        while 1:
            try:
                read = load(self.gfile)
                for recs in read:
                    if len(recs) == 3: # old type - no rooturl
                        sets = recs[2]
                    elif len(recs) == 4: # new type with rooturl
                        sets = recs[3]
                    else:
                        print 'Recs wrong len %d' % (len(recs))
                        sys.exit(1)
                    for r in sets: