## ------------------------------------------------------------------------
## PySilk - A python-based graphical user interface for SiLK
## ------------------------------------------------------------------------
## Copyright (C) 2006-2007 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 LICENSE file included in this distribution
## ------------------------------------------------------------------------

import wx, ws
import logging
import threading
import os, sys, string, tempfile, traceback, re, shutil
import analysisresult
from constants import *
from analysisresult import ProblemSet
from tools.query import QueryTool
from resultset import QueryResult
from jobmanager import JobManager
from querybuilder import QueryBuilder
from explorer import ProblemSetExplorer
import analysistool, tools, tools.shell
import optparse
from about import AboutBox
from config.editor import PreferencesDialog
# Application configuration modules
import config
import config.silk
import config.filteroptions
import config.fgloboptions
import config.editor, config.recordviewoptions
from infoframe import InfoFrame

wxEVT_INVOKE = wx.NewEventType()
class InvokeEvent(wx.PyEvent):
     def __init__(self, func, args, kwargs):
         wx.PyEvent.__init__(self)
         self.SetEventType(wxEVT_INVOKE)
         self.__func = func
         self.__args = args
         self.__kwargs = kwargs
     def invoke(self):
         self.__func(*self.__args, **self.__kwargs)
         
# self.Connect(-1, -1, wxEVT_INVOKE, self.OnInvoke)
#def InvokeLater(self, f, *args, **kwargs):
#    self.GetEventHandler().AddPendingEvent(InvokeEvent(f, args, kwargs))
#def LogMessage(self, msg):
#    print "Appending text to control"
#    self.InvokeLater(self.DoLogMessage, msg)

_running_ipython = False

class SilkApp(wx.App):
    

    def FilterEvent(self, event):
        print "FilterEvent: %s" % event
        
    
    def _processOptions(self):
        # Parse options
        p = optparse.OptionParser()
        p.set_description('''With no options are given, a new empty problem
set with a default name is created.  Use the --problemset option to use an
existing problem set directory or to create a new problem set.''')
        p.add_option("-w", default=None, action="store_true")
        p.add_option("--debug", default=False, action="store_true",
            help="display additional debugging information.")
        p.add_option("--query", default=None, action="store",
            help="run the remote query, QUERY, and create a " \
                 "new problem set for it")
        p.add_option("--import", default=None,
            action="store", dest="import_file",
            help="import the file IMPORT_FILE")
        p.add_option("--problem-set", default=None,
            action="store", dest="problem_set",
            help="add results to problem set NAME.pysilk in your " \
                 "problem set directory")
        (self.options, self.arguments) = p.parse_args()
        if self.options.query and self.options.import_file:
            print "--query and --import-file are mutually exclusive. " \
                  "Pick one or the other"
            sys.exit()
        if self.options.query and not self.options.problem_set:
            print "please specify a name for the new problem set "\
                  "with --problem-set"
            sys.exit()
        if self.options.query:
            # Parse value given for QUERY into a dictionary
            rArgs = re.compile(r'\s*([\w\-]*)=([\w,\-\/\:]*)')
            arg_list = rArgs.findall(self.options.query)
            print "arg_list = %s" % arg_list
            self.options.query_opts = {}
            for a in arg_list:
                (opt, val) = a
                print "setting %s = %s" % (opt, val)
                self.options.query_opts[opt] = val
            print "Query options = %s " % self.options.query_opts
        print "Running pysilk with the following options"
        print "options.debug = %s" % self.options.debug
        print "options.query = %s" % self.options.query
        print "options.import_file = %s" % self.options.import_file
        print "options.problem_set = %s" % self.options.problem_set
    
    ''' Create the main application object'''
    def __init__(self):
        global _running_ipython
        global _headless
        self._newWindowPos = (50,50)
        self._problemSets = {}
        self.quit = False
        self.SetMacPreferencesMenuItemId(wx.ID_PREFERENCES)
        wx.SystemOptions.SetOptionInt("mac.listctrl.always_use_generic", 0)
        wx.App.__init__(self, redirect=False, filename=None,
            useBestVisual=True)
        # FIXME: Is this the best way to do thsi?
        def _exceptionhook(type, value, tb):
            print "**** UNCAUGHT EXCEPTION *****"
            traceback.print_exception(type, value, tb)
            msg = "".join(traceback.format_exception(type, value, tb))
            entries = traceback.format_list(traceback.extract_tb(tb))
            title = "INTERNAL ERROR:\n%s" % entries[-1]
            dlg = wx.MessageDialog(self.GetTopWindow(), msg, title,
                    wx.ICON_EXCLAMATION)
            dlg.ShowModal()
        sys.excepthook = _exceptionhook
        
        # Handle command-line options
        if not _running_ipython:
            self._processOptions()

        # Initialize configuration information
        config.silk.init()
        config.filteroptions.init()
        config.fgloboptions.init()
        # Not sure what this does
        wx.InitAllImageHandlers()
        
        # Clean up any unsaved problem set files
        self.DeleteOrphans()
        
        # Initialize logging level
        logging.basicConfig(format="%(asctime)s %(levelname)8s %(message)s",
            filename = os.path.join(config.silk.opts["LOCAL_OUTPUT_DIR"],
                            "logfile.txt"),
            filemode='w')
        logging.getLogger().setLevel(logging.DEBUG)
        if not _headless:
            self._init_gui()

    def OnInvoke(self, evt): 
        evt.invoke()
            
    def InvokeLater(self, f, *args, **kwargs):
        print "Adding pending event"
        win = self.hiddenFrame
