## ------------------------------------------------------------------------
## 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 sys, os.path, copy
import ws
import constants
import config.silk, silk
import filemanager
import wx.lib.mixins.listctrl as listmix
import analysisresult
import remote
import textwrap
import pickle
import pdb
import re

from filemanager import FileReference, FileListCtrl

from analysistool import RunnableAnalysis
from tools.query import QueryTool
from tools.filter import FilterTool
from tools.uniq import UniqTool
from tools.count import CountTool

from analysisresult import ChangeEventSource

_test_analyses = [
    RunnableAnalysis(QueryTool, name="This is a query test",
        start_date="2004/12/15", end_date="2004/12/16:23", sport=22),
    RunnableAnalysis(FilterTool, name="This is filter #1", sport="22"),
    RunnableAnalysis(FilterTool, name="This is filter #2",
        sport="80,8080,443"),
    RunnableAnalysis(UniqTool, name="Unique sips by sensor",
        fields="sensor", sip_distinct=True)
]

_library = None

import inspect

def GetLibrary():
    return _library
    
    
class LibraryEvent(wx.PyCommandEvent):
   def __init__(self, evt_type, data):
       wx.PyCommandEvent.__init__(self, evt_type, -1)
       self.data = data
   def GetData(self):
       return self.data

        
class Library(ChangeEventSource):
    def __init__(self, libraryFrame):
        ChangeEventSource.__init__(self)
        self.libraryFrame = libraryFrame
        self.lib_path = config.silk.opts["REMOTE_LIBRARY_PATH"].split(":")
        self.refresh()
        # Use a fixed set of scripts for testing
        #self.analyses = _test_analyses
        self.analyses = []
        self.loadAnalyses()
        
    def getLibraryPath(self):
        return copy.copy(self.lib_path)
        
    # Refresh the contents of the library 
    def refresh(self, result=None):
        # Can't do anything until application is initialized
        if not wx.GetApp().IsInitialized():
            self.sets = []
            return
        
        # Do a remote directory listing to find sets in my path
        print "Looking for sets in %s" % self.lib_path
        self.sets = self._findFiles(self.lib_path, ".set")
    
    def saveAnalyses(self):
        filename = os.path.join(config.getPreferencesDir(),
            "library.dat")
        outf = open(filename, "w") 
        for a in self.analyses:
            pickle.dump(a, outf)
        outf.close()
        
    def loadAnalyses(self):
        filename = os.path.join(config.getPreferencesDir(),
            "library.dat")
        if os.path.exists(filename):
            inf = open(filename, "r")
            while True:
                try:
                    a = pickle.load(inf)
                    self.analyses.append(a)
                except EOFError:
                    break
            inf.close()
            return True
        else:
            return False 
    
    def addAnalysis(self, analysis):
        self.analyses.append(analysis)
        self.saveAnalyses()
        print "*Added Analysis to Library*"
        ChangeEventSource._postEvent(self,
            LibraryEvent(ws.EVTTYPE_LIBRARY_ANALYSIS_ADDED, analysis))
        #self.libraryFrame.UpdateContents(None)
    
    def removeAnalysis(self, analysis):
        self.analyses.remove(analysis)
        self.saveAnalyses()
        ChangeEventSource._postEvent(self,
            LibraryEvent(ws.EVTTYPE_LIBRARY_ANALYSIS_DELETED, analysis))
        #self.libraryFrame.UpdateContents(None)        
        
    # import a set file into the library
    def addSet(self, s, directory=None):
        # Copy the file to the set library path
        source_path = s.GetRemotePath()
        context = silk.SimpleRunContext(runRemote=True)
        if directory==None:
            directory = self.lib_path[0]
        silk.copy_remote_file(context, s.GetRemotePath(), directory)
        self.refresh()
        self.libraryFrame.UpdateContents(None, readSets=True)
        
    def getAnalyses(self):
        return self.analyses
        
    def getSets(self):
        return self.sets

    def _findFiles(self, lib_path, ext):
        file_list = filemanager.get_file_list(lib_path, ext)
        return file_list


