## ------------------------------------------------------------------------
## 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 logging
import threading
import pdb
import os, sys, string, tempfile, traceback, re, shutil, glob
import textwrap
import analysisresult
from constants import *
from analysisresult import ProblemSet, ResultRunContext
from tools.query import QueryTool
from resultset import QueryResult
from querybuilder import QueryBuilder
from explorer import ProblemSetExplorer
import analysistool, tools
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.viewoptions
from infoframe import InfoFrame
import library
import silk
from subprocess import Popen
import traceback
from tools import RegisterPlugin
from tools.plugins import CreateGenericPlugin, LoadGenericPlugin
import paramiko
import wx, ws
import wx.lib.mixins.listctrl as listmix

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)

_running_ipython = False
DISABLE_SSH_TEST = False

# Used for initial ssh testing.
# FIXME: This class probably belongs somewhere else
class SSHTestDialog(wx.ProgressDialog):
    def __init__(self, context):
        if not context:
            self.context = silk.SimpleRunContext(runRemote = True)
        else:
            self.context = context
        wx.ProgressDialog.__init__(self, "Testing ssh settings",
            "Running test commands",
            style = wx.STAY_ON_TOP|wx.PD_APP_MODAL|wx.PD_AUTO_HIDE)
        self.SetSize((320,-1))
        self.progress = 0
        self.message = ""
        self.error = None
    
    def GetError(self):
        return self.error
    
    def RunTests(self):
        class _testthread (threading.Thread):
           def __init__(self, dialog):
               threading.Thread.__init__(self)
               self.dialog = dialog
           def run(self):
               self.dialog._run_tests()
        self._thread = _testthread(self)
        self._thread.start()
        self.Show()
        self.MakeModal()
        while self._thread:
            self._thread.join(0.5)
            if self.error != None:
                self.progress = 100
                self.Update(self.progress, self.message)
                raise self.error
            elif self.progress == 100:
                self.Update(self.progress, self.message)
                return
            else:
                self.Update(self.progress, self.message)
    
    def _run_test(self, i, test_count, cmdLine, test):
        context = self.context
        print "Running test %d: %s" % (i,cmdLine)
        self.progress = 100*(i-1)/test_count
        self.message = "Running %s" % (cmdLine)
        if config.silk.useExternalSSH():
            cmdLine = _build_ssh_command(cmdLine)
            f = silk._local_exec_command(context, cmdLine)
        else:
            f = silk._paramiko_exec_command(context, cmdLine,
                combineStderr=True)
        # Consume the output and grab the exit code
        output_str = ""
        while True:
            l = f.readline()            
            output_str += l
            if not l:
                break
        exit_code = f.close()
        if exit_code != 0:
            raise silk.CommandError("SSH Test Failed",
                "Error %s:\n\n%s\n%s" % (test, cmdLine, output_str))
        self.progress = 100*(i)/test_count
        self.message = "Ran %s" % (cmdLine)
        return True
    
    def _run_tests(self):
        try:
            cmdLines = []
            binDir = config.silk.opts["REMOTE_BIN"]
            binDir = binDir.replace(" ", "")
            if binDir == "":
                binDir = None
            if binDir:
                verify_silk = remote.join(binDir, "rwfilter --help")
            else:
                verify_silk = "rwfilter --help"
            verify_output_dir = "ls -l %s" % config.silk.opts["REMOTE_OUTPUT_DIR"]
            verify_set_path = "ls -l %s" % config.silk.opts["REMOTE_LIBRARY_PATH"].replace(":", " ")
            self._run_test(1, 3, verify_silk, "running SiLK command")
            self._run_test(2, 3, verify_output_dir, "listing output directory")
            self._run_test(3, 3, verify_set_path, "listing set directories")
            self.progress = 100
            self.message = "Tests complete"
        except silk.CommandError, e:
            print "ssh test failed on command '%s':" % e.details
            self.error = e
            self.message = "SSH error during testing"
            return


