## ------------------------------------------------------------------------
## 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 wx
import ws
import sys, os
import pdb
import copy
import wx.lib.mixins.listctrl as listmix
from timeutils import datetime_obj
from timeutils import silk_datetime_range, silk_date, silk_datetime
import datetime, time
import analysisresult
import constants
import config.viewoptions
import silk

class RecordListCtrl(wx.ListCtrl, listmix.ColumnSorterMixin):
                     #listmix.ListCtrlAutoWidthMixin
                     #):
    def __init__(self, parent, dataset):
        wx.ListCtrl.__init__(self, parent, wx.NewId(), 
                 style=wx.LC_REPORT|wx.LC_VRULES | wx.LC_VIRTUAL|
                       wx.LC_SORT_ASCENDING | wx.BORDER_NONE)
        # Data fields
        self.dataset = dataset
        self._loaded = False
        self._selection = None
        # Display attributes for list items
        self.itemFont = ws.Font(wx.MODERN, wx.NORMAL, wx.NORMAL, False)
        self.evenItemAttr = wx.ListItemAttr(wx.BLACK, wx.WHITE, self.itemFont)
        higlight_color = wx.Color(237, 243, 254)
        highlight_color = wx.SystemSettings.GetColour(wx.SYS_COLOUR_HIGHLIGHT)
        rgbtuple = highlight_color.Get()
        print rgbtuple
        scale = 255.0/max(rgbtuple)
        rgbtuple = (scale*rgbtuple[0], scale*rgbtuple[1], scale*rgbtuple[2])
        print rgbtuple      
        rgbtuple = (240,240,240)  
        highlight_color = wx.Color(rgbtuple[0], rgbtuple[1], rgbtuple[2])
        self.oddItemAttr = wx.ListItemAttr(wx.BLACK,
            highlight_color, self.itemFont)  
        self.il = wx.ImageList(8, 8)        
        self.il.Add(wx.Bitmap("icons/sort_up.png", wx.BITMAP_TYPE_PNG))
        self.il.Add(wx.Bitmap("icons/sort_down.png", wx.BITMAP_TYPE_PNG))
        self.SetImageList(self.il, wx.IMAGE_LIST_SMALL)
        self.sm_up = 0
        self.sm_dn = 1
        self.Bind(wx.EVT_LEFT_DCLICK, self.OnDoubleClick)
        self.Bind(wx.EVT_LIST_ITEM_SELECTED, self.OnItemSelected)
        self.Bind(wx.EVT_LIST_ITEM_RIGHT_CLICK, self.OnRightClick)
        self.Bind(wx.EVT_LIST_COL_CLICK, self.OnColumnClick)
        self.Bind(wx.EVT_MENU, self.OnResolveHostNames, id=constants.ID_POPUP_RESOLVE_HOSTNAMES)
        self.Bind(wx.EVT_MENU, self.OnDrillDown, id=constants.ID_POPUP_DRILLDOWN_SIP)
        self.Bind(wx.EVT_MENU, self.OnDrillDown, id=constants.ID_POPUP_DRILLDOWN_DIP)
        self.Bind(wx.EVT_MENU, self.OnDrillDown, id=constants.ID_POPUP_DRILLDOWN_SPORT)
        self.Bind(wx.EVT_MENU, self.OnDrillDown, id=constants.ID_POPUP_DRILLDOWN_DPORT)
        self.Bind(wx.EVT_MENU, self.OnDrillDown, id=constants.ID_POPUP_DRILLDOWN_PROTO)
        self.Bind(wx.EVT_MENU, self.OnDrillDown, id=constants.ID_POPUP_DRILLDOWN_SENSOR)
        self.Bind(wx.EVT_MENU, self.OnDrillDown, id=constants.ID_POPUP_DRILLDOWN_SIP_AND_DIP)
        self.Bind(wx.EVT_MENU, self.OnDrillDown, id=constants.ID_POPUP_DRILLDOWN_SIP_AND_SPORT)        
        self.Bind(wx.EVT_MENU, self.OnDrillDown, id=constants.ID_POPUP_DRILLDOWN_DIP_AND_DPORT)   
        self.Bind(wx.EVT_MENU, self.OnQuery, id=constants.ID_POPUP_QUERY_SIP)
        self.Bind(wx.EVT_MENU, self.OnQuery, id=constants.ID_POPUP_QUERY_DIP)
        self.Bind(wx.EVT_MENU, self.OnQuery, id=constants.ID_POPUP_QUERY_SPORT)
        self.Bind(wx.EVT_MENU, self.OnQuery, id=constants.ID_POPUP_QUERY_DPORT)
        self.Bind(wx.EVT_MENU, self.OnQuery, id=constants.ID_POPUP_QUERY_PROTO)
        self.Bind(wx.EVT_MENU, self.OnQuery, id=constants.ID_POPUP_QUERY_SENSOR)        
        self.Bind(wx.EVT_MENU, self.OnQuery, id=constants.ID_POPUP_QUERY_SIP_AND_DIP)
        self.Bind(wx.EVT_MENU, self.OnQuery, id=constants.ID_POPUP_QUERY_SIP_AND_SPORT)        
        self.Bind(wx.EVT_MENU, self.OnQuery, id=constants.ID_POPUP_QUERY_DIP_AND_DPORT)                
        
    def OnDoubleClick(self, evt):
        rec = self._selection
        print "Selected %s" % rec
        opts = {}
        fields = ["proto", "sip", "dip", "sport", "dport", "date"]
        for f in fields:
            val = rec.getfield(f, None)
            if f == "date" and val:
                start = datetime_obj(val)
                end = start + datetime.timedelta(seconds=self.dataset.opts["bin-size"])
                val = silk_datetime_range(start,end)
                opts["stime"] = val
            else:
                if f == "sip": f = "saddress"
                if f == "dip": f = "daddress"
                if val: opts[f] = val
                
        if len(opts.keys()) == 0:
            return
        parent = wx.GetApp().GetSelectedFilterData()
        wx.GetApp().refine(parent=parent, name="Drilldown", **opts)
        
    def OnResolveHostNames(self, event):
        #rec = self._selection

        iplist = []

        rec = self.GetNextSelected(-1)
        while rec <> -1:
            iplist.append(self.dataset.GetRecordAt(rec).getfield("sip"))
            iplist.append(self.dataset.GetRecordAt(rec).getfield("dip"))
            rec = self.GetNextSelected(rec)

        self.dataset.cacheHostName(iplist)
        
    def OnRightClick(self, event):
        self._selectionIndex = event.m_itemIndex
        self._selection = self.dataset.GetRecordAt(event.m_itemIndex)
        rec = self._selection
        if rec == None:
            return
        ip = rec.getfield("ip", None)
        print "the ip is %s" % ip
        sip = rec.getfield("sip", None)
        dip = rec.getfield("dip", None)
        sport = rec.getfield("sport", None)
        dport = rec.getfield("dport", None)
        proto = rec.getfield("proto", None)
        sensor = rec.getfield("sensor", None)
        # Build the menu
        menu = wx.Menu()
        # Don't pop up menu if we don't the needed fields
        if (ip, sip, dip, sport, dport, proto, sensor) == (None, None, None, None, None, None):
            return
        if (sip, dip, ip) != (None, None, None):
            menu.Append(constants.ID_POPUP_RESOLVE_HOSTNAMES, "Resolve Host Names")
            menu.AppendSeparator()
        dd = wx.Menu()
        q = wx.Menu()
        menu.AppendSubMenu(dd, "Drill down")
        menu.AppendSubMenu(q, "Query")
        if sip and dip:
            dd.Append(constants.ID_POPUP_DRILLDOWN_SIP_AND_DIP, "%s => %s" % (sip, dip))
            q.Append(constants.ID_POPUP_QUERY_SIP_AND_DIP, "%s => %s" % (sip, dip))            
        if sip and sport:
            dd.Append(constants.ID_POPUP_DRILLDOWN_SIP_AND_SPORT, "%s:%s => *" % (sip, sport))
            q.Append(constants.ID_POPUP_QUERY_SIP_AND_SPORT, "%s:%s => *" % (sip, sport))                        
        if dip and dport:
            dd.Append(constants.ID_POPUP_DRILLDOWN_DIP_AND_DPORT, "* => %s:%s" % (dip, dport))
            q.Append(constants.ID_POPUP_QUERY_DIP_AND_DPORT, "* => %s:%s" % (dip, dport))
        
        if sip:
            dd.Append(constants.ID_POPUP_DRILLDOWN_SIP, "%s => x.x.x.x" % sip)
            q.Append(constants.ID_POPUP_QUERY_SIP, "%s => x.x.x.x" % sip)            
        if dip:
            dd.Append(constants.ID_POPUP_DRILLDOWN_DIP, "x.x.x.x => %s" % dip)
            q.Append(constants.ID_POPUP_QUERY_DIP, "x.x.x.x => %s" % dip)  
        
        if ip:
            print "Building menu based on ip = %s" % ip
            dd.Append(constants.ID_POPUP_DRILLDOWN_SIP, "%s => x.x.x.x" % ip)
            q.Append(constants.ID_POPUP_QUERY_SIP, "%s => x.x.x.x" % ip)
            dd.Append(constants.ID_POPUP_DRILLDOWN_DIP, "x.x.x.x => %s" % ip)
            q.Append(constants.ID_POPUP_QUERY_DIP, "x.x.x.x => %s" % ip)            
            
        if sip != None or dip !=None:        
            dd.AppendSeparator()        
            q.AppendSeparator()          
        dd.AppendSeparator()
        q.AppendSeparator()
        if sport:
            dd.Append(constants.ID_POPUP_DRILLDOWN_SPORT, "Port %s => (Any)" % sport)
            q.Append(constants.ID_POPUP_QUERY_SPORT, "Port %s => (Any)" % sport)            
        if dport:
            dd.Append(constants.ID_POPUP_DRILLDOWN_DPORT, "(Any) => Port %s" % dport)
            q.Append(constants.ID_POPUP_QUERY_DPORT, "(Any) => Port %s" % dport)    
        if sport != None or dport !=None:        
            dd.AppendSeparator()        
            q.AppendSeparator()                
        if proto:
            dd.Append(constants.ID_POPUP_DRILLDOWN_PROTO, "Protocol %s" % proto)
            q.Append(constants.ID_POPUP_QUERY_PROTO, "Protocol %s" % proto)
        if sensor:
            dd.Append(constants.ID_POPUP_DRILLDOWN_SENSOR, "Sensor %s" % sensor)
            q.Append(constants.ID_POPUP_QUERY_SENSOR, "Sensor %s" % sensor)
        self.PopupMenu(menu)
        
    def OnDrillDown(self, event):
        rec = self._selection
        print "OnDrillDown : %d" % event.GetId()
        opts = {}
        if event.GetId() == constants.ID_POPUP_DRILLDOWN_SIP:
            opts["saddress"] = rec.getfield("sip", rec.getfield("ip"))
            name = "Drilldown - %s" % opts["saddress"]
        if event.GetId() == constants.ID_POPUP_DRILLDOWN_DIP:            
            opts["daddress"] = rec.getfield("dip", rec.getfield("ip"))           
            name = "Drilldown - %s" % opts["daddress"]
        if event.GetId() == constants.ID_POPUP_DRILLDOWN_SPORT:
            opts["sport"] = rec.getfield("sport")
            name = "Drilldown - %s %s" % ("sport", rec.getfield("sport"))                         
        if event.GetId() == constants.ID_POPUP_DRILLDOWN_DPORT:
            opts["dport"] = rec.getfield("dport")            
            name = "Drilldown - %s %s" % ("dport", rec.getfield("dport"))                                     
        if event.GetId() == constants.ID_POPUP_DRILLDOWN_PROTO:
            opts["proto"] = rec.getfield("proto")
            name = "Drilldown - %s %s" % ("proto", rec.getfield("proto"))                                             
        if event.GetId() == constants.ID_POPUP_DRILLDOWN_SIP_AND_DIP:
            opts["saddress"] = rec.getfield("sip")
            opts["daddress"] = rec.getfield("dip")
            name = "Drilldown - %s to %s" % (rec.getfield("sip"), 
                                 rec.getfield("dip"))     
        if event.GetId() == constants.ID_POPUP_DRILLDOWN_SIP_AND_SPORT:
            opts["saddress"] = rec.getfield("sip")
            opts["sport"] = rec.getfield("sport")
            name = "Drilldown"
        if event.GetId() == constants.ID_POPUP_DRILLDOWN_DIP_AND_DPORT:
            opts["daddress"] = rec.getfield("dip")
            opts["dport"] = rec.getfield("dport")
            name = "Drilldown"
        if event.GetId() == constants.ID_POPUP_DRILLDOWN_SENSOR:
            opts["sensor"] = rec.getfield("sensor")
            name = "Drilldown - %s %s" % ("sensor", rec.getfield("sensor"))
        parent = wx.GetApp().GetSelectedFilterData()
        print "Got name %s" % name
        # In case ip was a cidr
        if "/" in name:
            (name,bits) = name.split("/")
        wx.GetApp().refine(parent=parent, name=name, **opts)
        
    def OnQuery(self, event):
        rec = self._selection
        opts = {}
        print "GetParent() = %s" % self.dataset.GetParent()
        parent = self.dataset.GetParentQuery()
        print "parent = %s" % parent
        print "query parent: %s" % parent.GetParentQuery()
        print "query opts = %s" % parent.opts
        
        parent_opts = silk._normalize_args(parent.opts)
        
        if parent:
            if "start-date" in parent_opts:
                opts["start-date"] = parent_opts["start-date"]
            if "end-date" in parent_opts:
                opts["end-date"] = parent_opts["end-date"]
        opts = silk._denormalize_args(opts)
        print "  --- opts = %s" % opts
        

        if event.GetId() == constants.ID_POPUP_QUERY_SIP:
            opts["saddress"] = rec.getfield("sip", rec.getfield("ip"))
        if event.GetId() == constants.ID_POPUP_QUERY_DIP:
            opts["daddress"] = rec.getfield("dip", rec.getfield("ip"))
        if event.GetId() == constants.ID_POPUP_QUERY_SPORT:
            opts["sport"] = rec.getfield("sport")
        if event.GetId() == constants.ID_POPUP_QUERY_DPORT:
            opts["dport"] = rec.getfield("dport")
        if event.GetId() == constants.ID_POPUP_QUERY_PROTO:
            opts["proto"] = rec.getfield("proto")
        if event.GetId() == constants.ID_POPUP_QUERY_SIP_AND_DIP:
            opts["saddress"] = rec.getfield("sip")
            opts["daddress"] = rec.getfield("dip")
        if event.GetId() == constants.ID_POPUP_QUERY_SIP_AND_SPORT:
            opts["saddress"] = rec.getfield("sip")
            opts["sport"] = rec.getfield("sport")
        if event.GetId() == constants.ID_POPUP_QUERY_DIP_AND_DPORT:
            opts["daddress"] = rec.getfield("dip")
            opts["dport"] = rec.getfield("dport")
        if event.GetId() == constants.ID_POPUP_QUERY_SENSOR:
            opts["sensor"] = rec.getfield("sensor")
        wx.GetApp().DoShowQueryBuilder(opts=opts)
    
    def OnItemSelected(self, event):
        ##print event.GetItem().GetTextColour()
        self._selectionIndex = event.m_itemIndex
        self._selection = self.dataset.GetRecordAt(event.m_itemIndex)

    def _addColumns(self, hide_fields = []):
        self.ClearAll()
