np_namecache.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 os
import sys
from sys import argv
import getopt
from signal import *

from nprobe import set_print_full_hostname, intoa_string, host_string, tcp_open
from nprobe import *
import np_file_util
import string
import os.path

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

CACHEFILE_DIR_DEF = '/usr/groups/nprobe/jch1003/ns'

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


class NameCache:
    """
    
    """
    def __init__(self, cachepath, load=0):
        """
        __init__(self)
        Initialize a new instance.
        """

        self.lookups_dict = {}
        if cachepath:
            self.cachefile_dir = cachepath
        else:
            self.cachefile_dir = CACHEFILE_DIR_DEF
        self.make_cachefile_name()
        self.cache_changed = 0
        self.use_long_hostnames = 0

        if load:
            self.load()
        
        return

    def make_cachefile_name(self):
    

        dir = os.path.dirname(self.cachefile_dir)

        if dir and dir[-1] != '/':
            dir += '/'

        self.cachefile_name =  dir + 'lookups.cache'
        

    def use_hostnames(self):

        self.use_long_hostnames = 1
        set_print_full_hostname(1)


    def addr2bin(self, addr):
        if type(addr) == type(0):
            return addr
        bytes = string.splitfields(addr, '.')
        if len(bytes) != 4:
            raise ValueError, 'bad IP address'
        n = 0
        for byte in bytes:
            n = n<<8 | int(byte)
        return n

    def addr2bin_nbo(self, addr):
        if type(addr) == type(0):
            return addr
        bytes = string.splitfields(addr, '.')
        if len(bytes) != 4:
            raise ValueError, 'bad IP address'
        n = 0
        bytes.reverse()
        for byte in bytes:
            n = n<<8 | int(byte)
        return n

    def bin2addr(self, n):
        return '%d.%d.%d.%d' % ((n>>24)&0xFF, (n>>16)&0xFF,
                                (n>>8)&0xFF, n&0xFF)

    def bin2addr_nbo(self, n):
        return '%d.%d.%d.%d' % (n&0xFF, (n>>8)&0xFF, (n>>16)&0xFF,
                                (n>>24)&0xFF)

    def is_dd(self, a):
        as = a.split('.')
        if len(as) != 4:
            return 0
        else:
            for d in as:
                try:
                    dd = int(d)
                except:
                    return 0
        return 1


    def hostname(self, addr, verb=0, retry=0):
        
        # print 'LOOKING UP %s ' % (intoa_string(addr)),
        lookups_dict = self.lookups_dict
        if self.use_long_hostnames:
            if lookups_dict.has_key(addr) and not retry:
                # print 'FOUND %s' % (lookups_dict[addr])
                return lookups_dict[addr]
            else:
                if verb:
                    print '%s -> ' % (intoa_string(addr)),
                    sys.stdout.flush()
                addrstr = host_string(addr)
                #print 'GOT %s' % (addrstr)			
                lookups_dict[addr] = addrstr
                self.cache_changed = 1
                if verb:
                    print '%s' % (addrstr)
            return addrstr
        else:
            return intoa_string(addr)

    def add2cache(self, addr, name):
        if not self.lookups_dict.has_key(addr):
            # print 'ADDED %s %s' % (intoa_string(addr), name)
            self.lookups_dict[addr] = name
            self.cache_changed = 1

    def file_cache(self):
        # print 'cach_lookups()'
        if self.cache_changed:
            # print 'Saving cached lookups'
            bin2addr_nbo = self.bin2addr_nbo
            f = open(self.cachefile_name, 'w')
            for entry in self.lookups_dict.items():
                f.write('%s %s\n' % (bin2addr_nbo(entry[0]), entry[1]))

    def load(self):

        print 'load'
        try:
            f = open(self.cachefile_name, 'r')
            print 'Using ns cachefile %s' % (self.cachefile_name)
        except:
            print 'No name cachefile %s - will create ...' \
                  % (self.cachefile_name)
            return
        lines = f.readlines()
        lookups_dict = self.lookups_dict
        addr2bin_nbo = self.addr2bin_nbo
        for line in lines:
            entry = string.split(line)
            # print entry
            lookups_dict[addr2bin_nbo(entry[0])] = entry[1]
            # print '%s = %s' % (entry[0], entry[1])
            # print
            # for e in lookups_dict.items():
            # print '%s -> %s' % (intoa_string(e[0]), e[1])


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

def usage(scriptname):
    print "usage: " + scriptname + " rep-file-dir"

    sys.exit(1)

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

def handle_sigint(n, f):


    ans = raw_input('SIGINT - save lookups? (y/n)')
    if ans.count('y'):
        nc.file_cache()
    sys.exit(0)

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

nc = None
	
def main():

    global nc

    scriptname =  os.path.basename(argv[0])
    cachepath = None
    verb = 0
    retry = 0
    servers = 0

    try:
	optlist, args = getopt.getopt(sys.argv[1:], 'vf:rS')

    except getopt.error, s:
        print '%s: %s' % (scriptname, s)
	usage(scriptname)
	sys.exit(1)

    for opt in optlist:
        if opt[0] == "-h":
            usage(scriptname)
        if opt[0] == "-f":
            cachepath = opt[1]
        if opt[0] == "-v":
            verb = 1
        if opt[0] == "-r":
            retry = 1
        if opt[0] == "-S":
            servers = 1
                

    signal(SIGINT, handle_sigint)
    openfilelist, counters, basepath = np_file_util.get_files(args)
    
    if cachepath == '.' or cachepath == None:
        cachepath = basepath
    # get_cached_lookups(cachedir)

    nc = NameCache(cachepath, load=1)
    #nc.load()

    nc.use_hostnames()
    hostname = nc.hostname
    dd = nc.is_dd
    
    #
    # To get TCP open records into
    #
    tcp_open_rec = tcp_open()

    for file in openfilelist:

        if len(openfilelist) > 1:
            print "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
            print "File %s - %d records" % (file.fnm, file.counters.nrecords)
            print
            sys.stdout.flush()

        while 1:
            rectype = file.next_rec()
            #print rectype
            if rectype == -1: #EOF
                break

            elif rectype == REC_TCP_HTTP_OPEN: # notification of TCP open
                tcp_open_rec.get_open_rec(file)
                if not servers:
                    caddr = tcp_open_rec.shost()
                    hostname(caddr, verb=verb)
                saddr = tcp_open_rec.dhost()
                hostname(saddr, verb=verb)

                for addr in [tcp_open_rec.shost(), tcp_open_rec.dhost()]:
                    nm = hostname(addr, verb=verb)
                    if retry and dd(nm):
                        print 'got %s - retrying' % (nm)
                        nm = hostname(addr, verb=verb, retry=1)
                    elif verb:
                        print 'got', nm
            else:
                file.advance()
                continue

    nc.file_cache()
	    
##############################################################################


# Call main when run as script
if __name__ == '__main__':
        main()