dirwalk.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.                                                                   #
#                                                                             #
###############################################################################


##############################################################################
## 
## 
## Walk a directory tree (pre-dates walk of later Pyton versions - with enhanced## functionality 
## 
##
## 
##
############################################################################ 
############################################################################ 

from __future__ import generators

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

from string import *
from sys import argv
import getopt
import os
import sys
from stat import *
from os import listdir
from os.path import normpath, basename, dirname, join, isfile, isdir, islink, \
realpath, expanduser

from sys import exit
import popen2
import re

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

def resolve_path(path):

    # patch any leading relative path element
    firstel = path.split('/')[0]
    if firstel  == '.' or firstel == '..' or path[0] not in '/~':
        path = join(os.getcwd(), path)

    # normalise any relative or user elements
    return normpath(expanduser(path))

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

def walk(root, topdown=1, onerror=None, depth=0, maxdepth=None,
         exlist=None, followlinks=0, visited=None, excluded=[]):

    def exclude_it(dir, res):

        if res and type(res[0]) == str:
            # first call
            for i in range(len(res)):
                try:
                    res[i] = re.compile(res[i]+'$')
                except:
                    print 'Can\'t compile RE \'%s\'' % (res[i])
                    sys.exit(1)
        
        for RE in res:
            if RE.search(dir):
                return 1

        return 0

    dirs = []
    files = []

    #print exlist

    if exlist != None and not exlist:
        exlist.extend([[], [], []])

    root = resolve_path(root)

    if followlinks:
        if visited == None:
            visited = {realpath(root): 1}
        else:
            visited[realpath(root)] = 1
            
    try:
        names = listdir(root)
    except OSError, s:
        if exlist:
            exlist[0].append(str(s))
        if onerror != None:
            onerror(s)
            yield(root, [], [])
            return
        else:
            raise s

    for nm in names:
        path = join(root, nm)
        if exclude_it(nm, excluded):
            if exlist:
                exlist[2].append(root)
            continue
        if isdir(path):
            dirs.append(nm)
        elif isfile(path):
            files.append(nm)
        else:
            if exlist:
                exlist[0].append('\'%s\' is not a file or directory' % (path))
            #print '\'%s\' is not a file or directory' % (path)


    dirs.sort()
    files.sort()

    if topdown:
        
        yield(root, dirs, files)
        
        if maxdepth == None or depth < maxdepth:
            for d in dirs:
                dir = join(root, d)
                if followlinks and visited.has_key(realpath(dir)):
                   ##  print '%s has already been visited as %s' % (dir,
##                                                             realpath(dir))
##                     raw_input('<RET> to continue')
                    if exlist:
                        exlist[1].append((dir, realpath(dir)))
                    continue
                if (not islink(dir)) or followlinks:
                    for t in walk(dir, topdown=topdown, onerror=onerror,
                                  depth=depth+1, maxdepth=maxdepth,
                                  exlist=exlist, followlinks=followlinks,
                                  visited=visited, excluded=excluded):
                        yield t
    else:
        
        if maxdepth == None or depth < maxdepth:
            for d in dirs:
                dir = join(root, d)
                if followlinks and visited.has_key(realpath(dir)):
                    ##print '%s has already been visited as %s' % (dir,
##                                                               realpath(dir))
##                     raw_input('<RET> to continue')
                    if exlist:
                        exlist[1].append((dir, realpath(dir)))
                    continue
                if (not islink(dir)) or followlinks:
                    for t in walk(dir, topdown=topdown, onerror=onerror,
                                  depth=depth+1, maxdepth=maxdepth,
                                  exlist=exlist, followlinks=followlinks,
                                  visited=visited, excluded=excluded):
                        yield t
        yield(root, dirs, files)
        

        
###############################################################################
###############################################################################
 
def ignore_OSError(s):
    pass

###############################################################################
 
def print_OSError(s):
    
    print '%s caught and ignored' % (str(s))

###############################################################################
 
def get_excluded(el):

    elist = el.split(',')
    rel = []
    for e in elist:
        rel.append(e)

    return rel

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


def main():

    topdown = 1
    maxdepth = DEFAULT_MAX_DEPTH
    exlist = []
    followlinks = 0
    excluded_re = []
    
    try:
        optlist, args = getopt.getopt(argv[1:], 'bd:le:')

    except getopt.error, s:
        print 'dirwalk test:', str(s)
        sys.exit(1)

    for opt in optlist:
        if opt[0] == '-b':
            topdown = 0
        if opt[0] == '-d':
            maxdepth = int(opt[1])
        if opt[0] == '-l':
            followlinks=1
        if opt[0] == '-e':
            excluded_re = get_excluded(opt[1])

    if len(args) ==0:
        dir = os.getcwd()
    elif len(args) ==1:
        dir = args[0]
    else:
        print 'dirwalk test - wrong number of args'
        sys.exit(1)
        
  ##   for root, dirs, files in walk(dir, topdown=topdown, onerror=print_OSError,
##                                   maxdepth=maxdepth, exlist=exlist):
        
    for root, dirs, files in walk(dir, topdown=topdown, onerror=print_OSError,
                                  maxdepth=maxdepth, followlinks=followlinks,
                                  exlist=exlist, excluded=excluded_re):
        #print root, dirs, files
        if files:
            print root
            for f in files:
                print f
            print

    if exlist:
        if exlist[0]:
            print '\nXXX The following errors occurred:'
            for err in exlist[0]:
                print err
        if exlist[1]:
            print '\nXXX The following directories were not visited multiple times:'
            for d in exlist[1]:
                print '\'%s\' -> \'%s\'' % (d[0], d[1])
        if exlist[2]:
            print '\nXXX The following dirs were excluded:'
            for d in exlist[2]:
                print d
            
             
##############################################################################


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