#        win = self.GetExplorers()[0]
        win.GetEventHandler().AddPendingEvent(InvokeEvent(f, args, kwargs))
        
    def _init_gui(self):
        # Initialize my controls module
        ws.init()
        # Make sure there's a menu displayed when nothing's showing
        menubar = self.CreateMenuBar(minimal=True)
        

        bogus = wx.Frame(None, size=(0,0), style=wx.FRAME_NO_TASKBAR)
        bogus.Connect(-1, -1, wxEVT_INVOKE, self.OnInvoke) 
        bogus.Show()
        self.hiddenFrame = bogus
        menubar.FindItemById(ID_SAVE_PROBLEMSET).Enable(False)
        menubar.FindItemById(ID_SAVEALL).Enable(False)
        # FIXME: no import in initial version
        #menubar.FindItemById(ID_IMPORT_DATASET).Enable(False)
        menubar.FindItemById(ID_CLOSE_WINDOW).Enable(False)
        bogus.SetMenuBar(menubar)

        # FIXME: disable in initial release
        # Displays all jobs (similar to Downloads in web browser)
        #self.dataSetBrowser = JobManager(self)
        # Create the query builder
        self.queryBuilder = QueryBuilder(None, # self.dataSetBrowser,
                                title="Query Builder")
        self.queryBuilder.Show(False)
        #self.SetTopWindow(self.dataSetBrowser)
        
        # Initiate a query if command line option is present
        if not _running_ipython and self.options.query:
            # If it exists already load it up
            dirName = self.options.problem_set + ".pysilk"
            filename= os.path.join(config.silk.opts["LOCAL_OUTPUT_DIR"],
                dirName, "index.xml")
            if os.path.exists(filename):
                # Open existing problem set
                problemSet = ProblemSet.Load(filename)
            else:
                # Create a new problem set
                problemSet = ProblemSet(uniqueId = self.options.problem_set)
            explorer = self.NewExplorer(problemSet)
            tool = QueryTool(QueryResult, problemSet,
                **self.options.query_opts)
            tool.Run(True)
        else:
            # Create an Explorer with an initial empty problem set
            if _running_ipython:
                problemSet = ProblemSet()
            else:
                print "*** Creating ProblemSet"
                problemSet = ProblemSet(createDirectory=False, 
                    uniqueId = self.options.problem_set)
            explorer = self.NewExplorer(problemSet)
        # Handle non-context specific menus
        self.Bind(wx.EVT_MENU, self.OnNew, id=ID_NEW_PROBLEMSET)
        self.Bind(wx.EVT_MENU, self.OnOpen, id=ID_OPEN_PROBLEMSET)
        self.Bind(wx.EVT_MENU, self.OnShowJobs, id=ID_WINDOW_SHOW_DATASETS)
        self.Bind(wx.EVT_MENU, self.OnWebSite, id=ID_HELP_SILK_HOME)
        self.Bind(wx.EVT_MENU, self.OnQuery, id=ID_WINDOW_SHOW_QUERY_BUILDER)
        self.Bind(wx.EVT_MENU, self.OnAbout,id=wx.ID_ABOUT)
        self.Bind(wx.EVT_MENU, self.OnPreferences,id=wx.ID_PREFERENCES)
        self.Bind(wx.EVT_MENU, self.OnQuit, id=wx.ID_EXIT)
    
    def DeleteOrphans(self):
        ''' Delete any unsaved problem set directories. '''
        documentsDir = config.silk.opts["LOCAL_OUTPUT_DIR"]
        print "Deleting orphans from %s" % documentsDir
        for d in os.listdir(documentsDir):
            if not os.path.splitext(d)[1] == ".pysilk": continue
            if not os.path.exists(os.path.join(documentsDir, d, "index.xml")):
                print "Removing %s" % os.path.join(documentsDir, d)
                #os.removedirs(os.path.join(documentsDir, d))
                shutil.rmtree(os.path.join(documentsDir, d))
    
    def GetProblemSets(self):
        return self._problemSets.keys()
    
    def GetExplorers(self):
        return self._problemSets.values()
    
    def GetExplorer(self, result):
        while result.GetParent():
            result = result.GetParent()
        problemSet = result
        if problemSet in self._problemSets:
            return self._problemSets[problemSet]
        else:
            return None
    
    def AddViews(self, result, selectResult=True):
        global _headless
        if _headless: return
        
        # Initialize corresponding user interface elements
        # FIXME: no jobs browser in intiail version
        #self.dataSetBrowser.AddDataSet(result)
        explorer = self.GetExplorer(result)
        # FIXME: create explorer if None above
        if explorer == None:
            explorer = self.NewExplorer(result.GetProblemSet())
        explorer.NewView(result)
        if selectResult:
            explorer.SelectDataSet(result)
        # Bind them to the result
        #result.AddChangeListener(self.dataSetBrowser)
        result.AddChangeListener(explorer)
    
    def NewExplorer(self, problemSet):
        explorer = ProblemSetExplorer(self, problemSet, self._newWindowPos)
        (x,y) = self._newWindowPos
        self._newWindowPos = (x+10, y+10)
        self._problemSets[problemSet] = explorer
        explorer.Show()
        self.queryBuilder.SetAvailableProblemSets(self.GetProblemSets())
        return explorer
    
    def ExplorerClosed(self, explorer):
        del self._problemSets[explorer.GetProblemSet()]
    
    def OnOpen(self, evt):
        dlg = wx.DirDialog(self.GetTopWindow(), "Open a Problem Set",
            defaultPath = config.silk.opts["LOCAL_OUTPUT_DIR"],
            style=wx.DD_DIR_MUST_EXIST)
        result = dlg.ShowModal()
        if result==wx.ID_OK:
            dataDir = dlg.GetPath()
            #lastDir = os.path.split(dataDir)[-1]
            #uniqueId = os.path.splitext(lastDir)[0]
            filename = os.path.join(dataDir, "index.xml")
            if not os.path.exists(filename):
                d = wx.MessageDialog(self.GetTopWindow(),
                  "Not a problem set directory (can't find %s)" % filename,
                  "Can't open directory", wx.ICON_EXCLAMATION)
                d.ShowModal()
                return
            problemSet = ProblemSet.Load(filename)
            wx.GetApp().NewExplorer(problemSet)
        return
    
    def OnNew(self, evt):
        # Create an Explorer with an initial empty problem set
        problemSet = ProblemSet()
        wx.GetApp().NewExplorer(problemSet)
    
    def OnQuery(self, evt):
         wx.GetApp().DoShowQueryBuilder()
    
    def OnQuit(self, evt):
        # FIXME: save any state
        self.DeleteOrphans()
        self.Exit()
    
    def OnShowCommandLog(self, evt):
        self.logWindow.Show(True)
    
    def OnWebSite(self, evt):
        # FIXME: make this non-mac specific (how?!?)
        os.system("open %s" % "http://tools.netsa.cert.org/silk")
    
    def OnPreferences(self, evt):
        d = PreferencesDialog(config.silk.opts,
                self.GetTopWindow(),
                title="Preferences")
        d.CenterOnScreen()
        result = d.ShowModal()
        if result == wx.ID_OK:
            d.SavePreferences()
    
    def OnShowJobs(self, evt):
        self.DoShowDataSets()
    
    def OnAbout(self, evt):
        dlg = AboutBox(None)
        dlg.ShowModal()
        dlg.Destroy()
    
    '''Add the data set to the browser and launch a query'''
    def StartQuery(self, problemSet, **opts):
        # Create the tool
        tool = QueryTool(QueryResult, problemSet, **opts)
        # Create UI elements
        self.AddViews(tool.result)
        # Start the analsis thread
        tool.Run(True)
    
    def AddDataSet(self, dataset):
        pass
        # FIXME: no job browser for now
        #self.dataSetBrowser.AddDataSet(dataset)
    
    # Show the data sets browser
    def DoShowDataSets(self):
        pass
        # FIXME: no job browser for now       
        #self.dataSetBrowser.Show()
        #self.dataSetBrowser.Raise()
    
    # Explore the selected data set
    def DoShowExplorer(self):
        self.explorer.Show(show)
        self.explorer.Raise()
        
    def GetSelectedFilterData(self):
        d = self.GetSelectedData()
        while d != None and not d.IsFilterResult():
            d = d.GetParent()
        if d.IsFilterResult():
            return d
        else:
            raise RuntimeException("INTERNAL ERROR: Couldn't find binary parent")
        
    def GetSelectedData(self):
        return ProblemSetExplorer.GetLastActive().GetSelection()
    
    def GetActiveProblemSet(self):
        return ProblemSetExplorer.GetLastActive().GetProblemSet()
    
    # Show the query bouilder
    def DoShowQueryBuilder(self, opts=None, show=True):
        self.queryBuilder.Show(show)
        self.queryBuilder.Raise()
        self.queryBuilder.SelectProblemSet(self.GetActiveProblemSet())
        if opts:
            self.queryBuilder.SetFilterOptions(opts)
        
    def GetViewMenu(self):
        return self.menuView
    
    def CreateMenuBar(self, minimal=False):
        # Populate the menu bar -- use the same menu bar for the application
        menuBar = wx.MenuBar()
        menuFile = wx.Menu()
        menuFile.Append(ID_NEW_PROBLEMSET, "&New\tCTRL-N",
            "New problem set")
        menuFile.Append(ID_OPEN_PROBLEMSET, "&Open...\tCTRL-O",
            "Open problem set")
        menuFile.Append(ID_SAVE_PROBLEMSET, "&Save\tCTRL-S",
            "Save problem set")
        # FIXME: disabling for intial remote-only release
        # menuFile.Append(ID_CLEANUP, "Clean up",
        #    "Delete orphaned files in the problem set directory")
        #menuFile.AppendSeparator()
        #menuFile.Append(ID_IMPORT_DATASET,
        #    "&Import Data File...\tCTRL-SHIFT-I", "Import binary data file")
        #menuFile.AppendSeparator()
        #menuFile.Append(ID_EXPORT_AS_SCRIPT, "&Export As Script...",
        #    "Export selection as reusable script")
        #menuFile.Append(ID_EXPORT_AS_TOOL, "&Export As Plug-in Tool...",
        #    "Export selection as a tool")
        #menuFile.AppendSeparator()
        menuFile.Append(ID_CLOSE_WINDOW, "Close\tCTRL-W",
            "Close problem set")
        menuFile.Append(ID_SAVEALL, "Save All", "Save all problem sets")
        menuFile.Append(wx.ID_PREFERENCES, "&Preferences...")
        menuBar.Append(menuFile, "File")
        if not minimal:
            menuAnalysis = wx.Menu()
            menuAnalysis.Append(ID_ANALYSIS_REFINE, "rwfilter\tCTRL-SHIFT-F")
            menuAnalysis.Append(ID_ANALYSIS_COUNT,  "rwuniq\tCTRL-SHIFT-U")
            menuAnalysis.Append(ID_ANALYSIS_TREND,  "rwcount\tCTRL-SHIFT-C")
            menuAnalysis.Append(ID_ANALYSIS_SET,  "rwset\tCTRL-SHIFT-S")            
            #menuAnalysis.AppendSeparator()
            #menuAnalysis.Append(ID_ANALYSIS_SPLIT, "Split")
            menuBar.Append(menuAnalysis, "Tools")
            menuGraph = wx.Menu()
            menuGraph.Append(ID_ANALYSIS_SMART_GRAPH,
                "Quick Graph\tCTRL-SHIFT-G", "Pick a graph based on the data")
            menuBar.Append(menuGraph, "Graph")
            menuView = wx.Menu()
            menuView.Append(ID_WINDOW_SHOW_QUERY_BUILDER,
                "Query &Builder\tCTRL-SHIFT-B","Show Query Builder")