class RunnableAnalysisList(wx.TreeCtrl):
    def __init__(self, parent, tool_list, *args, **kargs):
        wx.TreeCtrl.__init__(self, parent, -1,
            style=wx.TR_NO_LINES|wx.TR_FULL_ROW_HIGHLIGHT|wx.TR_DEFAULT_STYLE|wx.TR_EDIT_LABELS|wx.TR_SINGLE|wx.TR_HIDE_ROOT)
        self.SetFont(ws.SmFont(wx.SWISS, wx.NORMAL, wx.NORMAL))
        self.tool_list = None
        self.SetList(tool_list)
        self.SetIndent(0)
        il = wx.ImageList(16, 16)
        categories = ("Query", "Filter", "Uniq", "Count")
        img1 = wx.Image("icons/tool.gif", wx.BITMAP_TYPE_GIF)
        img2 = wx.Image("icons/cog-query.gif", wx.BITMAP_TYPE_GIF)
        img3 = wx.Image("icons/cog-f.gif", wx.BITMAP_TYPE_GIF)
        img4 = wx.Image("icons/cog-u.gif", wx.BITMAP_TYPE_GIF)
        img5 = wx.Image("icons/cog-c.gif", wx.BITMAP_TYPE_GIF)        
        il.Add(img1.Rescale(16,16).ConvertToBitmap())
        il.Add(img2.Rescale(16,16).ConvertToBitmap())
        il.Add(img3.Rescale(16,16).ConvertToBitmap())        
        il.Add(img4.Rescale(16,16).ConvertToBitmap())        
        il.Add(img5.Rescale(16,16).ConvertToBitmap())  
        self.AssignImageList(il)

    def SetList(self, tool_list):
        self.DeleteAllItems()
        self._items = {}
        self.tool_list = tool_list
        self.root = self.AddRoot("All")
        for tool in self.tool_list:
            self.AddAnalysis(tool)
        if len(self.tool_list) > 0:
            self.SelectAnalysis(self.tool_list[0])

    def RemoveAnalysis(self, tool):
        item = self._items[tool]
        sel = self.GetPrevSibling(item)
        if not sel.IsOk():
            sel = self.GetNextSibling(item)
        self.Delete(item)
        self.SelectItem(sel)

    def AddAnalysis(self, tool):
        name = tool.GetName()
        name = name.replace("_", " ")
        icon = 0
        categories = ["Query", "Filter", "Uniq", "Count"]
        classes = [QueryTool, FilterTool, UniqTool, CountTool]
        if tool.GetToolClass().GetCategory() in categories:
            icon = categories.index(tool.GetToolClass().GetCategory()) + 1
        item = self.AppendItem(self.root, name, icon)
        self._items[tool] = item
        self.SetPyData(item, tool)

    def SelectAnalysis(self, analysis):
        if analysis == None:
            return
        item = self._items[analysis]
        self.SelectItem(item)

    def GetSelectedAnalysis(self):
        item = self.GetSelection()
        if not item.IsOk(): return None
        return self.GetPyData(item)


