#! /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:
                        set = []
                        for rr in r:
                            files[rr[2]] = 1
                                    
            except EOFError:
                break

        print 'Files are:'
        for f in files.keys():
            print f

    def move_all(self, newrecfile):

        recfiles = {}
        #self.getfile('read')
        nfd = os.open(newrecfile, (os.O_CREAT | os.O_WRONLY))
        ngfnm = os.path.dirname(self.fnm) + '/' \
                + os.path.basename(self.fnm).replace('.grabs',
                                                     '.relocated.grabs')
        ngf = open(ngfnm, 'w')

        n = 0
        while 1:
            try:
                read = load(self.gfile)
                for rec in read:
                    if len(rec) == 3: # old type - no rooturl
                        seti = 2
                    elif len(rec) == 4: # new type with rooturl
                        seti = 3
                    else:
                        print 'Recs wrong len %d' % (len(recs))
                        sys.exit(1)
                    newsets = []
                    sets = rec[seti]
                    for r in sets:
                        newr = []
                        for rr in r:
                            rr = list(rr)
                            try:
                                file = recfiles[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)
                                recfiles[rr[2]] = file
                            rr[2] = newrecfile
                            rr[3] = file.copy_rec(rr[3], nfd, n)
                            rr = tuple(rr)
                            newr.append(rr)
                            n += 1
                        newsets.append(newr)
                    rec[seti] = newsets
                    dump([rec], ngf, BINARY_DUMP)
                                    
            except EOFError:
                break
            
        file.write_pseudocounters(nfd, n)

    def draw_objects(self, what, obrefs, obdict):

        oblist = []
        tlist = []
        clist = []

        #
        # If redrawing reload relevant modules - this allows source to be
        # edited and re-run without exiting
        #

        if what.find('Redraw') == 0:
            #reload(np_obnode)
            reload(np_WebHost)
            reload(np_TCOConn)
            reload(np_HTTPTrans)
            what = what.split()[1]
            if what == 'Connections':
                reload(np_TCPDisplay)
            elif what == 'Trees':
                reload(np_Obtree)
                reload(np_TScreen)
        
        for m in obrefs:
            #print m
            ref = m[0]
            if type(ref) == types.TupleType:
                #print 'tuple'
                for o in ref:
                    #print o
                    oblist.append(o)
            else:
                #print 'ob', ref
                oblist.append(ref)
                
        for o in oblist:
            rec = obdict[o][0]
            ob = rec.reconstruct(trace=1)
            if ob.Class == 'WebClient' or ob.Class == 'WebServer' \
                   or ob.Class == 'WebHosts':
                ob.rooturl = obdict[o][1]
                tlist.append(ob)
            elif ob.Class == 'TCPConn':
                clist.append(ob)

##         if len(tlist) and len(clist):
##             print 'Selection includes TCP connections and HTTP Clients'

        if len(tlist) and what == 'Trees':
            np_TScreen.Tscreen(tlist, standalone='no', trace=1)

        if len(clist) and what == 'Connections':
            for c in clist:
                c.adjust_tm_offsets(c.abstart)
            np_TCPDisplay.TcpDisplay(clist, standalone='no', path='', trace=1)

    def reconstruct(self, quickmode):

        self.log = []
        obdict = {}
        tcp_imps = TCP_Imps()

        # collect objects for quick mode
        conns = []
        trees = []

        if quickmode:
            logfun = sys.stdout.write
        else:
            logfun = self.write_log

        logfun('')
        logfun('#######################################################\n')
        logfun('GrabRec - Multi-Log Composite:\n\n')
        logfun('#######################################################\n')

        for grab in self.grabs:

            obj = grab[0].reconstruct(logfun=logfun)
            obdict[obj.id] = (grab[0], grab[2])
            if obj.Class == 'WebHosts' or  obj.Class == 'WebServer' \
                   or  obj.Class == 'WebClient':
                obj.rooturl = grab[2]
                if quickmode:
                    obj.trace = 1
                    trees.append(obj)
                else:
                    obdict[intoa_string(obj.addr)] = (grab[0], grab[2])
                    obj.make_tree(logfun=logfun)
                    # logfun('WebClient #%s %s:' % (intoa_string(obj.addr),
                    #                                     grab[1]))
                    logfun('XMsg Interesting Clients: WebClient: %s: %s' \
                           % (intoa_string(obj.addr), grab[1]))
            elif obj.Class == 'TCPConn':
                if quickmode:
                    conns.append(obj)
                else:
                    obdict['%d' % (obj.id)] = (grab[0], grab[2])
                    #logfun('TCPConn #%d %s:' % (obj.id, grab[1]))
                    logfun('XMsg Interesting Connections: TCPConn: %s: %s'\
                                   % (obj.id, grab[1])) 
                    obj.adjust_tm_offsets(obj.abstart)
                    obj.model(tcp_imps, None)
            else:
                print 'Can\t handle object type \'%s\'' % (obj.Class)

        #for s in self.log:
            #print s

        if quickmode:
            if trees:
                np_TScreen.Tscreen(trees, standalone='yes', trace=1)
            if conns:
                np_TCPDisplay.TcpDisplay(conns, standalone='yes', path='',
                                         trace=1)

        else:

            # display options
            draw_menu = [
                    ['Connections', ['TCPConn'], 0],
                    ['Trees', ['WebClient'], 1],
                    ['Redraw Connections', ['TCPConn'], 1],
                    ['Redraw Trees', ['WebHosts', 'WebClient', 'WebServer'], 1]
                    ]

            if not self.root:
                self.root = Tk()
            Sellist(draw_menu, self.log, obdict, self.draw_objects,
                    root=self.root)

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

def usage(s):
    
    print "usage: " + s + '[-h(elp)] [-q(uick)] [-s(howfiles)] [grabfile]'

    sys.exit(1)
    
##############################################################################
        
def main(argv):
    
    import getopt
    
    scriptname = os.path.basename(argv[0])
    quickmode = 0
    show_files = 0
    moveto = ''
    
    try:
        optlist, args = getopt.getopt(sys.argv[1:], 'hqst:')
        
    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] == '-q':
            quickmode = 1
        if opt[0] == '-s':
            show_files = 1
        if opt[0] == '-t':
            moveto = opt[1]

    if args:
        gfnm = args[0]
    else:
        gfnm = None

    if show_files:
        gr = GrabRec(mode='read', fnm=gfnm)
        gr.show_files()
        return

    if moveto:
        gr = GrabRec(mode='read', fnm=gfnm)
        gr.move_all(moveto)
        return
        
    
    #print 'main'
    gr = GrabRec(mode='read', fnm=gfnm)
    #print 'got grabrec'
    if gr.gfile:
        gr.load()
        # print 'loaded'
        
        # for r, c in gr.grabs:
        #print r
        
        gr.reconstruct(quickmode)


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

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

    main(sys.argv)
