#! /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 glob
import os
import sys
import nprobe
from sys import argv
import getopt
from signal import *
import atexit
import gc
from traceback import print_stack
from Tkinter import *

import Numeric

from nprobe import REC_TCP_HTTP, REC_TCP_HTTP_OPEN, REC_TCP_HTTP_HDRS, \
     OPEN_BIT, HDRS_BIT
from nprobe import SERVER, CLIENT, tcp_open, tcp_hdrs
from np_file_util import get_files
from np_http_util import get_http_rec
from np_http_util import allocate_http_reusable_objects, get_http_rec_and_trans
import np_ns_utils
from np_lookups import get_cached_lookups
from np_TCPConn import TCPConn, TCPHdr
from np_TCPDisplay import TcpDisplay
from  np_longutil import ull2l, Longstring, tv2l
#from np_lookups import *
from np_tfilter import TFilter
#from np_TCP import TCP_Machine, TCPNoModel, TCPModelPkts, TCPModelNoTrans
from TCP_Imp import TCP_Imps
from np_atypes import Atype
from np_filerec import FileRec
from np_tcpstats import TCPStats
from minmax import BIGNUMBER, MIN, MAX, MAX3, MAX4

#from np_notesel import Sellist
from np_ectrs import E_Counters

from print_col import whoops, overprint


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

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

    sys.exit(1)

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

def exitfun():
    
    from np_grabrec import grabobj
    
    #print 'exitfun called'
    #cache_lookups()
    print grabobj
    if grabobj:
        grabobj.close()



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

def handle_sigint(n, f):

    global logf

    print 'SIGINT'
    if logf != None:
	logf.flush()
    sys.exit(0)

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

#
# Save logged events to file
#

log = []
logpath = None
logf = None

def save_log():

    global logf, logpath

    

    if logf == None:
	logf = open(logpath, 'w')

    for s in log:
        #if not ':' in s:
            #print s
        logf.write(s+'\n')

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

#
# Log 'difficult' events noted
#

def write_log(txt):

    global log

    log.append(txt[:])

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

#
# Null logfun
#

def logfun_null(txt):

    pass

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

#
# Aggregate cl list of connections wanted

def get_conns_wanted(s):
    connids = []
    for f in string.splitfields(s, ','):
	#print f
	#print '%d' % (string.atoi(f))
	connids.append(string.atoi(f))

    return connids

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

#
# Call-back for Sellist if desired to examine object(s)
#   - reconstructs from file records and displays

def drawfun(what, oblist, obdict):

##     keys =  obdict.keys()
##     keys.sort()
    print what, oblist

    if what == 'Connections':
        oblist.sort()
        clist = []
        for m in oblist:
            type = m[2]
            if type == 'TCPConn':
                #conn = obdict[m[0]][0].reconstruct(trace=1)
                conn = obdict[m[0]].reconstruct(trace=1)
                conn.adjust_tm_offsets(conn.abstart)
                clist.append(conn)
            elif type == 'Rank Server':
                for c in obdict[m[0]]:
                    conn = obdict[c].reconstruct(trace=1)
                    conn.adjust_tm_offsets(conn.abstart)
                    clist.append(conn)
            else:
                whoops('Can\'t draw TCP Connection type %s\n' % (type))
                sys.exit (0)
                
        t = TcpDisplay(clist, standalone='no', trace=1)
        del(t.display)
        del(t)
    else:
        stats = obdict['stats']
        stats.drawfun(what, oblist, obdict)
        
        

#############################################################################
	
