#! /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
import getopt
from random import shuffle

from DMuxer import FlowDmuxer, FlowDmuxerException
from FlowDmuxerDmux import FlowDmuxerDmux

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

class FlowDmuxerCopy(FlowDmuxer):

    '''
    Class FlowDmuxerCopy (file dmux.py) - sub-class of FlowDmuxer

      Copy packets passing filter from input file into output file - i.e. same
      as \'tcpdump -rxxx -wyyy expr\'.

      Options: None
      '''

    def __init__(self, infile, filtexp, opath=None, raise_natural_exceptions=0):

        FlowDmuxer.__init__(self, infile, filtexp, opath, raise_natural_exceptions=raise_natural_exceptions)

        if self.ofile:
            d = '.'
        else:
            d = ''

        if self.ofile != '-':
            passfile = os.path.join(self.odir, self.ofile + d + 'passed')
        else:
            passfile = self.ofile


        try:
            self.pcap.dump_open(passfile)
        except (IOError, TypeError), s:
            self.exception(str(s))
        
############################################################################

    def __del__(self):
        
        try:
            self.pcap.dump_close()
        except IOError, s:
            self.exception(str(s))
        
        FlowDmuxer.__del__(self)

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

    def loop(self, count):
        
        try:
            self.pcap.loop(count)
        except TypeError, s:
            self.exception(str(s))

        
        
############################################################################
            
    def check(self):

         raise FlowDmuxerException('Check not implemented')
        
############################################################################
############################################################################

class FlowDmuxerSplit(FlowDmuxer):

    '''
    Class FlowDmuxerSplit (file dmux.py) - sub-class of FlowDmuxer

      Copy packets from input file into two output files - \'passed\' and
      \'not_passed\' oas selected by filter expression.

      Options: None
      '''

    def __init__(self, infile, filtexp, opath=None, raise_natural_exceptions=0):

        self.of = None

        FlowDmuxer.__init__(self, infile, '', opath, raise_natural_exceptions=raise_natural_exceptions)

        if self.ofile:
            d = '.'
        else:
            d = ''
        passfile = os.path.join(self.odir, self.ofile + d + 'passed')
        failfile = os.path.join(self.odir, self.ofile + d + 'not_passed')

        pcap = self.pcap

        try:
            pcap.py_pcap_compile(filtexp)
            pcap.dump_open(passfile)
            self.of = pcap.alt_dump_open(failfile)
        except (IOError, TypeError), s:
            self.exception(str(s))
        
############################################################################

    def __del__(self):
        
        try:
            if self.of:
                self.of.close()
            self.pcap.dump_close()
        except IOError, s:
            self.exception(str(s))
        
        FlowDmuxer.__del__(self)

        
        
############################################################################
            
    def process_pkt(self, (len, data, ts)):
         
         try:
             if self.pcap.filter():
                 self.pcap.dump()
             else:
                 self.pcap.alt_dump(self.of)
         except IOError, s:
             self.exception(str(s))

        
        
############################################################################
            
    def check(self):

         raise FlowDmuxerException('Check not implemented')
        
############################################################################
############################################################################

class FlowDmuxerJumble(FlowDmuxer):

    '''
    Class FlowDmuxerJumble (file dmux.py) - sub-class of FlowDmuxer

      Output packets, optionally selected from input file by filter
      expression, shuffled to simulate gross reordering in network.

      Options: None
      '''

    def __init__(self, infile, filtexp, opath=None, raise_natural_exceptions=0):

        self.of = None

        FlowDmuxer.__init__(self, infile, filtexp, opath, raise_natural_exceptions=raise_natural_exceptions)

        if self.ofile:
            d = '.'
        else:
            d = ''
        ofile = os.path.join(self.odir, self.ofile + d + 'jumbled')
        print ofile

        try:
            self.of = self.pcap.alt_dump_open(ofile)
        except (IOError, TypeError), s:
            self.exception(str(s))

        self.jumbled = []
        self.tss = []
        self.pcap_next = self.pcap.next_o_tts
        
############################################################################

    def __del__(self):
        
        try:
            if self.of:
                self.of.close()
            self.pcap.dump_close()
        except IOError, s:
            self.exception(str(s))
        
        FlowDmuxer.__del__(self)

        
        
############################################################################
            
    def process_pkt(self, (off, len, data, ts)):

         self.jumbled.append(off)
         self.tss.append(ts)

        
        
############################################################################
            
    def tidy(self):

        shuffle(self.jumbled)

        jumbled = zip(self.jumbled, self.tss)

        pcap = self.pcap
        seek = pcap.fseek
        next = pcap.next_fts
        dump = pcap.alt_dump
        f = self.of
        for off, ts in jumbled:
            seek(off)
            next()
            print ts[0], ts[1]
            pcap.set_ts(ts)
            dump(f)
            

        

        
        