class SilkApp(wx.App):
    def FilterEvent(self, event):
        print "FilterEvent: %s" % event

    def GetVersion(self):
        return ISILK_APP_VERSION

    # Import all of the python modules in the given directory
    def _loadplugins(self, directory):
        # Must explicitly add the plugin directory to the path
        sys.path.insert(0, directory)
        files = glob.glob(os.path.join(directory, "*.py"))
        for path in files:
            (d, f) = os.path.split(path)
            modulename = f.replace(".py", "")
            if modulename != "__init__":
                __import__(modulename)
        files = glob.glob(os.path.join(directory, "*.plugin"))        
        for path in files:
            cls = LoadGenericPlugin(path)
            print "Registering plugin '%s'(%s)" % (cls._GetToolName(), type(cls._GetToolName()))
            RegisterPlugin(cls._GetToolName(), "icons/tool.gif",
                cls, cls.GetCategory())
    
    # Import plugins
    def LoadPluginTools(self):
        # Load plugins installed in <INSTALL_DIR>/config/plugins
        self._loadplugins(os.path.join(os.getcwd(), "config", "plugins"))
        # Load plugins installed in <PREFSDIR>/plugins
        self._loadplugins(os.path.join(config.getPreferencesDir(),
            "plugins"))

    def _processOptions(self):
        # Parse options
        p = optparse.OptionParser()
        p.add_option("--debug", default=False, action="store_true",
            help="display additional debugging information.")
        p.add_option("--open", default=None, action="store", metavar="LOCAL_DIR",
            help="open the given problem set directory")
        p.add_option("--download", default=False, action="store_true",
            help="download the result of query immediately")
        p.add_option("--import-files", default=None, action="store",
            dest="import_files", metavar="REMOTE_FILE_LIST",
            help="import one or more remote binary files into the problem set, file names separated by commas")
        p.add_option("--result-names", default="Untitled Result", action="store", metavar="NAME",
            help="Name to give to result(s) (query result or imported files)")            

        rwfilter_opts =  ["class", "type", "sensors", "start-date", "end-date",
                          "saddress", "daddress", "any-address", "dport", "sport", 
                         "aport", "proto", "stime", "etime", "dur", "flags-all", 
                         "bytes", "packets", "icmp-type", "icmp-code",
                         "dipset", "sipset", "sval", "dval"]
        g = optparse.OptionGroup(p, "rwfilter pass-through options", 
'''To initiate a SiLK query, use one or more of the rwfilter
pass through options.  By default a new query will be initiated
on the analysis server but will not be downloaded.  If you wish
download it immediately use the --download option.''')

        for rwfopt in rwfilter_opts:
            opt = "--" + rwfopt
            g.add_option(opt, action="store")
        p.add_option_group(g)
        (self.options, self.arguments) = p.parse_args()
        # Process the rwfilter options and set them aside
        self.rwfilter_options = {}
        for rwfopt in rwfilter_opts:
            rwfopt = rwfopt.replace("-", "_")
            if self.options.__dict__[rwfopt] != None:
                self.rwfilter_options[rwfopt] = self.options.__dict__[rwfopt]
        print "rwfilter_options = %s" % self.rwfilter_options
            
        print "Running isilk with the following options"
        print "options.debug = %s" % self.options.debug
        print "options.download = %s" % self.options.download
        print "options.open = %s" % self.options.open
    
    ''' Create the main application object'''
    def __init__(self):
        global _running_ipython
        global _headless
        self.library = None
        self._newWindowPos = wx.DefaultPosition # (0,25)
        self._problemSets = {}
        self.quit = False
        self._initialized = False
        
        wx.SystemOptions.SetOptionInt("mac.listctrl.always_use_generic", 1)
        wx.SystemOptions.SetOptionInt("msw.window.no-clip-children", 1)
        wx.App.__init__(self, redirect=False, filename=None,
            useBestVisual=True)

        # FIXME: Is this the best way to do thsi?
        def _exceptionhook(type, value, tb):
            msg = "".join(traceback.format_exception(type, value, tb))
            entries = traceback.format_list(traceback.extract_tb(tb))
            title = "INTERNAL ERROR"
            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 (non)
        self.config_file_exists = config.silk.init()
        config.filteroptions.init()
        config.fgloboptions.init()
        config.viewoptions.init()
        
        # Create the isilk directory if it doesn't exist
        documentsDir = config.silk.opts["LOCAL_OUTPUT_DIR"]
        if not os.path.exists(documentsDir):
            print "Creating documents directory: %s" % documentsDir
            os.mkdir(documentsDir)
        else:
            # Clean up any unsaved problem set files
            self.DeleteOrphans()
        
        # Not sure what this does
        wx.InitAllImageHandlers()
        
        # Load plugin modules
        self.LoadPluginTools()
        
        # 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')
        # Make sure there's a configuration file before doing 
        # anything else.  Force the user to create a good configuration.
        firstTime = False
        if not self.config_file_exists:
            firstTime = True
            self.ConfigureSSH(None, firstTime=True, 
                runTest=False, lastTestSucceeded=False)

        # Initialize the user interface
        # Create a new problem set if appropriate
        if not _headless:
            print "Initializing user interface"
            problemSet = self._init_gui()

        # Verify SSH configuration for the problem set is correct.
        # FIXME: allow offline use?
        if not firstTime:
            test_succeeded = True
            error_message = ""
            if self.config_file_exists:
                (test_succeeded, error_message) = self.TestSSH()
                if not test_succeeded:
                      d = wx.MessageDialog(None, error_message, "SSH Test Failed",
                         wx.OK|wx.ICON_EXCLAMATION)
                      d.CenterOnScreen()
                      d.ShowModal()
            if not self.config_file_exists or not test_succeeded:
                 print "Configuring"
                 self.ConfigureSSH(problemSet,
                     not self.config_file_exists,
                     runTest=False, lastTestSucceeded=test_succeeded,
                     error_message = error_message)
        
        if len(self.rwfilter_options) > 0:
            self._run_initial_query(problemSet)
            
        if self.options.import_files:
            self._run_initial_import(problemSet)
        
        # If we're running in py2exe on Windows, 
        # redirect stdout and stderr
        redirect_streams = False            
        if os.path.split(sys.executable)[-1] == "isilk.exe":
            redirect_streams = True
        if redirect_streams:
            prefsDir = wx.StandardPaths.Get().GetDocumentsDir()
            print "prefsDir = %s" % prefsDir
            log_filename = os.path.join(prefsDir, "isilk.log")
            print "Redirecting stderr and stdout to %s" % log_filename
            sys.stdout = sys.stderr = open(log_filename, "w")
            
        self._initialized = True
        
    def IsInitialized(self):
        return self._initialized
            
    def _run_initial_import(self, problemSet):
        results = self.options.import_files.split(",")
        names = self.options.result_names.split(",")
        for (i, filename) in enumerate(results):
            try:
                name = names[i]
            except:
                name = "Untitled Imported Result"
            q = analysistool.ImportFile(problemSet, filename, name, QueryResult)
            if self.options.download:
                q.DownloadFile()
    
    def _run_initial_query(self, problemSet):
        result_names = self.options.result_names.split(",")
        try:
            name = result_names[0]
        except:
            name = None
        if not name:
            name = "Untitled Query Result"
        q = analysistool.RunTool(True, False, tools.query.QueryTool, 
                problemSet, name=name,
                **self.rwfilter_options)
        print "Created initial query: %s" % q
        if self.options.download:
            q.DownloadFile()
    
    def ConfigureSSH(self, context, firstTime=False, runTest=True,
        lastTestSucceeded=True, error_message=""):
        print "Displaying ssh configuration dialog..."
        test_succeeded = lastTestSucceeded
        #if not runTest:
        #    test_succeeded = True # assume last test worked
        #else:
        #    test_succeeded = False
        # Run a test if we've already configured
        if self.config_file_exists and runTest:
            (test_succeeded, error_message) = self.TestSSH()
        print "Not running test"
        if context:
            print "Editing ssh options for %s"
            print "%s" % context.GetSSHOptions()
            self.sshOptions = context.GetSSHOptions()
        else:
            self.sshOptions = config.silk.opts

        while True:
            # If the configuration file doesn't exist
            # assume user is running for the first time
            if not test_succeeded and not firstTime:
                print "**** Creating preferences dialog"
                d = PreferencesDialog(self.sshOptions,
                        self.GetTopWindow(),
                        requireEdits=True,
                        headingText = "SSH Settings - SSH Test Failed",
                        helpText = error_message,
                        testFailed=True,
                        title="SSH Settings")
            elif firstTime:
                d = PreferencesDialog(self.sshOptions,
                        self.GetTopWindow(),
                        requireEdits=True,
                        headingText = "Welcome to iSiLK",
                        helpText = (
                            "Before you use iSiLK for the first time, "
                            "you need to set the SSH configuration for "
                            "accesssing the remote analysis system.  "
                            "Future problem sets will use this server "
                            "for analysis and storage."),
                        title="SSH Settings",
                        defaultsOnly=True)
            else:
                d = PreferencesDialog(self.sshOptions,
                        self.GetTopWindow(),
                        requireEdits=False,
                        headingText = "SSH Settings",
                        title="SSH Settings")
            d.CenterOnScreen()
            print "Calling ShowModal"
            result = d.ShowModal()
            print "Returned from ShowModal"
            if result == wx.ID_OK:
                d.SavePreferences()
                (test_succeeded, error_message) = self.TestSSH(context)
                if test_succeeded:
                    if context:
                        self.GetExplorer(context).RefreshContents()
                    return
                else:
                    d = wx.MessageDialog(None, error_message,
                       "SSH Test Failed",
                       wx.OK|wx.ICON_EXCLAMATION)
                    d.CenterOnScreen()
                    d.ShowModal()
            else:
                if not test_succeeded or firstTime:
                    self.quit = True
                return
    
    def TestSSH(self, context=None, showProgressDialog=True):
        if DISABLE_SSH_TEST:
            return (True, "")
        # Verify that ssh works properly
        dlg = SSHTestDialog(context)
        try:
            print "Running ssh test"
            dlg.RunTests()
            print "ssh test succeeded"
            dlg.Destroy()
            return (True, "Test succeeded")
        except silk.CommandError, e:
            print "Caught exception"
            print e
            print "ssh test failed: %s" % e.details
            dlg.Destroy()
            return (False, e.details)
    
    def DidQuit(self):
        return self.quit
    
    def OpenExcel(self, file):
        if wx.Platform == "__WXMAC__":
            os.system("open -a Excel %s" % file)
        elif wx.Platform == "__WXGTK__":
            os.startfile(file)
        elif wx.Platform == "__WXMSW__":
            os.startfile(file)
    
    def ShowFile(self, file):
        # FIXME: For other platforms...
        if wx.Platform == "__WXMAC__":
            os.system("open %s" % file)
        elif wx.Platform == "__WXGTK__":
            os.system("gnome-open %s" % file)
        elif wx.Platform == "__WXMSW__":
            print "Running explorer: %s " % file
            os.startfile(file)
    
    def OnInvoke(self, evt):
        evt.invoke()
    
    def InvokeLater(self, f, *args, **kwargs):
        win = self.hiddenFrame
        win.GetEventHandler().AddPendingEvent(InvokeEvent(f, args, kwargs))
        print "Added pending event"
    
    def _init_gui(self):
    
        # Initialize my controls module
        ws.init()
        
        # FIXME: disable in initial release
        # Displays all jobs (similar to Downloads in web browser)
        # Create the query builder
        self.queryBuilder = QueryBuilder(None, # self.dataSetBrowser,
            title="Query Builder",
            style=wx.CAPTION|wx.CLOSE_BOX|wx.SYSTEM_MENU)
        self.queryBuilder.Show(False)
        
        # Create the problem set and an explorer for it
        if self.options.open:
            problemSet = ProblemSet.Load(self.options.open)
        else:
            problemSet = ProblemSet()
        explorer = self.NewExplorer(problemSet)

        # Create a hidden window so there's always a menu bar (on Mac)
        # and we always have a window we can use to process
        # any custom events we create (e.g., InvokeLater).
        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_RESULT).Enable(False)
        menubar.FindItemById(ID_CLOSE_WINDOW).Enable(False)
        bogus.SetMenuBar(menubar)
        explorer.Raise()
        explorer.SetFocus()

        # FIXME: add ipython support back!

        # 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.OnSaveAll, id=ID_SAVEALL)
        self.Bind(wx.EVT_MENU, self.OnShowJobs, id=ID_WINDOW_SHOW_DATASETS)
        self.Bind(wx.EVT_MENU, self.OnSilkHome, id=ID_HELP_SILK_HOME)
        self.Bind(wx.EVT_MENU, self.OnISilkHome, id=ID_HELP_ISILK_HOME)        
        self.Bind(wx.EVT_MENU, self.OnHelp, id=ID_HELP_RWFILTER)
        self.Bind(wx.EVT_MENU, self.OnHelp, id=ID_HELP_RWUNIQ)
        self.Bind(wx.EVT_MENU, self.OnHelp, id=ID_HELP_RWCOUNT)
        self.Bind(wx.EVT_MENU, self.OnHelp, id=ID_HELP_RWSET)
        self.Bind(wx.EVT_MENU, self.OnHelp, id=ID_HELP_RWSETBUILD)        
        self.Bind(wx.EVT_MENU, self.OnHelp, id=ID_HELP_RWSETTOOL)
        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.OnSSHSettings,id=ID_SSH_SETTINGS)
        self.Bind(wx.EVT_MENU, self.OnQuit, id=wx.ID_EXIT)
        
        return problemSet

    
    def DeleteOrphans(self):
        ''' Delete any unsaved problem set directories. '''
        documentsDir = config.silk.opts["LOCAL_OUTPUT_DIR"]
        if not os.path.exists(documentsDir):
            return False
        print "Deleting orphans from %s" % documentsDir
        deleted = False
        for d in os.listdir(documentsDir):
            if not os.path.splitext(d)[1] == ".isilk": continue
            if not os.path.exists(os.path.join(documentsDir, d, "index.xml")):
                print "Removing %s" % os.path.join(documentsDir, d)
                shutil.rmtree(os.path.join(documentsDir, d))
                deleted = True
        return deleted
    
    def GetProblemSets(self):
        return self._problemSets.keys()
    
    def GetExplorers(self):
        return self._problemSets.values()
        
    def GetActiveQueries(self):
        queries = []
        for e in self.GetExplorers():
            q = e.GetActiveQuery()
            queries.append(q)
        return queries
    
    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()]
        # FIXME: why is this necessary!?!?
        if len(self._problemSets) == 0:
            print "Quitting"
            self.DoQuit()
    
    def OnOpen(self, evt):
        values = []
        directories = []
        for f in os.listdir(config.silk.opts["LOCAL_OUTPUT_DIR"]):
            print "Looking at file %s" % f
            if os.path.isdir(os.path.join(config.silk.opts["LOCAL_OUTPUT_DIR"], f)):
                print "Adding directory %s to list" % f
                try:
                    name = analysisresult._extractNameFromFileName(f)
                except:
                    name = "No name"
                directories.append(f)
                values.append("%s (%s)" % (name, f))
                
                
        dlg = ProblemSetPickerDialog(self.GetTopWindow())    
        #dlg = wx.SingleChoiceDialog(self.GetTopWindow(), "Choose a problem set",
        #    "Open a problem set", values)
        