def main():

    global logpath, logf

    scriptname = os.path.basename(argv[0])
    do_draw = 0
    nrecs = 0
    startconn = 0
    oneconn = None
    conn_ids_got = 0
    connstr = ''
    fspec = 0
    fstr = ''
    tstr = ''
    trace = 0
    quiet = 0
    tfilt = TFilter()


    try:
	optlist, args = getopt.getopt(sys.argv[1:], 'hvdtr:F:T:s:c:q')

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

    optstr = ''

    for opt in optlist:
	if opt[0] == "-h":
	    usage(scriptname)
	if opt[0] == "-v":
	    nprobe.set_print_full_hostname(1)
	if opt[0] == "-d":
	    do_draw = 1
	if opt[0] == "-t":
	    trace = 1
        if opt[0] == "-q":
            quiet = 1
            optstr += '-q'
	if opt[0] == "-r":
	    nrecs = string.atoi(opt[1])
	    optstr = optstr + '-r' + opt[1]
	if opt[0] == "-s":
	    startconn = string.atoi(opt[1])
	    optstr = optstr + '-s' + opt[1]
	if opt[0] == "-c":
	    #print opt[1]
	    #sys.exit(0)
	    #oneconn = string.atoi(opt[1])
	    oneconn = get_conns_wanted(opt[1])
	    #connstr = '.#' + opt[1]
	    optstr = optstr + '-c' + opt[1]
	    print 'analysing only connection(s) ',
	    print oneconn
	if opt[0] == "-F":
	    if opt[1] == 'h' or opt[1] == 'H':
		nprobe.filter_help()
		sys.exit(0)
	    fspec = string.atoi(opt[1])
	    #fstr = '.F' + opt[1]
	    optstr = optstr + '-F' + opt[1] 
	if opt[0] == "-T":
	    if opt[1] == 'h' or opt[1] == 'H':
		tfilt.help()
		sys.exit(0)
	    #tspec = string.atoi(opt[1])
	    tfilt.build_filters(opt[1])
	    #tstr = '.T' + opt[1]
	    optstr = optstr + '-T' + opt[1]

    if optstr and optstr[0] == '-':
        optstr = optstr[1:]

    signal(SIGINT, handle_sigint)
    get_cached_lookups()

    atexit.register(exitfun)

    openfilelist, counters, basepath = get_files(args)
    run_start = tv2l(counters.run_start())
    

    #counters.printself("")
    filepath = basepath + 'counters'
    counters.printself_tofile(filepath, '')
    del counters

    basep = basepath + optstr
    if os.path.basename(basep):
        basep += '.'
    basepath = basep
    
    if do_draw:
	logpath = '/dev/null'
    else:
	logpath = (basepath + '_machlog')

    #
    # Stats collector
    #
    invoke_args = ''
    for a in optlist:
        for aa in a:
            invoke_args += '%s' % (aa)
        invoke_args += ' '
    for a in args:
        invoke_args += '%s ' % (a)
        
    stats = TCPStats(invoke_args, quiet=quiet, logpath=logpath)
    stats.tfilt = tfilt
    stats.basepath = basepath

    #
    # Dictionary of tcp connections keyed by conn_id 
    # - used to collect headers
    #
    conn_ids = stats.conn_ids = {}

    #
    # Dictionary of objects 
    # - used to collect stats etc
    #
    obdict = stats.obdict = {}

    #
    # List of TCP connections found
    #
    conn_list = []

    #
    # Dict of servers seen
    #
    serv_dict = stats.serv_dict = {}

    #
    # Dict of clients seen
    #
    cli_dict = stats.cli_dict = {}

    #
    # Dict of agent types seen
    #
    adict = stats.adict = Atype()

    #
    # Bad analysis failure
    #
    bad_fails = []
    
    ectrs = stats.ectrs = E_Counters()

    nconns = 0
    # rejected by id filter
    id_rej = 0
    # rejected by connection filter
    cf_rej = 0
    # one way only
    ow_rej = 0
    # rejected by tfilter
    tf_rej = 0

    tcp_imps = TCP_Imps()

    nread = 0
    done = 0

    if do_draw:
        logfun = logfun_null
    else:
        logfun = stats.write_log

    #logfun = None

    gc.set_debug(gc.DEBUG_UNCOLLECTABLE)
    #gc.set_debug(gc.DEBUG_STATS)
    #print gc.get_threshold()
    gc.set_threshold(100, 10, 10)
    #raw_input('...')
    
    #
    # To get TCP records into
    #
    tcp_open_rec = tcp_open()
    hdrs = tcp_hdrs(1)
    connrec, translist = allocate_http_reusable_objects()

    for file in openfilelist:

        flen = os.stat(file.fnm)[6]

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

	while not done:

	    if nrecs and nread > nrecs:
		break

	    rectype = file.next_type_rec(REC_TCP_HTTP)
            offset = file.curr_offset # get before read advances
            #print 'offset %d' % (offset)
            if offset >= flen:
                print 'ERROR bad offset %d - file size %d' % (offset, flen)
	    #print rectype
	    if rectype == -1: #EOF
		break

	    elif rectype & OPEN_BIT: # notification of TCP open
		tcp_open_rec.get_open_rec(file)
		id = tcp_open_rec.conn_id
		#print 'open %d' % (id)
                ectrs.nconns += 1
		if oneconn != None:
		    if oneconn.count(id):
			print 'got conn %d' % (id)
			conn_ids[id] = ([], []) #set up for list of headers
			#conn_ids_got += 1
		    else:
                        ectrs.id_rej += 1
			#print 'id_rej' 
		elif oneconn == None and id >= startconn:
		    if nprobe.accept_conn(tcp_open_rec.flow, fspec):
                        # set up for hdr and file/offset lists
			conn_ids[id] = ([], [])
		    else:
                        ectrs.cf_rej += 1
			#print 'cf_rej'
		else:
		    print 'REJECTED'
		
		continue

	    elif rectype & HDRS_BIT: # a block of hdrs
		#hdrs = nprobe.tcp_hdrs(0)
		hdrs.get_hdrs_p(file)
		conn_id = hdrs.conn_id
		if conn_ids.has_key(conn_id):
                    ent = conn_ids[conn_id]
                    ent[1].append((file.indx, REC_TCP_HTTP_HDRS, file, offset))
		    hlist = ent[0]
                    nhdr = len(hlist)-1
		    base_tm = ull2l(hdrs.get_hdrs_abstm())
                    for i in range(hdrs.nheld):
                        hlist.append(TCPHdr(hdrs, i, base_tm, nhdr))
                        nhdr += 1

		continue

	    elif rectype == REC_TCP_HTTP: #TCP/HTTP
                nread = nread + 1
		#read the lot
		#indx, connrec, translist =  get_http_rec(file)
                ntrans = get_http_rec_and_trans(file, connrec, translist)
                tlist = translist[:ntrans]

	    else: # any other TCP close record
		# get tconn record and its data
		connrec = nprobe.tcp_conn()
		connrec.get_conn(file)


	    # get the conn_id
	    conn_id = connrec.get_conn_id()

	    # only want where both directions seen
	    if not (connrec.server_seen() and connrec.client_seen()):
		if conn_ids.has_key(conn_id):
		    # remove dictionary item
		    del conn_ids[conn_id]
                    ectrs.ow_rej += 1
		    #print 'ow_rej'
		continue


	    if conn_ids.has_key(conn_id):
		#print 'got'
		if tfilt.pss(connrec, tlist): # one we're interested in


		    if oneconn != None and oneconn.count(id):
			print 'got conn %d' % (id)
			conn_ids_got += 1
			if conn_ids_got == len(oneconn):
			    print 'done'
			    done = 1

                    # find any headers already dumped
                    ent = conn_ids[conn_id]
                    ent[1].append((file.indx, REC_TCP_HTTP, file, offset)) 
		    hlist = ent[0]
		    #add any with the close record
		    base_tm = ull2l(connrec.get_hdrs_abstm())
                    nhdr = len(hlist)-1
                    for i in range(connrec.get_nhdrs_held()):
                        hlist.append(TCPHdr(connrec, i, base_tm, nhdr))
                        nhdr += 1
                    
                    # create Filerec for connection
                    fr = FileRec('TCPConn', ent[1])
                    #fr.add(ent[1])
                    obdict[conn_id] = fr
                    
		    # remove dictionary item
		    del conn_ids[conn_id]
		    #connrec.printself()

		    #conn.adjust_tm_offsets(conn.abstart)
                    #print 'start %.2f offset %.2f' % (conn.abstart/1000000.0, (conn.abstart-run_start)/1000000.0)

		    if do_draw:
			conn_list.append(fr)
			
		    else:
			#
			# run through the state machine
			#
                        conn = TCPConn(connrec, hlist, logfun, trace=trace)
                        conn.tlist = tlist
                        conn.adjust_tm_offsets(run_start)
                        conn.dur = conn.close - conn.open
                        conn.FileRec = fr
                        t = conn.model(tcp_imps, adict)
                        if t:
                            stats.add_TCPStats(t)
			# associate by server
			server = conn.server
			#print 'server %s' % (nprobe.intoa_string(server))
			if not serv_dict.has_key(server):
			    serv_dict[server] = []
			serv_dict[server].append(conn_id)
	## 		# associate by client
