<s> '''<STR_LIT>''' <EOL> import os <EOL> import sys <EOL> import logging <EOL> import struct <EOL> import datetime <EOL> import dshell <EOL> import util <EOL> class Output ( object ) : <EOL> '''<STR_LIT>''' <EOL> _DEFAULT_FORMAT = '<STR_LIT>' <EOL> _DEFAULT_TIMEFORMAT = '<STR_LIT>' <EOL> _DEFAULT_DELIM = '<STR_LIT:U+0020>' <EOL> _NULL = None <EOL> _FILTER_EXTRA = False <EOL> def __init__ ( self , * a , ** kw ) : <EOL> '''<STR_LIT>''' <EOL> self . logger = kw . get ( '<STR_LIT>' , logging ) <EOL> self . setformat ( kw . get ( '<STR_LIT>' , self . _DEFAULT_FORMAT ) ) <EOL> self . timeformat = ( kw . get ( '<STR_LIT>' , self . _DEFAULT_TIMEFORMAT ) ) <EOL> self . delim = ( kw . get ( '<STR_LIT>' , self . _DEFAULT_DELIM ) ) <EOL> self . nobuffer = ( kw . get ( '<STR_LIT>' , False ) ) <EOL> if '<STR_LIT>' in kw : <EOL> self . pcapwriter = PCAPWriter ( kw [ '<STR_LIT>' ] ) <EOL> else : <EOL> self . pcapwriter = None <EOL> self . extra = kw . get ( '<STR_LIT>' , False ) <EOL> if '<STR_LIT>' in kw : <EOL> self . sessionwriter = SessionWriter ( ** kw ) <EOL> else : <EOL> self . sessionwriter = None <EOL> def log ( self , msg , level = logging . INFO , * args , ** kw ) : <EOL> '''<STR_LIT>''' <EOL> self . logger . log ( level , msg , * args , ** kw ) <EOL> def setformat ( self , formatstr = None , typemap = None ) : <EOL> '''<STR_LIT>''' <EOL> if formatstr : <EOL> self . format = formatstr + "<STR_LIT:\n>" <EOL> else : <EOL> self . format = self . _DEFAULT_FORMAT + "<STR_LIT:\n>" <EOL> self . fields = [ ] <EOL> self . fieldnames = [ ] <EOL> self . fieldmap = { } <EOL> e = <NUM_LIT:0> <EOL> while True : <EOL> s = self . format . find ( '<STR_LIT:%>' , e ) + <NUM_LIT:1> <EOL> if s < <NUM_LIT:1> or self . format [ s ] != '<STR_LIT:(>' : <EOL> break <EOL> e = self . format . find ( '<STR_LIT:)>' , s ) <EOL> if e < <NUM_LIT:0> : <EOL> break <EOL> fname = self . format [ s + <NUM_LIT:1> : e ] <EOL> fspec = '<STR_LIT>' <EOL> for i in xrange ( e + <NUM_LIT:1> , len ( self . format ) ) : <EOL> if self . format [ i ] in '<STR_LIT>' : <EOL> fspec += self . format [ i ] <EOL> else : <EOL> break <EOL> ftype = self . format [ i ] <EOL> i += <NUM_LIT:1> <EOL> if i < len ( self . format ) and self . format [ i ] == '<STR_LIT::>' : <EOL> e = self . format . find ( '<STR_LIT:U+0020>' , i ) <EOL> ftype , outtype = self . format [ i - <NUM_LIT:1> : e ] . split ( '<STR_LIT::>' ) <EOL> else : <EOL> outtype = None <EOL> e = i <EOL> try : <EOL> fname , fmap = fname . split ( '<STR_LIT::>' ) <EOL> except : <EOL> fmap = fname <EOL> if typemap and ftype in typemap and not outtype : <EOL> try : <EOL> ( ftype , outtype ) = typemap [ ftype ] <EOL> except : <EOL> outtype = typemap [ ftype ] <EOL> self . fields . append ( ( fname , ftype , fspec ) ) <EOL> self . fieldnames . append ( fname ) <EOL> if outtype : <EOL> self . fieldmap [ fname ] = ( fmap , outtype ) <EOL> def parse ( self , * args , ** kw ) : <EOL> '''<STR_LIT>''' <EOL> for ts in [ k for k in kw if k == '<STR_LIT>' or k . endswith ( '<STR_LIT:time>' ) ] : <EOL> dt = ts [ : - <NUM_LIT:4> ] + '<STR_LIT>' <EOL> kw [ dt ] = datetime . datetime . fromtimestamp ( <EOL> float ( kw [ ts ] ) ) . strftime ( self . timeformat ) <EOL> if kw . get ( '<STR_LIT>' ) is '<STR_LIT>' : <EOL> kw [ '<STR_LIT>' ] = '<STR_LIT>' <EOL> elif kw . get ( '<STR_LIT>' ) is '<STR_LIT>' : <EOL> kw [ '<STR_LIT>' ] = '<STR_LIT>' <EOL> else : <EOL> kw [ '<STR_LIT>' ] = '<STR_LIT>' <EOL> if '<STR_LIT>' in kw : <EOL> self . setformat ( kw [ '<STR_LIT>' ] ) <EOL> del kw [ '<STR_LIT>' ] <EOL> rec = dict ( ( f , self . _NULL ) for f in self . fieldnames ) <EOL> if self . _FILTER_EXTRA : <EOL> rec . update ( <EOL> dict ( ( f , kw [ f ] ) for f in self . fieldnames if ( f in kw and kw [ f ] != None ) ) ) <EOL> if self . extra : <EOL> rec [ '<STR_LIT>' ] = self . delim . join ( [ '<STR_LIT>' % ( f , kw [ f ] ) for f in sorted ( <EOL> kw . keys ( ) ) if f not in self . fieldnames and f != '<STR_LIT>' ] ) <EOL> else : <EOL> rec . update ( kw ) <EOL> if args : <EOL> rec [ '<STR_LIT:data>' ] = self . delim . join ( map ( str , args ) ) <EOL> return rec <EOL> def dump ( self , pkt = None , ** kw ) : <EOL> '''<STR_LIT>''' <EOL> pktdata = str ( pkt ) <EOL> pktlen = kw . get ( '<STR_LIT>' , len ( pktdata ) ) <EOL> if self . pcapwriter : <EOL> self . pcapwriter . write ( pktlen , pktdata , kw [ '<STR_LIT>' ] ) <EOL> else : <EOL> self . log ( util . hexPlusAscii ( str ( pkt ) ) , level = logging . DEBUG ) <EOL> def close ( self ) : <EOL> if self . pcapwriter : <EOL> self . pcapwriter . close ( ) <EOL> def dispatch ( self , m , * args , ** kwargs ) : <EOL> '''<STR_LIT>''' <EOL> if m == '<STR_LIT>' : <EOL> self . write ( * args , ** kwargs ) <EOL> if m == '<STR_LIT>' : <EOL> self . alert ( * args , ** kwargs ) <EOL> if m == '<STR_LIT>' : <EOL> self . dump ( * args , ** kwargs ) <EOL> def setup ( self ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> class FileOutput ( Output ) : <EOL> def __init__ ( self , * args , ** kw ) : <EOL> '''<STR_LIT>''' <EOL> Output . __init__ ( self , * args , ** kw ) <EOL> f = None <EOL> if '<STR_LIT>' in kw : <EOL> self . fh = kw [ '<STR_LIT>' ] <EOL> return <EOL> elif '<STR_LIT:file>' in kw : <EOL> f = kw [ '<STR_LIT:file>' ] <EOL> elif args : <EOL> f = args [ <NUM_LIT:0> ] <EOL> if f : <EOL> if '<STR_LIT>' in kw : <EOL> mode = kw [ '<STR_LIT>' ] <EOL> else : <EOL> mode = '<STR_LIT:w>' <EOL> if mode == '<STR_LIT>' : <EOL> mode = '<STR_LIT:w>' <EOL> try : <EOL> while os . stat ( f ) : <EOL> p = f . split ( '<STR_LIT:->' ) <EOL> try : <EOL> p , n = p [ : - <NUM_LIT:1> ] , int ( p [ - <NUM_LIT:1> ] ) <EOL> except ValueError : <EOL> n = <NUM_LIT:0> <EOL> f = '<STR_LIT:->' . join ( p + [ '<STR_LIT>' % ( int ( n ) + <NUM_LIT:1> ) ] ) <EOL> except OSError : <EOL> pass <EOL> self . fh = open ( f , mode ) <EOL> else : <EOL> self . fh = sys . stdout <EOL> def write ( self , obj , ** kw ) : <EOL> '''<STR_LIT>''' <EOL> if self . sessionwriter : <EOL> self . sessionwriter . write ( obj , ** kw ) <EOL> elif self . fh : <EOL> self . fh . write ( str ( obj ) ) <EOL> if self . nobuffer : <EOL> self . fh . flush ( ) <EOL> def close ( self ) : <EOL> '''<STR_LIT>''' <EOL> if self . fh != sys . stdout : <EOL> self . fh . close ( ) <EOL> Output . close ( self ) <EOL> class TextOutput ( FileOutput ) : <EOL> '''<STR_LIT>''' <EOL> _DEFAULT_FORMAT = "<STR_LIT>" <EOL> _NULL = '<STR_LIT>' <EOL> _FILTER_EXTRA = True <EOL> def __init__ ( self , * args , ** kw ) : <EOL> if '<STR_LIT>' in kw : <EOL> self . _DEFAULT_FORMAT += "<STR_LIT>" <EOL> FileOutput . __init__ ( self , * args , ** kw ) <EOL> def alert ( self , * args , ** kw ) : <EOL> '''<STR_LIT>''' <EOL> rec = self . parse ( * args , ** kw ) <EOL> if rec : <EOL> self . fh . write ( self . format % rec ) <EOL> if self . nobuffer : <EOL> self . fh . flush ( ) <EOL> class DBOutput ( Output ) : <EOL> '''<STR_LIT>''' <EOL> _DEFAULT_FORMAT = "<STR_LIT>" <EOL> _NULL = None <EOL> _TYPEMAP = { '<STR_LIT:s>' : '<STR_LIT>' , '<STR_LIT:r>' : '<STR_LIT>' , '<STR_LIT:c>' : '<STR_LIT>' , <EOL> '<STR_LIT:x>' : '<STR_LIT>' , '<STR_LIT:X>' : '<STR_LIT>' , '<STR_LIT:o>' : '<STR_LIT>' , <EOL> '<STR_LIT:d>' : '<STR_LIT>' , '<STR_LIT:i>' : '<STR_LIT>' , '<STR_LIT:u>' : '<STR_LIT>' , <EOL> '<STR_LIT:e>' : '<STR_LIT>' , '<STR_LIT:E>' : '<STR_LIT>' , <EOL> '<STR_LIT:f>' : '<STR_LIT>' , '<STR_LIT:F>' : '<STR_LIT>' , <EOL> '<STR_LIT:g>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT:b>' : ( '<STR_LIT:d>' , '<STR_LIT>' ) , <EOL> '<STR_LIT:t>' : ( '<STR_LIT:f>' , '<STR_LIT>' ) , '<STR_LIT:D>' : ( '<STR_LIT:s>' , '<STR_LIT>' ) , <EOL> '<STR_LIT:A>' : ( '<STR_LIT:s>' , '<STR_LIT>' ) , '<STR_LIT:H>' : ( '<STR_LIT:s>' , '<STR_LIT>' ) , '<STR_LIT:N>' : ( '<STR_LIT:s>' , '<STR_LIT>' ) , '<STR_LIT:M>' : ( '<STR_LIT:s>' , '<STR_LIT>' ) } <EOL> _DBCONNPARAMS = [ '<STR_LIT:host>' , '<STR_LIT:user>' , '<STR_LIT>' , <EOL> '<STR_LIT:password>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:port>' , '<STR_LIT>' ] <EOL> _DBTYPE_PLACEHOLDER_MAP = { '<STR_LIT>' : '<STR_LIT:?>' } <EOL> def __init__ ( self , * args , ** kw ) : <EOL> '''<STR_LIT>''' <EOL> self . dbconfig = kw . copy ( ) <EOL> if '<STR_LIT>' in self . dbconfig : <EOL> import ConfigParser <EOL> config = ConfigParser . ConfigParser ( ) <EOL> config . read ( self . dbconfig [ '<STR_LIT>' ] ) <EOL> sections = config . sections ( ) <EOL> if len ( sections ) > <NUM_LIT:0> : <EOL> self . dbconfig [ '<STR_LIT>' ] = sections [ <NUM_LIT:0> ] <EOL> for k , v in config . items ( sections [ <NUM_LIT:0> ] , raw = True ) : <EOL> self . dbconfig [ k ] = v <EOL> self . db = __import__ ( self . dbconfig [ '<STR_LIT>' ] ) <EOL> self . dbconn = self . db . connect ( <EOL> * args , ** dict ( ( k , self . dbconfig [ k ] ) for k in self . _DBCONNPARAMS if k in self . dbconfig ) ) <EOL> Output . __init__ ( self , * args , ** self . dbconfig ) <EOL> def createtable ( self , table = None ) : <EOL> '''<STR_LIT>''' <EOL> if not table and '<STR_LIT>' in self . dbconfig : <EOL> table = self . dbconfig [ '<STR_LIT>' ] <EOL> try : <EOL> cursor = self . dbconn . cursor ( ) <EOL> sqlfields = [ ] <EOL> for fname , ftype , fspec in [ f for f in self . fields if f [ <NUM_LIT:1> ] != '<STR_LIT:T>' ] : <EOL> ctype = self . fieldmap [ fname ] [ <NUM_LIT:1> ] <EOL> if ctype == '<STR_LIT>' and not fspec : <EOL> ctype = '<STR_LIT>' <EOL> fdef = self . fieldmap [ fname ] [ <NUM_LIT:0> ] + '<STR_LIT:U+0020>' + ctype <EOL> if fspec : <EOL> fdef += '<STR_LIT:(>' + fspec . strip ( '<STR_LIT>' ) . replace ( '<STR_LIT:.>' , '<STR_LIT:U+002C>' ) + '<STR_LIT:)>' <EOL> sqlfields . append ( fdef ) <EOL> sql = '<STR_LIT>' + table + '<STR_LIT>' + '<STR_LIT:U+002C>' . join ( sqlfields ) + '<STR_LIT:)>' <EOL> self . log ( sql , logging . DEBUG ) <EOL> return cursor . execute ( sql ) <EOL> except : <EOL> raise <EOL> def close ( self ) : <EOL> '''<STR_LIT>''' <EOL> self . dbconn . close ( ) <EOL> Output . close ( self ) <EOL> def alert ( self , * args , ** kw ) : <EOL> '''<STR_LIT>''' <EOL> rec = self . parse ( self , * args , ** kw ) <EOL> if rec : <EOL> self . insert ( rec ) <EOL> def setformat ( self , formatstr = None ) : <EOL> '''<STR_LIT>''' <EOL> ph = self . dbconfig . get ( '<STR_LIT>' , <EOL> self . _DBTYPE_PLACEHOLDER_MAP . get ( <EOL> self . dbconfig [ '<STR_LIT>' ] , '<STR_LIT>' ) <EOL> ) <EOL> self . format_types = self . dbconfig . get ( '<STR_LIT>' , '<STR_LIT:x>' ) <EOL> Output . setformat ( self , formatstr , typemap = self . _TYPEMAP ) <EOL> self . tablefield = '<STR_LIT>' <EOL> for fname , ftype , fspec in self . fields : <EOL> if ftype == '<STR_LIT:T>' : <EOL> self . tablefield = fname <EOL> sqlfields = [ self . fieldmap [ fname ] [ <NUM_LIT:0> ] <EOL> for ( fname , ftype , fspec ) in self . fields if fname in self . fieldmap ] <EOL> self . insertsql = '<STR_LIT>' % ( <EOL> '<STR_LIT:U+002C>' . join ( sqlfields ) , '<STR_LIT:U+002C>' . join ( [ ph ] * len ( sqlfields ) ) ) <EOL> def insert ( self , rec , table = None ) : <EOL> '''<STR_LIT>''' <EOL> if not table : <EOL> if '<STR_LIT>' in self . dbconfig : <EOL> table = self . dbconfig [ '<STR_LIT>' ] <EOL> elif rec [ self . tablefield ] : <EOL> table = rec [ self . tablefield ] <EOL> try : <EOL> sqlvalues = [ ] <EOL> cursor = self . dbconn . cursor ( ) <EOL> for fname , ftype , fspec in self . fields : <EOL> if fname in self . fieldmap : <EOL> if ftype in self . format_types : <EOL> sqlvalues . append ( ( '<STR_LIT:%>' + fspec + ftype ) % rec [ fname ] ) <EOL> else : <EOL> sqlvalues . append ( rec [ fname ] ) <EOL> sql = self . insertsql % table <EOL> self . log ( sql + '<STR_LIT>' % sqlvalues , logging . DEBUG ) <EOL> except : <EOL> raise <EOL> fail = False <EOL> while True : <EOL> try : <EOL> cursor . execute ( sql , sqlvalues ) <EOL> self . dbconn . commit ( ) <EOL> break <EOL> except Exception , e : <EOL> self . log ( e , level = logging . WARNING ) <EOL> if fail : <EOL> raise <EOL> else : <EOL> fail = True <EOL> try : <EOL> self . createtable ( table ) <EOL> except : <EOL> raise <EOL> class PCAPWriter ( FileOutput ) : <EOL> '''<STR_LIT>''' <EOL> def __init__ ( self , * args , ** kw ) : <EOL> FileOutput . __init__ ( self , * args , ** kw ) <EOL> if self . fh : <EOL> self . fh . write ( <EOL> struct . pack ( '<STR_LIT>' , <NUM_LIT> , <NUM_LIT:2> , <NUM_LIT:4> , <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT> , <NUM_LIT:1> ) ) <EOL> def write ( self , pktlen , pktdata , ts ) : <EOL> if self . fh : <EOL> self . fh . write ( <EOL> struct . pack ( '<STR_LIT>' , int ( ts ) , int ( ( ts - int ( ts ) ) * <NUM_LIT> ) ) ) <EOL> self . fh . write ( struct . pack ( '<STR_LIT>' , len ( pktdata ) , pktlen ) ) <EOL> self . fh . write ( pktdata ) <EOL> class SessionWriter ( Output ) : <EOL> '''<STR_LIT>''' <EOL> def __init__ ( self , session = None , ** kw ) : <EOL> self . file = kw . get ( '<STR_LIT>' , session ) <EOL> self . dir = kw . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> self . mode = kw . get ( '<STR_LIT>' , '<STR_LIT:a>' ) <EOL> self . timeformat = ( kw . get ( '<STR_LIT>' , self . _DEFAULT_TIMEFORMAT ) ) <EOL> self . fieldnames = [ ] <EOL> def write ( self , obj , ** kwargs ) : <EOL> out = None <EOL> kw = dict ( ** kwargs ) <EOL> try : <EOL> kw . update ( ** obj . info ( ) ) <EOL> kw = self . parse ( ** kw ) <EOL> if self . dir == '<STR_LIT>' : <EOL> ds = [ None ] <EOL> elif self . dir == '<STR_LIT>' : <EOL> ds = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> else : <EOL> ds = [ self . dir ] <EOL> for d in ds : <EOL> kw . update ( direction = d if d else '<STR_LIT>' ) <EOL> out = FileOutput ( self . file % kw , mode = self . mode ) <EOL> out . fh . write ( obj . data ( direction = d ) ) <EOL> out . close ( ) <EOL> except : <EOL> out = FileOutput ( self . file % kw , mode = self . mode ) <EOL> out . fh . write ( str ( obj ) ) <EOL> out . close ( ) <EOL> class QueueOutput ( Output ) : <EOL> '''<STR_LIT>''' <EOL> def __init__ ( self , q , ** kwargs ) : <EOL> self . queue = q <EOL> Output . __init__ ( self , ** kwargs ) <EOL> def write ( self , * args , ** kw ) : self . dispatch ( '<STR_LIT>' , * args , ** kw ) <EOL> def alert ( self , * args , ** kw ) : self . dispatch ( '<STR_LIT>' , * args , ** kw ) <EOL> def dump ( self , * args , ** kw ) : self . dispatch ( '<STR_LIT>' , * args , ** kw ) <EOL> def dispatch ( self , m , * args , ** kw ) : <EOL> self . queue . put ( ( m , args , kw ) ) <EOL> def close ( self ) : <EOL> self . queue . close ( ) <EOL> Output . close ( self ) <EOL> obj = TextOutput </s>