#        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:
            print "Selected item %s - %s" % (dlg.GetSelectedName(), dlg.GetSelectedDir())
            #dirName = directories[dlg.GetSelection()]
            dirName = dlg.GetSelectedDir()
            print "dirName = %s" % dirName
            dataDir = os.path.join(config.silk.opts["LOCAL_OUTPUT_DIR"], dirName)
 #           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\n(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 OnSaveAll(self, evt):
        for p in self.GetProblemSets():
            p.Save()
    
    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 GetLibrary(self):
        return library.GetLibrary()
    
    def OnQuit(self, evt):
        print "OnQuit"
        self.DoQuit()
    
    def DoQuit(self):
        print "In DoQuit"
        canClose = True
        for e in self.GetExplorers():
            if not e.CanClose():
                canClose = False
                break
        if not canClose:
            d = wx.MessageDialog(None,
                "There are tools currently running or other unsaved changes.",
                 "Are you sure you want to exit?",
                 wx.YES_NO | wx.ICON_EXCLAMATION)
            result = d.ShowModal()
            if result != wx.ID_YES:
                print "NOT EXITING"
                return
        self.DeleteOrphans()
        self.Exit()
    
    def OnShowCommandLog(self, evt):
        self.logWindow.Show(True)
        
    def LaunchBrowser(self, url):
        if wx.Platform == "__WXMAC__":
            os.system("open %s" % url)
        elif wx.Platform == "__WXGTK__":
            os.startfile(url)
        elif wx.Platform == "__WXMSW__":
            os.startfile(url)
    
    def OnSilkHome(self, evt):
        # FIXME: make this non-mac specific (how?!?)
        self.LaunchBrowser("http://tools.netsa.cert.org/silk")

    def OnISilkHome(self, evt):
        # FIXME: make this non-mac specific (how?!?)
        self.LaunchBrowser("http://tools.netsa.cert.org/isilk")
        
    def OnHelp(self, evt):
        tool = None
        if evt.GetId() == ID_HELP_RWFILTER:
            tool = "rwfilter"
        elif evt.GetId() == ID_HELP_RWUNIQ:
            tool = "rwuniq"
        elif evt.GetId() == ID_HELP_RWCOUNT:
            tool = "rwcount"
        elif evt.GetId() == ID_HELP_RWSET:
            tool = "rwset"
        elif evt.GetId() == ID_HELP_RWSETBUILD:
            tool = "rwsetbuild"
        elif evt.GetId() == ID_HELP_RWSETTOOL:            
            tool = "rwsettool"
        self.LaunchBrowser("http://tools.netsa.cert.org/silk/%s.html" % tool)
    
    def OnSSHSettings(self, evt):
        self.ConfigureSSH(self.GetActiveProblemSet(),
            runTest=False)
        if self.quit:
            self.DoQuit()
    
    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()
    
    def GetActiveExplorer(self):
        result = self.GetActiveProblemSet()
        return self.GetExplorer(result)
    
    # Show the query bouilder
    def DoShowQueryBuilder(self, opts=None, show=True):
        print "In DoShowQueryBuilder"
        # FIXME: this shouldn't be necessary but for some strange
        # reason this dialog is disabled sometimes!
        self.queryBuilder.Enable(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 RefreshTools(self):
        for e in self.GetExplorers():
            e.RefreshTools()
    
    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_CLOSE_WINDOW, "Close\tCTRL-W",
            "Close problem set")
        menuFile.Append(ID_SAVE_PROBLEMSET, "&Save\tCTRL-S",
            "Save problem set")
        menuFile.Append(ID_SAVEALL, "Save All",
            "Save all problem sets")
        menuFile.AppendSeparator()
        menuFile.Append(ID_SAVE_AS, "&Rename Problem Set",
            "Rename problem set directory")            
        menuFile.AppendSeparator()        
        menuFile.Append(ID_IMPORT_RESULT,
            "&Import Remote Data File...\tCTRL-SHIFT-I",
            "Import binary data file")
        menuFile.AppendSeparator()       
        menuFile.Append(ID_SHOW_FILES,
            "Show Local Files", "Show local files on desktop")
        menuFile.AppendSeparator()     
        menuFile.Append(wx.ID_EXIT, "Exit")
        menuBar.Append(menuFile, "File")
        if not minimal:
            menuEdit = wx.Menu()
            menuEdit.Append(ID_EDIT_COPY_RECS_AS_TEXT, "Copy as tabular data")
            menuEdit.Append(ID_EDIT_COPY_RECS_AS_LATEX, "Copy as LaTex table")
            menuBar.Append(menuEdit, "Edit")
            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")
            
            # Put plug-ins in the menu
            def _addmenu(category):
                pluginMenu = wx.Menu()
                if category == "Query":
                    pluginMenu.Append(ID_WINDOW_SHOW_QUERY_BUILDER,
                        "Query &Builder\tCTRL-SHIFT-B","Show Query Builder")
                elif category == "Filter":
                    pluginMenu.Append(ID_ANALYSIS_REFINE, "rwfilter...\tCTRL-SHIFT-F")
                elif category == "Uniq":
                    pluginMenu.Append(ID_ANALYSIS_COUNT,  "rwuniq...\tCTRL-SHIFT-U")  
                elif category == "Count":
                    pluginMenu.Append(ID_ANALYSIS_TREND,  "rwcount...\tCTRL-SHIFT-C")                                      
                elif category == "Set":
                    pluginMenu.Append(ID_ANALYSIS_SET,  "rwset...")
                    pluginMenu.Append(ID_ANALYSIS_SET_TOOL,  "rwsettool...")
                    pluginMenu.Append(ID_ANALYSIS_SET_BUILD,  "rwsetbuild...")
                if category != "Analysis":
                    pluginMenu.AppendSeparator()
                
                if len(tools.GetPlugins(category)) == 0:
                    pluginMenu.Append(ID_NO_COMMAND, "(No plug-ins defined)")
                else:
                    for (i, plugin) in enumerate(tools.GetPlugins(category)):
                        print "plugin.toolName = '%s'(%s)" % (plugin.toolName,type(plugin.toolName))
                        shortcut = ""
                        pluginMenu.Append(ID_ANALYSIS_CUSTOM+plugin.toolId,
                            plugin.toolName + shortcut)
                if category == "Analysis":
                    menuAnalysis.AppendSubMenu(pluginMenu, "Other analysis")
                else:    
                    menuAnalysis.AppendSubMenu(pluginMenu, category)

            _addmenu("Query")
            _addmenu("Filter")
            _addmenu("Uniq")
            _addmenu("Count")
            _addmenu("Set")
            _addmenu("Analysis")
