#! /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
from sys import argv
import getopt
from signal import *

from nprobe import set_print_full_hostname, intoa_string, host_string, tcp_open
from nprobe import *
import np_file_util
import string
import os.path

lookups_dict = {}
CACHEFILE_DIR_DEF = '/usr/groups/nprobe/jch1003/ns'
cachefile_path = None
cache_changed = 0
use_long_hostnames = 0

def make_cachefile_path(dir):

    global cachefile_path
    
    if dir == None:
        dir = CACHEFILE_DIR_DEF

    dir = os.path.dirname(dir)

    if dir[-1] != '/':
        dir += '/'

    cachefile_path =  dir + 'lookups.cache'

    return

def use_hostnames():

    global use_long_hostnames

    use_long_hostnames = 1
    set_print_full_hostname(1)

def addr2bin(addr):
    if type(addr) == type(0):
	return addr
    bytes = string.splitfields(addr, '.')
    if len(bytes) != 4: raise ValueError, 'bad IP address'
    n = 0
    for byte in bytes: n = n<<8 | string.atoi(byte)
    return n

def addr2bin_nbo(addr):
    if type(addr) == type(0):
	return addr
    bytes = string.splitfields(addr, '.')
    if len(bytes) != 4: raise ValueError, 'bad IP address'
    n = 0
    bytes.reverse()
    for byte in bytes: n = n<<8 | string.atoi(byte)
    return n

def bin2addr(n):
    return '%d.%d.%d.%d' % ((n>>24)&0xFF, (n>>16)&0xFF,
			    (n>>8)&0xFF, n&0xFF)

def bin2addr_nbo(n):
    return '%d.%d.%d.%d' % (n&0xFF, (n>>8)&0xFF, (n>>16)&0xFF, (n>>24)&0xFF)
    

def np_hostname(addr, verb=0):
    global cache_changed, use_long_hostnames
    #print 'LOOKING UP %s ' % (intoa_string(addr)),
    if use_long_hostnames:
        if lookups_dict.has_key(addr):
            #print 'FOUND %s' % (lookups_dict[addr])
            return lookups_dict[addr]
        else:
            if verb:
                print '%s -> ' % (intoa_string(addr)),
                sys.stdout.flush()
            addrstr = host_string(addr)
            #print 'GOT %s' % (addrstr)			
            lookups_dict[addr] = addrstr
            cache_changed = 1
            if verb:
                print '%s' % (addrstr)
	return addrstr
    else:
        return intoa_string(addr)

def add2lookups_dict(addr, name):
    global cache_changed
    if not lookups_dict.has_key(addr):
	#print 'ADDED %s %s' % (intoa_string(addr), name)
	lookups_dict[addr] = name
	cache_changed = 1

def cache_lookups():

    global cachefile_path
    #print 'cach_lookups()'
    if cache_changed:
	#print 'Saving cached lookups'
	f = open(cachefile_path, 'w')
	for entry in lookups_dict.items():
	    f.write('%s %s\n' % (bin2addr_nbo(entry[0]), entry[1]))

def get_cached_lookups(cachefile_dir=None):
    
    global cachefile_path
    make_cachefile_path(cachefile_dir)
    print 'Using ns cachefile %s' % (cachefile_path)
    try:
	f = open(cachefile_path, 'r')
    except:
	print 'Error opening %s - creating ...' % (cachefile_path)
	return
    lines = f.readlines()
    for line in lines:
	entry = string.split(line)
	#print entry
	lookups_dict[addr2bin_nbo(entry[0])] = entry[1]
	#print '%s = %s' % (entry[0], entry[1])
    #print
    #for e in lookups_dict.items():
	#print '%s -> %s' % (intoa_string(e[0]), e[1])


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

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

    sys.exit(1)

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

def handle_sigint(n, f):


    ans = raw_input('SIGINT - save lookups? (y/n)')
    if ans.count('y'):
        cache_lookups()
    sys.exit(0)

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

    

    scriptname = 'np_lookups.py'
    cachepath = None
    verb = 0

    try:
	optlist, args = getopt.getopt(sys.argv[1:], 'vf:')

    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] == "-f":
            cachepath = opt[1]
        if opt[0] == "-v":
            verb = 1
                

    signal(SIGINT, handle_sigint)
    openfilelist, counters, basepath = np_file_util.get_files(args)
    
    if cachepath == '.' or cachepath == None:
        namecache = basepath
    get_cached_lookups(namecache)

    use_hostnames()
    
    #
    # To get TCP open records into
    #
    tcp_open_rec = tcp_open()

    for file in openfilelist:

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

        while 1:
            rectype = file.next_rec()
            #print rectype
            if rectype == -1: #EOF
                break

            elif rectype == REC_TCP_HTTP_OPEN: # notification of TCP open
                tcp_open_rec.get_open_rec(file)
                caddr = tcp_open_rec.shost()
                saddr = tcp_open_rec.dhost()
                np_hostname(caddr, verb=verb)
                np_hostname(saddr, verb=verb)

            else:
                file.advance()
                continue

    cache_lookups()
	    
##############################################################################


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