#        all_fields = ["#"]
        all_fields = []
        all_fields.extend(self.dataset.fields)#(copy.copy(self.dataset.fields))
        i = 0
        self.fields = []
        self.reverse_flags = []
        for f in all_fields:
            if f in hide_fields:
                continue
            self.fields.append(f)
            self.reverse_flags.append(True)
            label = ws.GetLabelText(f)
            if label == "recnum":
                label = "#"
            # Insert it and compute the width, then delete
            self.InsertColumn(i, ws.GetTextPattern(f))
            self.SetColumnWidth(i, -2)
            width = self.GetColumnWidth(i)
            self.DeleteColumn(i)
            # Insert and set this width
            self.InsertColumn(i, label)
            li = wx.ListItem()
            li.SetAlign(ws.GetTextAlignment(f))                    
            self.SetColumn(i, li)
            #print "Ideal width for %s is %d" % (f, width)
            # FIXME: this approach isn't working properly on windows
            # (very wide columsn in some cases).  This is a temporary fix.
            if wx.Platform == "__WXMSW__":
                width = 6 * (len(ws.GetTextPattern(f))) + 15
            self.SetColumnWidth(i, width)
            i += 1
        # Initialize the item count
        self.SetItemCount(self.dataset.GetRecordCount())
        # Initialize the sorter mix in (what does this really buy me?)
        listmix.ColumnSorterMixin.__init__(self, len(self.fields))
        #if len(self.fields) > 4:
        #    listmix.ListCtrlAutoWidthMixin.__init__(self)

    def OnGetItemText(self, item, c):
        index = item
       # if c == 0:
    #        s = self.dataset.GetRecFieldAsString(index, "recnum")
    #        return str(s)
        try:
            s = self.dataset.GetRecFieldAsString(index, self.fields[c],
                lookupPortAndProtocol = self.dataset.GetViewOption("RECVIEW_SHOW_NAMES", True),
                lookupHostName = self.dataset.GetViewOption("RECVIEW_SHOW_HOSTNAMES", True)) 
        except Exception, e:
            print e
            return "<Err>"
        return s
        
    def GetClipboardObject(self, format=constants.CLIPBOARD_FORMAT_DEFAULT):
        if format == constants.CLIPBOARD_FORMAT_LATEX_TABLE:
            latex_str = self.GetLatexMarkup()
            return wx.TextDataObject(latex_str)
        else:
            fieldsep = "\t"
            hdr_str = ""
            for c in range(1, len(self.fields)):
                hdr_str += self.fields[c] + fieldsep
            recs_str = self.GetTableMarkup()
            clip_str = hdr_str + "\n" + recs_str
            return wx.TextDataObject(clip_str)

    def GetTableMarkup(self, fieldsep="\t", linesep="\n"):
        clip_str = ""
        MAX_CLIPBOARD_ROWS = 25000
        count = min(MAX_CLIPBOARD_ROWS, self.dataset.GetRecordCount())
        for index in range(0, count):
            sel_state = self.GetItemState(index, wx.LIST_STATE_SELECTED)
            if (self.GetSelectedItemCount()==0) or (sel_state == wx.LIST_STATE_SELECTED):
                for c in range(1, len(self.fields)):
                    s = self.dataset.GetRecFieldAsString(index, self.fields[c],
                        lookupPortAndProtocol = self.dataset.GetViewOption("RECVIEW_SHOW_NAMES", True),
                        lookupHostName = self.dataset.GetViewOption("RECVIEW_SHOW_HOSTNAMES", True)) 
                    clip_str += s
                    if c < (len(self.fields)-1):
                        clip_str += fieldsep
                clip_str += linesep
        # remove the last newline
        clip_str = clip_str.rstrip("\n")
        return clip_str
        
    def GetLatexMarkup(self):
        fieldsep = " & "
        hdr_str = ""
        hdr_format_str = ""
        for c in range(1, len(self.fields)):
            hdr_str += self.fields[c]
            if c < (len(self.fields)-1):
                hdr_str += fieldsep
            hdr_format_str += "c|"
        clip_str  = "\\begin{tabular}"
        clip_str += "{|"
        clip_str += hdr_format_str
        clip_str += "}\n"
        clip_str += "\\hline\n"
        clip_str += hdr_str
        clip_str += " \\\\\n"
        clip_str += "\\hline\n\\hline\n"
        clip_str += self.GetTableMarkup(fieldsep=" & ", linesep=" \\\\\n")
        clip_str += "\n\\hline\n"
        clip_str += "\\end{tabular}"
        return clip_str

    def OnGetItemAttr(self, index):
        if index %2: return self.oddItemAttr
        else: return self.evenItemAttr
        
    def GetListCtrl(self):
        return self

    def GetSortImages(self):
        return (self.sm_dn, self.sm_up)  
        
    # Don't use the sorter mixins implementation since it requires
    # yet another copy of the data.  Instead ask the data to sort
    # itself.
    def OnColumnClick(self, evt):
        self.DoSort()
        
    # Note: still rely on some superclass to set _col
    def DoSort(self):
         busy = wx.BusyInfo("Sorting records...")
         if self._col == 0:
             field = "recnum"
         else:
             field = self.fields[self._col]
             
         reverse_flag = self.reverse_flags[self._col]
         self.reverse_flags[self._col] = not reverse_flag
         self.dataset.Sort(field, reverse_flag)
         self.Refresh()

    def GetSelectedRecord(self):
        index = self.GetSelection()
        return self.GetRecordAt(index)
        
    def RebuildView(self):
        hide_fields = self.dataset.GetViewOption("RECVIEW_HIDE_COLUMNS", [])
        #print "Hiding fields: %s" % hide_fields
        self._addColumns(hide_fields = hide_fields)
        self.RefreshContents()

    def RefreshContents(self):
        if self.dataset.IsLocalFileAvailable() and not self._loaded:
            if self.dataset.GetRecordCount() == 0:
                return
            self._loaded = True
            self.Show(True)
            # Build our column headers now that we've got the data
            hide_fields = self.dataset.GetViewOption("RECVIEW_HIDE_COLUMNS", [])
            print "Hiding fields: %s" % hide_fields
            self._addColumns(hide_fields = hide_fields)
            self.SetItemCount(self.dataset.GetRecordCount())
        else:
            if not self._loaded:
                self.Show(False)
            #print "Refreshing"
            #print "viewOptions = %s" % self.dataset.GetViewOptions()
            if False:
                hide_fields = self.dataset.GetViewOption("RECVIEW_HIDE_COLUMNS", [])
                #print "Hiding fields: %s" % hide_fields
                self._addColumns(hide_fields = hide_fields)
            self.Refresh()