#            menuAnalysis.AppendSeparator()
# FIXME: re-enable when recipe features fixed
            menuAnalysis.Append(ID_ANALYSIS_SAVE_IN_LIBRARY,
                "Create recipe based on result\tCTRL-SHIFT-R")
            menuAnalysis.AppendSeparator()                
            menuAnalysis.Append(ID_ANALYSIS_MANAGE_SET_LIBRARY,
                "Manage library...")
            menuAnalysis.AppendSeparator()
            menuAnalysis.Append(ID_ANALYSIS_SHELL_TOOL,
                "Run shell command...")
            menuAnalysis.Append(ID_ANALYSIS_MANAGE_CUSTOM_TOOLS,
                "Manage plug-ins...\tCTRL-SHIFT-M")
            menuAnalysis.AppendSeparator()                
            menuAnalysis.Append(ID_SSH_SETTINGS,
                "SSH Settings..\tCTRL-SHIFT-A")
            
            #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")
            menuView.AppendSeparator()
# FIXME: re-enable when recipes work
            menuView.AppendCheckItem(ID_VIEW_SHOW_LIBRARY,
                "Recipe Library\tCTRL-SHIFT-L", "Show Analysis Library")
            menuView.AppendCheckItem(ID_VIEW_SHOW_COMMAND_LOG,
                "Command Log", "Show Command Log")

            menuView.AppendSeparator()
            menuView.Append(ID_SHOW_INFO, "Show File 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 iSiLK", "")
        menuHelp.Append(ID_HELP_ISILK_HOME,
            "Visit http://tools.netsa.cert.org/isilk", "")
        menuHelp.Append(ID_HELP_SILK_HOME,
            "Visit http://tools.netsa.cert.org/silk", "")
            
