## ------------------------------------------------------------------------
## 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 config.silk
import silk
import tools
import resultset
import wx, ws
import copy
from controls.previewcmdpanel import PreviewCommandPanel
from filemanager import FilePicker

class UniqTool(tools.AnalysisTool):
    def __init__(self, **kargs):
        tools.AnalysisTool.__init__(self, **kargs)
        
    def CreateDialog(parentResult, parent, title):
        return UniqDialog(parentResult, parent, title=title)
    CreateDialog = staticmethod(CreateDialog)
    
    def CanRun(parentResult):
        if not parentResult:
            return False
        return parentResult.IsFilterResult()
    CanRun = staticmethod(CanRun)
    
    def _GetToolName():
        return "rwuniq"
    _GetToolName = staticmethod(_GetToolName)
    
    def _GetResultClass(self):
        return resultset.UniqResult

    def _GetResultName(self):
        return "Untitled rwuniq count"

    def _RunAnalysis(self, context, inputFileName, outputFileName):
        silk.uniq(context, inputFileName, outputFileName, **self.opts)
        
    def _GetCommandLines(inputFile, outputFile, **kargs):
        cmd = silk.build_uniq_command(False, "", inputFile, outputFile,
            commandsOnly=True, **kargs)
        return cmd
    _GetCommandLines = staticmethod(_GetCommandLines)
        
        