## 			client = conn.client
## 			#print 'client %s' % (nprobe.intoa_string(client))
## 			if not cli_dict.has_key(client):
## 			    cli_dict[client] = []
## 			cli_dict[client].append(conn_id)
			
		else:
		    # not interested - remove from dict
		    del(conn_ids[conn_id])
                    ectrs.tf_rej += 1
		    #print 'tfilt_rej'

            if not nread%100:
                overprint('Record %d' % (nread))

            if done:
                break
        
    overprint('%d relevant records' % (nread))

    # register this lot for future reference
    obdict['filepath'] = basepath
    obdict['stats'] = stats
    obdict['serv_dict'] = serv_dict

    #sys.exit(0)

    #
    # In drawing mode - fire up display and pass it connection list
    # - display will run conns through state machine so we can watch
    #
    if do_draw:
        ectrs.printself(tfilt)
        drawlist = []
        for c in conn_list:
            C = c.reconstruct(trace=1)
            C.adjust_tm_offsets(C.abstart)
            drawlist.append(C)
        TcpDisplay(drawlist, standalone='yes', path=basepath,
                   logfun=stats.nullf, trace=1)
	
    else:

        stats.results()
	    
	    
##############################################################################

def sort_a(a, b):
    return len(b[1])-len(a[1])
	    
##############################################################################


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