## ------------------------------------------------------------------------
## iSiLK - A python-based graphical user interface for SiLK
## ------------------------------------------------------------------------
## Copyright (c) 2007-2008 Carnegie Mellon University. All Rights Reserved.
## ------------------------------------------------------------------------
## Authors: Andrew Kompanek (ajk@cert.org)
## ------------------------------------------------------------------------
## GNU General Public License (GPL) Rights pursuant to Version 2, June 1991
## Government Purpose License Rights (GPLR) pursuant to DFARS 252.225-7013
## ------------------------------------------------------------------------
## See the LICENSE file included in this distribution
## ------------------------------------------------------------------------

import os, string
import shutil
import time
from silk import CommandError

# Types for purposes of comparison in a sort
fieldTypes = {
    "recnum": int,
    "nhip"  : str,
    "class" : str,
    "type" : str,
    "icmp_type" : int,
    "icmp_code" : int,
    "proto" : int,
    "sip"   : lambda ip : map(int,ip.split(".")),
    "dip"   : lambda ip : map(int,ip.split(".")),
    "sport" : int,
    "dport" : int,
    "stime" : str,
    "etime" : str,
    "dur"   :  float,
    "flags" : str,
    "bytes" : float,
    "packets": float,
    "flows"  : float,
    "records": float,
    "sensor" : str,
    "date"   : str,
    "unique_sip" : int,
    "unique_dip" : int
}

def _format_count(val):
    val = float(val)
    fraction = float(val) - int(val)
    l = list(str(int(val)))
    l.reverse()
    s = []
    for i,d in enumerate(l):
        if i > 0 and i%3==0: 
            s.append(",")
        s.append(d)
    s.reverse()
    s = string.join(s,"")
    if fraction != 0:
        s = s + "." + "%02d" % (fraction*100)
    return s
    
def formatFieldAsString(fieldName, val):
    if fieldName in ["bytes", "packets", "flows", "records"]:
        #print "Formatting field %s value %d" % (fieldName, val)
        return _format_count(val)
    else:
        return val

class datarec(object):
    __slots__ = ('datacache', 'vals')
    def __init__(self, datacache, vals):
        self.datacache = datacache
        self.vals = vals
    def getfield(self, field, default="*ERROR*"):
        try:
            index = self.datacache.fields.index(field)
            return self.vals[index]
        except:
            return default
    def __str__(self):
        return self.vals.__str__()    


class SimpleDataCache(object):
    def __init__(self):
        self.cacheFileName = None
        self.recCount = 0
        self._cached = False
        self._records = []
        self.fields = []
        
    def _cacheData(self, cacheFileName, recCount=-1):
        self.cacheFileName = cacheFileName
        if not os.path.exists(cacheFileName):
            return False
        
        # Read the header fields into self.fields
        ascf = open(self.cacheFileName)
        
        
        l = ascf.readline()
        # Verify it looks like ascii
        if l.strip(string.printable) != "":
            msg = "Couldn't read data.  Not a pipe-delimited ascii file?"
            raise CommandError("Error processing file", msg)
        l = l.replace("\n", "")
        l = l.replace(" ", "")
        self.fields = ["recnum"]
        self.fields.extend(l.split("|"))
        print "Fields in cache %s is now %s" % (cacheFileName, self.fields)
        if "" in self.fields:
            self.fields.remove("")
        # Normalize field names
        for i in range(0, len(self.fields)):
           # All lowercase
           self.fields[i] = self.fields[i].lower()
           # Use consistent abbreviation
           if self.fields[i].find("pro") == 0: self.fields[i] = "proto"
           if self.fields[i].find("cla") == 0: self.fields[i] = "class"
           if self.fields[i].find("sen") == 0: self.fields[i] = "sensor"
           if self.fields[i].find("ity") == 0: self.fields[i] = "icmp_type"
           if self.fields[i].find("ico") == 0: self.fields[i] = "icmp_code"
           if self.fields[i].find("sva") == 0: self.fields[i] = "sval"
           if self.fields[i].find("dva") == 0: self.fields[i] = "dval"           
           
        # Read in the records
        self.recCount = 0
        while 1:
            l = ascf.readline()
            if not l : break
            self._records.append(self._parseRec(l, self.recCount))
            self.recCount += 1
        ascf.close()
        self._cached = True

    def GetRecFieldAsString(self, index, field):
        rec = self.GetRecordAt(index)
        val = rec.getfield(field)
        if field == "dport" and rec.getfield("proto") == "1":
            val = int(val)
            icmp_type  = val>>8
            icmp_code = val & 0xFF
            return "%d / %d" % (icmp_type, icmp_code)
        
        return formatFieldAsString(field, val)
    
    def GetRecordCount(self):
        if not self._cached: 
            return 0
        else:
            return self.recCount
           
    def Sort(self, fld, reverse=False):
          # DSU makes a significant performance difference
          fld_type = fieldTypes.get(fld, str)
          aux = [(fld_type(x.getfield(fld)), i) for (i,x) in enumerate(self._records)] 
          aux.sort(reverse=reverse)
          _new_records = [ self._records[i] for (x,i) in aux]
          self._records = _new_records

    def GetRecordAt(self, index):
        return self._records[index]
    
    def _parseRec(self, l, recnum):
        # Split by "|"
        l = l.replace("\n", "")
        l = l.rstrip("|") # remove trailing pipe
        vals = [str(recnum)]
        vals.extend(l.split("|"))
        for i,v in enumerate(vals):
            v = v.rstrip(" ")
            v = v.lstrip(" ")
            vals[i] = v
        rec = datarec(self, vals)
        return rec