#        menuHelp.Append(ID_HELP_RWFILTER, "rwfilter manual")
        menuManuals = menuHelp# wx.Menu()
        menuHelp.AppendSeparator()
        menuManuals.Append(ID_HELP_RWFILTER, "rwfilter manual")
        menuManuals.Append(ID_HELP_RWUNIQ, "rwuniq manual")
        menuManuals.Append(ID_HELP_RWCOUNT, "rwcount manual")
        menuManuals.Append(ID_HELP_RWSET, "rwset manual")
        menuManuals.Append(ID_HELP_RWSETBUILD, "rwsetbuild manual")
        menuManuals.Append(ID_HELP_RWSETTOOL, "rwsettool manual")
#        menuHelp.AppendSubMenu(menuManuals, "SiLK Manuals")
        menuBar.Append(menuHelp, "&Help")
        return menuBar
    
    def _DoShowError(self, win, msg, title):
        # FIMXE: no browser in initial version
        d = wx.MessageDialog(None, msg, title, wx.OK|wx.ICON_EXCLAMATION)
        d.ShowModal()
    
    def ShowError(self, title, msg):
        global _headless
        if _headless: return
        #pdb.set_trace()
        self.InvokeLater(self._DoShowError, self, msg, title)
    
    def refine(self, parent=None, showDialog=False, name=None, **kargs):
        if parent == None:
            parent = wx.GetApp().GetSelectedData()
        result = analysistool.RunTool(config.silk.workRemotely(), False,
                        tools.filter.FilterTool, parent,
                        name = name,
                        **kargs)
                        