class LibraryPanel(wx.Panel):
    def __init__(self, parent):
        global _library
        wx.Panel.__init__(self, parent) #, title="Analysis Library",
                # sizePanel350,-1), style=wx.CLOSE_BOX|wx.CAPTION|wx.RESIZE_BORDER)
        if _library == None:
            _library = Library(self)
        _library.AddChangeListener(self)
        self.lib_path = config.silk.opts["REMOTE_LIBRARY_PATH"].split(":")
        self.nbfont = ws.XSmFont(wx.SWISS, wx.NORMAL, wx.NORMAL)
        self.nb = wx.Notebook(self, style=wx.NB_RIGHT)
        self.nb.SetFont(self.nbfont)
        self._createRecipesPanel()
        #self._createGraphsPanel()
        #self._createSetsPanel()
        #self._createPMapsPanel()
        
        szFrame = wx.GridBagSizer()
        szFrame.Add(self.nb, (0,0), (1,1), wx.TOP|wx.BOTTOM|wx.EXPAND, border=4)
        szFrame.AddGrowableRow(0)
        szFrame.AddGrowableCol(0)
        self.SetSizerAndFit(szFrame)
        self.UpdateContents(None, readSets=True)
        self.Bind(wx.EVT_CLOSE, self.OnClose)
        #self.Bind(wx.EVT_SIZE, self.OnSize) 
        self.Bind(wx.EVT_MENU, self.OnInfo, id=constants.ID_SHOW_INFO)

    def OnInfo(self, evt):
        print "Info selected"
        if self.GetSelection() == self.setsPanel:
            print self.setList.GetSelection() 

    def OnSize(self, evt):
        return
        (w,h) = self.GetSize()
        self.setList.SetSize((-1, h-75))
        evt.Skip(True)

    def OnKeyDown(self, evt):
        if evt.GetKeyCode() in (wx.WXK_BACK, wx.WXK_DELETE):
            self.DeletedSelectedAnalysis()

    def DeletedSelectedAnalysis(self):
        d = wx.MessageDialog(self, "This will permanently remove the analysis from your library.",
                     "Are you sure?", wx.YES_NO | wx.ICON_QUESTION)
        res = d.ShowModal()
        if res == wx.ID_YES:        
            analysis = self.GetSelectedAnalysis()
            _library.removeAnalysis(analysis)
            
    def OnRightDown(self, evt):
        menu = wx.Menu()
        if self.GetSelectedAnalysis() != None:
            menu.Append(constants.ID_LIBRARY_RENAME_ANALYSIS, "Rename")        
            menu.AppendSeparator()
            menu.Append(constants.ID_LIBRARY_DELETE_ANALYSIS, "Delete")
            self.PopupMenu(menu)
            
    def OnDeleteAnalysis(self, evt):
        self.DeletedSelectedAnalysis()
        
    def OnAnalysisDeleted(self, evt):
        tool = evt.GetData()
        self.analysisList.RemoveAnalysis(tool)
        
    def OnAnalysisAdded(self, evt):
        tool = evt.GetData()
        self.analysisList.AddAnalysis(tool)
        
    def OnEditLabel(self, evt):
        analysis = self.GetSelectedAnalysis()
        if evt.IsEditCancelled(): 
            return
        r = re.compile(r'^[\w\_][\w\s\-\_]*$')
        if r.match(evt.GetLabel()):
            analysis.SetName(evt.GetLabel())
            # FIXME: belongs here?
            _library.saveAnalyses()
        else:
            dlg = wx.MessageDialog(self, "Illegal character in name",
                "Illegal name", wx.ICON_EXCLAMATION)
            dlg.ShowModal()
            evt.Veto()
            return
        
    def OnRenameAnalysis(self, evt):
        print "OnRenameAnalysis - not yet implemented"
        self.analysisList.EditLabel(self.analysisList.GetSelection())
    
    def GetSelectedAnalysis(self):    
        return self.analysisList.GetSelectedAnalysis()
        sel = self.analysisList.GetSelection()
        if sel == -1:
            return None
        analyses = GetLibrary().getAnalyses()
        if len(analyses) == 0:
            return None
        analysis = analyses[sel]
        return analysis
            
    def SelectAnalysis(self, analysis):
        self.analysisList.SelectAnalysis(analysis)
            
    def OnSelectAnalysis(self, evt):
        if self.GetSelectedAnalysis() == None:
            enable = False
        else:
            enable = self.GetSelectedAnalysis().CanRun(self.selectedResult)
        self.runAnalysisButton.Enable(enable)
        
    def OnRunAnalysis(self, evt):
        analysis = self.analysisList.GetSelectedAnalysis()
        if analysis==None:
            return
        explorer = wx.GetApp().GetActiveExplorer()
        parentResult = explorer.viewbook.GetSelectedData()
        if analysis.CanRun(parentResult):    
            analysis.Run(config.silk.workRemotely(), True, parentResult, analysis.GetName())
     
    def _createGraphsPanel(self):
        panel = wx.Panel(self.nb)
        nyi = wx.StaticText(panel, -1, " Not yet implemented")
        self.nb.AddPage(panel, "Graphs")
        
    def _createRecipesPanel(self):
        panel = wx.Panel(self.nb)
        self.analysisList = RunnableAnalysisList(panel,
            GetLibrary().getAnalyses(),
            style=wx.SIMPLE_BORDER)
        self.runAnalysisButton = run = ws.XSmButton(panel, -1, "Apply Recipe")
        sz = wx.GridBagSizer()
        sz.Add(self.analysisList, (0,0), (1,1), flag=wx.LEFT|wx.TOP|wx.EXPAND, border=8)
        sz.Add(run, (1,0), (1,1), flag=wx.ALL, border=8)
        panel.SetSizerAndFit(sz)
        sz.AddGrowableRow(0)
        sz.AddGrowableCol(0)
        self.nb.AddPage(panel, "Recipe Library")
        self.Bind(wx.EVT_TREE_ITEM_ACTIVATED, self.OnRunAnalysis,
            self.analysisList)

        self.Bind(wx.EVT_BUTTON, self.OnRunAnalysis, run)
        self.Bind(wx.EVT_TREE_SEL_CHANGED, self.OnSelectAnalysis,
            self.analysisList)
        self.analysisList.Bind(wx.EVT_KEY_DOWN, self.OnKeyDown)
        self.analysisList.Bind(wx.EVT_RIGHT_UP, self.OnRightDown)
        self.Bind(wx.EVT_MENU, self.OnDeleteAnalysis,
            id=constants.ID_LIBRARY_DELETE_ANALYSIS)
        self.Bind(wx.EVT_MENU, self.OnRenameAnalysis,
            id=constants.ID_LIBRARY_RENAME_ANALYSIS)
        self.Bind(wx.EVT_TREE_END_LABEL_EDIT, self.OnEditLabel,
            self.analysisList)
        self.Bind(ws.EVT_LIBRARY_ADD_ANALYSIS,
            self.OnAnalysisAdded)
        self.Bind(ws.EVT_LIBRARY_ANALYSIS_DELETED,
            self.OnAnalysisDeleted)        
        
    def _createSetsPanel(self):
        self.setsPanel = panel = wx.Panel(self.nb)
        self.setList = FileListCtrl(panel, GetLibrary().getSets())
        self.refreshSetsButton  = ws.XSmButton(panel, -1, "Refresh")
        sz = wx.GridBagSizer()
        sz.Add(self.refreshSetsButton, (0,0), (1,1), flag=wx.ALL|wx.ALIGN_LEFT, border=8)
        sz.Add(self.setList, (1,0), (1,1), flag=wx.ALL|wx.EXPAND, border=8)
        panel.SetSizerAndFit(sz)
        sz.AddGrowableRow(1)
        sz.AddGrowableCol(0)        
        self.nb.AddPage(panel, "Remote Sets")
        self.Bind(wx.EVT_BUTTON, self.OnRefreshSets,
            self.refreshSetsButton)
    
    # FIXME: removing sets from right panel
    def OnRefreshSets(self, evt):
        return
        self.items = GetLibrary().getSets()
        self.setList.AddSets(self.items)
        
    def _createPMapsPanel(self):
        panel = wx.Panel(self.nb)
        nyi = wx.StaticText(panel, -1, " Not yet implemented")
        self.nb.AddPage(panel, "Prefix Maps")
        
    def OnLibraryChange(self, evt):
        self.UpdateContents(None)

    # FIXME: removing sets from right panel        
    def UpdateContents(self, result, readSets=False):
        self.selectedResult = result
        self.analysisList.SetList(GetLibrary().getAnalyses())
        analysis = self.GetSelectedAnalysis()
        if analysis == None:
            return
        enable = analysis.CanRun(self.selectedResult)
        self.runAnalysisButton.Enable(enable)
        readSets = False
        # Populate set list
        if readSets:
            self.items = GetLibrary().getSets()
            self.setList.AddSets(self.items)
        self.OnSelectAnalysis(None)
        
    def OnClose(self, evt):
       ''' User closed the window.  Either hide it or quit. '''
       if wx.GetApp().quit:
           evt.Skip(False)
       else:
           self.Hide()