############################################################################
            
    def check(self):

         raise FlowDmuxerException('Check not implemented')

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


    
def help(level, status):
    
    if not level:
        return

    if level == 'dmux':
        print main.__doc__
    elif level == 'd':
        print FlowDmuxer.__doc__
    elif level == 'j':
        print FlowDmuxerJumble.__doc__
    elif level == 's':
        print FlowDmuxerSplit.__doc__
    elif level == 'e':
        print FlowDmuxerException.__doc__
    elif level == 'f':
        print FlowDmuxerDmux.__doc__
    elif level == 'c':
        print FlowDmuxerCopy.__doc__

    sys.exit(status)
    
############################################################################

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

def main():

    '''
    dmux.py

       Interface script to py_libpcap via FlowDmuxer class (DMuxer.py)
       - also contains FlowDmuxer_Copy, FlowDmuxer_Split and FlowDmuxer_Jumble
         sub-classes.

       Usage:

         dmux.py [option(s)] [expression]

         options:
            -r<file>  Read from tcpdump file <file> (pylibpcap does not
                      currently support a live packet filter). '-' = stdin.
            -c<count> Read in <count> packets.
            -w<path>  Specify path for output file(s) - if not specified will
                      use input file directory, or cwd for stdin input.
            -e        Throw natural exceptions (default is FlowDmuxerException
                      for everything).
            -h        This help.
            -H<x>     Help on Dmuxer (sub)class x:
                         d FlowDmuxer base class
                         c FlowDmuxerCopy sub-class
                         s FlowDmuxerSplit
                         j FlowDmuxerJumble
                         f FlowDmuxerDmux
                         e FlowDmuxerException
            --copy   Invoke FlowDmuxer_Copy class - just copy packets passing
                     filter defined by expression.
            --split  Invoke FlowDmuxer_Split class - split input into files
                     passing/not passing filter.
            --jumble Invoke FlowDmuxer_Jumble class - as copy but shuffles
                     output packets to simulate gross reordering in network.
            --dmux   Invoke FlowDmuxerDmux - dmux all flows into tree ordered
                     by protocol, address, port etc. as appropriate.
            For class-specific options use -H<x>

         [expression] Optional packet filter expression - use and syntax
                      exactly as for tcpdump.
       '''
    
    def die(msg):
    
        sys.stderr.write(scriptname + ' ERROR: ' + msg + '\n')
        sys.exit(1)
    
    scriptname = os.path.basename(sys.argv[0])
    readfile =  None
    odir = None
    cnt = -1
    copy = 0
    split = 0
    dmux = 0
    rne = 0
    check = 0
    jumble = 0
    npe = 0
    helplevel = ''

    shortopts = 'r:c:w:ehH:'
    longopts = ['split', 'copy', 'dmux', 'jumble', 'accept_nosyns', 'timeout=', 'float_timestamps', 'check', 'nprobe-errors=']

    try:
        optlist, args = getopt.getopt(sys.argv[1:], shortopts, longopts)
    except getopt.error, s:
        die(str(s))

    dmuxer_opts = []

    for opt, par in optlist:
        #print opt, par
        if opt == '-r':
            readfile = par
        elif opt == '-w':
            odir = par
        elif opt == '-c':
            cnt = int(par)
        elif opt == '-e':
            rne = 1
        elif opt == '-h':
            helplevel = 'dmux'
        elif opt == '-H':
            helplevel = par
        elif opt == '--split':
            split = 1
        elif opt == '--copy':
            copy = 1
        elif opt == '--dmux':
            dmux = 1
        elif opt == '--check':
            check = 1
        elif opt == '--jumble':
            jumble = 1
        elif opt == '--nprobe-errors':
            npe = int(par)
        else:
            if par:
                eq = '='
            else:
                eq = ''
            dmuxer_opts.append(opt+eq+par)

    help(helplevel, 0)

    if copy + split + dmux + jumble != 1:
        die('one and only one of long args --copy, --split, --dmux or --jumble must be given')

    if args:
        exp = args[0]
        for arg in args[1:]:
            exp += ' ' + arg
    else:
        exp = ''

    try:
        if split:
            dmuxer = FlowDmuxerSplit(readfile, exp, odir, raise_natural_exceptions=rne)
        elif copy:
            dmuxer = FlowDmuxerCopy(readfile, exp, odir, raise_natural_exceptions=rne)
        elif dmux:
            dmuxer = FlowDmuxerDmux(readfile, exp, odir, raise_natural_exceptions=rne, save_check_info=check, nprobe_errors=npe)
        elif jumble:
            dmuxer = FlowDmuxerJumble(readfile, exp, odir, raise_natural_exceptions=rne)
            
    except FlowDmuxerException, s:
        die('FlowDmuxer*** instantiation: ' + str(s))

    try:
        dmuxer.getopts(dmuxer_opts)
    except FlowDmuxerException, s:
        die('FlowDMuxer*** getopt(): ' + str(s))

    
    try:
        dmuxer.loop(cnt)
    except FlowDmuxerException, s:
        die('DMuxer loop: ' + str(s))

    try:
        dmuxer.tidy()
    except FlowDmuxerException, s:
        die('DMuxer tidy: ' + str(s))

    if check:
        try:
            dmuxer.check()
        except FlowDmuxerException, s:
            die('DMuxer check: ' + str(s))

    
    dmuxer.report()
        
            
##############################################################################


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