class ProblemSetListCtrl(wx.ListCtrl, listmix.ListCtrlAutoWidthMixin):
    def __init__(self, parent):
        wx.ListCtrl.__init__(self, parent, 
                 style = wx.LC_SINGLE_SEL|wx.LC_REPORT|wx.NO_BORDER|wx.LC_HRULES)
        self.InsertColumn(0, "Name", width=150)
        self.InsertColumn(1, "Directory", width=150)

class ProblemSetPickerDialog(wx.Dialog):
    def __init__(self, parent, *args, **kargs):
        wx.Dialog.__init__(self, parent, -1,
            title="Open a problem set", *args, **kargs)
        self.values = []
        self.directories = []
        self.list = ProblemSetListCtrl(self)
        sz = wx.GridBagSizer()
        sz.SetEmptyCellSize((0,0))
        szok = self.CreateButtonSizer(wx.OK|wx.CANCEL)
        sz.Add(self.list, (0,0), (1,2),
            flag=wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, border=8)
        sz.Add(szok, (3,0), (1,2),
            flag=wx.EXPAND|wx.BOTTOM|wx.RIGHT|wx.TOP, border=8)
        sz.AddGrowableCol(0)
        sz.AddGrowableRow(0)
        self.SetSize((-1,480))
        self.SetSizer(sz)            
        self.FindWindowById(wx.ID_OK).SetDefault()
        self.FindWindowById(wx.ID_OK).SetLabel("Open")
        if self.list.GetSelectedItemCount() == 0:
            self.FindWindowById(wx.ID_OK).Enable(False)
        self.list.Bind(wx.EVT_LIST_ITEM_SELECTED, self.OnItemSelected)
        self.list.Bind(wx.EVT_LIST_ITEM_DESELECTED, self.OnItemDeSelected)
        self.list.Bind(wx.EVT_LIST_ITEM_ACTIVATED, self.OnDoubleClick)
        self._populateList()
        
    def _populateList(self):
        for f in os.listdir(config.silk.opts["LOCAL_OUTPUT_DIR"]):
            print "Looking at file %s" % f
            if os.path.isdir(os.path.join(config.silk.opts["LOCAL_OUTPUT_DIR"], f)):
                (name, ext) = os.path.splitext(f)
                if ext != ".isilk":
                    continue
                print "Adding directory %s to list" % f
                try:
                    name = analysisresult._extractNameFromFileName(f)
                except:
                    name = "No name"
                self.directories.append(f)
                self.values.append("%s (%s)" % (name, f))
                index = self.list.InsertStringItem(sys.maxint, name)
                self.list.SetStringItem(index, 1, f)
        # Very ugly
        listmix.ListCtrlAutoWidthMixin.__init__(self.list)
        
    def OnItemSelected(self, event):
        self._selectionIndex = event.m_itemIndex
        self.FindWindowById(wx.ID_OK).Enable(True)
        
    def OnDoubleClick(self, event):
        self._selectionIndex = event.m_itemIndex
        self.EndModal(wx.ID_OK)
        
    def OnItemDeSelected(self, event):
        self._selectionIndex = None
        self.FindWindowById(wx.ID_OK).Enable(False)
        
    def GetSelectedDir(self):
        return self.directories[self._selectionIndex]        
        
    def GetSelectedName(self):
        return self.values[self._selectionIndex]        

# Utilities
# Dump out default preferences
def dump_default_site_prefs():
    config.filteroptions.init(False)
    config.fgloboptions.init(False)
    config.viewoptions.init(False)
    config.fgloboptions._save_options()
    config.filteroptions._save_options()
    config.viewoptions._save_options()

def main(ipython=False, headless=False):
    global _running_ipython
    global _headless
    
    _running_ipython = ipython
    _headless = headless
    
    # Create application class
    app = SilkApp()
    if app.DidQuit():
        return
    app.MainLoop()
    return app

if __name__ == "__main__":
    main()