class UniqDialog(wx.Dialog):
    def __init__(self, parentResult, parent, *args, **kargs):
        wx.Dialog.__init__(self, parent, -1, *args, **kargs)
        self.parentResult = parentResult
        # Dialog consists of options, command line preview and the buttons
        self.optionsPanel = self._createKeyOptionsPanel()
        self.countOptionsPanel  = self._createCountOptionsPanel()
        self.previewPanel = PreviewCommandPanel(self,  "Not yet initialized")
        szButtons = self._createButtons()
        
        # Create name editor
        self._createNameEditor()
        # Layout everything
        self.gbs = gbs = wx.GridBagSizer(0,0)
        self.gbs.Add(self.optionsPanel, (0,0), (1,1), wx.LEFT|wx.RIGHT|wx.TOP|wx.EXPAND, border=10)
        self.gbs.Add(self.countOptionsPanel, (0,1), (1,1), wx.LEFT|wx.RIGHT|wx.TOP|wx.EXPAND, border=10)        
        #self.gbs.Add(ws.SmStaticLabel(self, -1, "Command lines"), (1,0), (1,1),
        #    wx.LEFT|wx.ALIGN_BOTTOM, border=10)
        self.gbs.Add(self.previewPanel, (1,0), (1,2), wx.LEFT|wx.RIGHT|wx.EXPAND, border=10)
        self.gbs.Add(self.szName, (2,1), (1,1), wx.ALL|wx.ALIGN_RIGHT, border=10)
        self.gbs.Add(szButtons, (3,0), (1,2), wx.ALL|wx.ALIGN_RIGHT, border=10)
        self.gbs.AddGrowableCol(0)
        self.gbs.AddGrowableCol(1)        
        self.gbs.AddGrowableRow(3)
        self.SetSizerAndFit(self.gbs)
        # Initialize command line preview
        self.RefreshCommandLine()
        self.Bind(wx.EVT_BUTTON, self.OnOK, self.btnCreate) 
        self.Bind(wx.EVT_CHECKBOX, self.OnCheck)

    def _createNameEditor(self):
        self.name = wx.TextCtrl(self, -1, "Untitled rwuniq result", 
            validator=ws.NameValidator(), size=(200,-1))
        self.szName = szName = wx.BoxSizer(wx.HORIZONTAL)
        szName.Add(ws.SmStaticText(self, -1, "Name"),
            flag=wx.RIGHT|wx.ALIGN_CENTER, border=6)
        szName.Add(self.name, flag=wx.ALIGN_CENTER|wx.RIGHT, border=10)

    def GetResultName(self):
        return self.name.GetValue()
        
    def OnCheck(self, evt):
        self.RefreshCommandLine()
        
    def _createButtons(self):
        # Create buttons
        szButtons = wx.BoxSizer(wx.HORIZONTAL)
        szButtons.Add(wx.Button(self, wx.ID_CANCEL, "Cancel"), flag=wx.ALL, border=6)
        self.btnCreate = wx.Button(self, wx.ID_OK, "Run Analysis") #, (110,-1), (100,-1))
        szButtons.Add(self.btnCreate, flag=wx.ALL, border=6)
        # Create is default button
        self.btnCreate.SetDefault()
        # Escape key or close causes cancel
        self.SetEscapeId(wx.ID_CANCEL)
        return szButtons
        
        
    def _create_pmap_controls(self, parent):
        CONTROL_LABEL_WIDTH=-1
        # Create controls
        self.lblpmap = ws.SmStaticText(parent, -1, "Apply a prefix map", size=(CONTROL_LABEL_WIDTH,-1), style=wx.ALIGN_RIGHT)
        self.txtpmap = ws.SmTextCtrl(parent, -1, "(Choose a prefix map)", size=(230,-1), style=wx.TE_READONLY)
        self.btnclrpmap = ws.SmButton(parent, -1, "Clear")
        self.btnpmap = ws.SmButton(parent, -1, "Choose...")
        # Pack buttons
        szpick = wx.FlexGridSizer(1, 2)
        szpick.Add(self.btnclrpmap, flag=wx.BOTTOM|wx.LEFT|wx.RIGHT|wx.ALIGN_RIGHT, border=6)
        szpick.Add(self.btnpmap, flag=wx.BOTTOM|wx.LEFT|wx.RIGHT|wx.ALIGN_RIGHT, border=6)
        # and file name
        sz = wx.FlexGridSizer(2,1)
        sz.Add(self.lblpmap, flag=wx.ALIGN_LEFT|wx.ALIGN_CENTER_VERTICAL|wx.BOTTOM, border=6)
        sz.Add(self.txtpmap, flag=wx.EXPAND|wx.ALIGN_LEFT, border=18)
        sz.Add((0,0))
        sz.Add(szpick, flag=wx.ALL|wx.ALIGN_RIGHT, border=6)
        self.Bind(wx.EVT_BUTTON, self.OnPick, self.btnpmap)
        self.Bind(wx.EVT_BUTTON, self.OnClr, self.btnclrpmap)
        return sz

    def OnClr(self, evt):
       val = "(Choose a file)"
       self.txtpmap.SetValue(val)
       self.cbsval.Enable(False)
       self.cbdval.Enable(False)

    def OnPick(self, evt):
       # Create a set picker
       dlg = FilePicker(self, ext=[".pmap"], cancel_label="No Prefix Map",
        title="Choose prefix map file")
       dlg.CenterOnParent()
       if dlg.ShowModal() == wx.ID_OK:
           val = dlg.GetResult()
           self.cbsval.Enable(True)
           self.cbdval.Enable(True)
       else:
           val = "(Choose a prefix map)"
           self.cbsval.Enable(False)
           self.cbdval.Enable(False)           
       self.txtpmap.SetValue(val)

    def _createKeyOptionsPanel(self):
        # Create top-level options Panel
        optionsPanel = wx.Panel(self)
        box = wx.StaticBox(optionsPanel, -1, "Count by")
        self.cbsip = wx.CheckBox(optionsPanel, -1, "Source IP (sip)")
        self.cbscc = wx.CheckBox(optionsPanel, -1, "Source Country (scc)")
        self.cbsport = wx.CheckBox(optionsPanel, -1, "Source Port (sport)")
        self.cbdip = wx.CheckBox(optionsPanel, -1,"Destination IP (dip)")
        self.cbdcc = wx.CheckBox(optionsPanel, -1, "Destination Country (dcc)")
        self.cbdport = wx.CheckBox(optionsPanel, -1,"Destination Port (dport)")
        self.cbproto = wx.CheckBox(optionsPanel, -1,"Protocol (proto)")
        self.cbsensor = wx.CheckBox(optionsPanel, -1, "Sensor (sensor)")
        self.cblabel = wx.CheckBox(optionsPanel, -1, "Next Hop IP (nhip)")
        
        self.szpmap = self._create_pmap_controls(optionsPanel)
        self.cbsval = wx.CheckBox(optionsPanel, -1, "Source pmap Value (sval)")
        self.cbdval = wx.CheckBox(optionsPanel, -1, "Destination pmap Value (dval)")        
        self.cbsval.Enable(False)
        self.cbdval.Enable(False)
        
        self.cbsport.SetValue(wx.CHK_CHECKED)
        self.gbs = gbs = wx.GridBagSizer(0,0)
        self.gbs.SetEmptyCellSize((0,0))
        szOptions = wx.StaticBoxSizer(box, wx.VERTICAL)
        flag = wx.EXPAND|wx.TOP|wx.LEFT|wx.RIGHT
        szOptions.Add(self.cbsip, flag=flag, border=8)
        szOptions.Add(self.cbscc, flag=flag, border=8)        
        szOptions.Add(self.cbsport, flag=flag, border=8)
        szOptions.Add(self.cbdip, flag=flag, border=8)
        szOptions.Add(self.cbdcc, flag=flag, border=8)                
        szOptions.Add(self.cbdport, flag=flag, border=8)
        szOptions.Add(self.cbproto, flag=flag, border=8)
        szOptions.Add(self.cbsensor, flag=flag, border=8)
        szOptions.Add(self.cblabel, flag=flag, border=8)
        szOptions.Add(wx.StaticLine(optionsPanel, style=wx.HORIZONTAL), flag=flag, border=8)
        szOptions.Add(self.szpmap, flag=flag, border=8)
        szOptions.Add(self.cbsval, flag=flag, border=8)
        szOptions.Add(self.cbdval, flag=flag, border=8)
        
        optionsPanel.SetSizerAndFit(szOptions)
        return optionsPanel
        
    def _createCountOptionsPanel(self):
        # Create top-level options Panel
        optionsPanel = wx.Panel(self)
        box = wx.StaticBox(optionsPanel, -1, "Volume fields")
        self.cbbytes = wx.CheckBox(optionsPanel, -1, "Bytes")
        self.cbpackets = wx.CheckBox(optionsPanel, -1, "Packets")
        self.cbflows = wx.CheckBox(optionsPanel, -1, "Flow Records")
        self.cbuniquesips = wx.CheckBox(optionsPanel, -1, "Unique Source IPs")
        self.cbuniquedips = wx.CheckBox(optionsPanel, -1, "Unique Destination IPs")
        self.cbbytes.SetValue(wx.CHK_CHECKED)
        self.cbpackets.SetValue(wx.CHK_UNCHECKED)
        self.cbflows.SetValue(wx.CHK_UNCHECKED)
        self.gbs = gbs = wx.GridBagSizer(0,0)
        szOptions = wx.StaticBoxSizer(box, wx.VERTICAL)
        flag = wx.EXPAND|wx.TOP|wx.LEFT|wx.RIGHT
        szOptions.Add(self.cbbytes, flag=flag, border=8)
        szOptions.Add(self.cbpackets, flag=flag, border=8)
        szOptions.Add(self.cbflows, flag=flag, border=8)
        szOptions.Add(self.cbuniquesips, flag=flag, border=8)
        szOptions.Add(self.cbuniquedips, flag=flag, border=8)                                
        optionsPanel.SetSizerAndFit(szOptions)
        return optionsPanel

    def SetOptions(self, opts):
        def _setcheck(cb, opt):
            if opt in opts and opts[opt]: cb.SetValue(wx.CHK_CHECKED)
            else: cb.SetValue(wx.CHK_UNCHECKED)
        _setcheck(self.cbbytes, "bytes")
        _setcheck(self.cbpackets, "packets")
        _setcheck(self.cbflows, "flows")
        _setcheck(self.cbuniquesips, "sip_distinct")
        _setcheck(self.cbuniquedips, "dip_distinct")
        if not "fields" in opts:
            return
        fields = opts["fields"].split(",")
        def _setfldscheck(cb, fld):
            if fld in fields:
                cb.SetValue(wx.CHK_CHECKED)
            else:
                cb.SetValue(wx.CHK_UNCHECKED)
        _setfldscheck(self.cblabel, "nhIP")
        _setfldscheck(self.cbsip, "sip")
        _setfldscheck(self.cbdip, "dip")
        _setfldscheck(self.cbscc, "scc")
        _setfldscheck(self.cbdcc, "dcc")        
        _setfldscheck(self.cbsval, "sval")
        _setfldscheck(self.cbdval, "dval")                
        _setfldscheck(self.cbdport, "dport")
        _setfldscheck(self.cbsport, "sport")
        _setfldscheck(self.cbproto, "proto")
        _setfldscheck(self.cbsensor, "sensor")
        if "pmap-file" in opts:
            self.txtpmap.SetValue(opts["pmap-file"])

    def GetOptions(self):
        opts = {}
        if self.cbbytes.GetValue() == wx.CHK_CHECKED: opts["bytes"] = True
        if self.cbpackets.GetValue() == wx.CHK_CHECKED: opts["packets"] = True
        if self.cbflows.GetValue() == wx.CHK_CHECKED: opts["flows"] = True
        if self.cbuniquesips.GetValue() == wx.CHK_CHECKED: opts["sip_distinct"] = True                        
        if self.cbuniquedips.GetValue() == wx.CHK_CHECKED: opts["dip_distinct"] = True                                  
        fields = []
        if self.cblabel.GetValue() == wx.CHK_CHECKED: fields.append("nhIP")
        if self.cbsip.GetValue() == wx.CHK_CHECKED: fields.append("sip")
        if self.cbscc.GetValue() == wx.CHK_CHECKED: fields.append("scc")        
        if self.cbsval.GetValue() == wx.CHK_CHECKED: fields.append("sval")                
        if self.cbsport.GetValue() == wx.CHK_CHECKED: fields.append("sport")
        if self.cbdip.GetValue() == wx.CHK_CHECKED: fields.append("dip")
        if self.cbdcc.GetValue() == wx.CHK_CHECKED: fields.append("dcc")                
        if self.cbdval.GetValue() == wx.CHK_CHECKED: fields.append("dval")                        
        if self.cbdport.GetValue() == wx.CHK_CHECKED: fields.append("dport")
        if self.cbproto.GetValue() == wx.CHK_CHECKED: fields.append("proto")
        if self.cbsensor.GetValue() == wx.CHK_CHECKED: fields.append("sensor")
        opts["fields"] = ",".join(fields)
        if self.txtpmap.GetValue()[0] != "(":
            opts["pmap-file"] = self.txtpmap.GetValue()
        return opts

    def OnCheck(self, evt):
        self.RefreshCommandLine()
        
    def OnOK(self, evt):
        if self.GetOptions()["fields"] == "":
            d = wx.MessageDialog(self,
                "No key field selected.  You must select at least one key field\nto use for counting",
                "Validation Error", wx.ICON_EXCLAMATION)
            d.ShowModal()
        elif len(self.GetOptions().keys()) < 2:
            d = wx.MessageDialog(self,
                "No count field selected.   You must select at least one volume field\nto compute.",
                "Validation Error", wx.ICON_EXCLAMATION)
            d.ShowModal()            
        else:
            evt.Skip(True)
            
    def GetCommandLine(self):
        opts = copy.copy(self.GetOptions())
        if config.silk.opts["WORK_REMOTE"]:
            cmd = silk.build_uniq_command(False, self.parentResult.GetRemoteDir(),
                    self.parentResult.GetFileName(), "$output", **opts)            
        else:
            cmd = silk.build_uniq_command(False, self.parentResult.GetLocalDir(),
                    self.parentResult.GetFileName(), "$output", **opts)
        return cmd

    def RefreshCommandLine(self):
        cmd = self.GetCommandLine()
        self.previewPanel.SetContent(cmd)
    