# FIXME: no browser in initial version
#            menuView.Append(ID_WINDOW_SHOW_DATASETS,
#                "Anaysis &Jobs\tCTRL-J", "Show Analysis Jobs")
            menuView.AppendSeparator()
            menuView.Append(ID_SHOW_INFO, "Show Info\tCTRL-I")
            menuView.AppendSeparator()
            menuView.Append(ID_VIEW_OPTIONS, "View Options...",
                "Customize selected view.")
            menuBar.Append(menuView, "View")            
        menuHelp = wx.Menu()
        menuHelp.Append(wx.ID_ABOUT, "&About SiLK", "")
        menuHelp.Append(ID_HELP_SILK_HOME,
            "Visit http://tools.netsa.cert.org/silk", "")
        menuBar.Append(menuHelp, "&Help")
        return menuBar
    
    def _DoShowError(self, win, msg, title):
        # FIMXE: no browser in initial version
        d = wx.MessageDialog(None, #self.dataSetBrowser, 
                            msg, title,
                            wx.OK | wx.ICON_EXCLAMATION)
        d.ShowModal()
    
    def ShowError(self, title, msg):
        global _headless
        if _headless: return
        print "Calling invoke later"
        self.InvokeLater(self._DoShowError, self, msg, title)

    def refine(self, parent=None, showDialog=False, **kargs):
        if parent == None:
            parent = wx.GetApp().GetSelectedData()
        return analysistool.RunTool(config.silk.workRemotely(),
                    tools.refine.RefineTool, parent, **kargs)

def main(ipython=False, headless=False):
    global _running_ipython
    global _headless
    print "Starting up"
    
    _running_ipython = ipython
    _headless = headless
    # Create application class
    app = SilkApp()
    app.MainLoop()
    return app

if __name__ == "__main__":
    main()
