#!/usr/bin/python
######################################################################
##
## Copyright (C) 2006,  Blekinge Institute of Technology
##
## Author:        Simon Kagstrom <simon.kagstrom@gmail.com>
## Description:   The main program
##
## Licensed under the terms of GNU General Public License version 2
## (or later, at your option). See COPYING file distributed with Dissy
## for full text of the license.
##
######################################################################
import pygtk, pango, getopt, sys, os, cgi, re

sys.path.append(".")
sys.path = ['/home/ska/projects/dissy/trunk/'] + sys.path

pygtk.require('2.0')
import gtk, gobject
gtk.gdk.threads_init()

from dissy.Config import *
from dissy.File import File, ExportedFile
from dissy.File import linuxKernelCrashRegexp
from dissy.Entity import Entity
from dissy.StrEntity import StrEntity
from dissy.Instruction import Instruction
from dissy.Function import Function
from dissy.PreferencesDialogue import PreferencesDialogue
from dissy.history import History
from dissy import FunctionModel
from dissy import InstructionModel
from dissy import InstructionModelHighlighter
from dissy.InfoBox import InfoBox

#Value analysis requires the WALi python bindings
enable_value_analysis = False
try:
    import wali
    print 'Enabling value analysis'
    enable_value_analysis = True
    from dissy.ValueAnalysisInfoProvider import ValueAnalysisInfoProvider
except:
    pass

NUM_JUMP_COLUMNS=3

# Navigation history
history = History()

class GUI_Controller:
    """ The GUI class is the controller for Dissy """

    def __init__(self, inFile=None):
        if inFile == None:
            inFile = ""
        self.fileContainer = File(baseAddress=baseAddress)

        icon = None
        icon_name = lookupFile('gfx/icon.svg')
        if icon_name != None:
            icon = gtk.gdk.pixbuf_new_from_file(icon_name)
            icon = icon.scale_simple(64, 64, gtk.gdk.INTERP_BILINEAR)

        #GUI setup from gtk.Builder
        builder = gtk.Builder()
        ui_file_name = lookupFile("dissy.ui")
        if ui_file_name == None:
            sys.stderr.write("Cannot find the dissy UI definition file (dissy.ui)\n")
            sys.exit(1)
        builder.add_from_file(ui_file_name)
        self.root = builder.get_object('mainWindow')
        self.root.set_title("%s - %s" % (PROGRAM_NAME, inFile))
        self.root.set_icon(icon)
        builder.connect_signals(self)
        self.root.show_all()

        #setup infobox
        self.vboxInfoBox = builder.get_object('vboxInfoBox')
        self.infoBox = InfoBox()
        scrwInfoBox = builder.get_object('scrwInfoBox')
        scrwInfoBox.add(self.infoBox)
        scrwInfoBox.show_all()

        lookupCombo = builder.get_object('lookupCombo')
        liststore = gtk.ListStore(gobject.TYPE_STRING)
        lookupCombo.set_model(liststore)
        lookupCombo.set_text_column(0)
        lookupCombo.child.connect("activate", self.pasteBinEntryActivated, lookupCombo)
        lookupCombo.connect("changed", self.pasteBinChanged, lookupCombo)

        highlightEntry = builder.get_object('highlightEntry')
        highlightEntry.connect("activate", self.patternMatchBinCallback, highlightEntry)

        self.vpaned = builder.get_object('vpaned')
        self.vpaned.set_position(650/3)
        self.hpaned_down = builder.get_object('hpaned')
        self.hpaned_down.set_position(900/3*2)

        #Set shortcuts, that doesn't work from Gtk.Builder
        accelgroupMainWin = builder.get_object('accelgroupMainWin')
        builder.get_object('mnuShowSource').add_accelerator('activate', accelgroupMainWin, ord('u'),
            gtk.gdk.CONTROL_MASK, gtk.ACCEL_VISIBLE)
        builder.get_object('navBack').add_accelerator('activate', accelgroupMainWin,
            gtk.gdk.keyval_from_name('Left'),
            gtk.gdk.MOD1_MASK, gtk.ACCEL_VISIBLE)
        builder.get_object('navForward').add_accelerator('activate', accelgroupMainWin,
            gtk.gdk.keyval_from_name('Right'),
            gtk.gdk.MOD1_MASK, gtk.ACCEL_VISIBLE)

        # Get the model and attach it to the view
        self.searchwordHighlighter = InstructionModelHighlighter.SearchwordHighlighter()
        self.conditionFlagHighlighter = InstructionModelHighlighter.ConditionFlagHighlighter()
        self.highlighters = [
            self.searchwordHighlighter,
            self.conditionFlagHighlighter,
        ]

        self.enableValueAnalysis(builder)

        # Setup focus chains (no tab-to-focus in the information box)
        builder.get_object('vboxMain').set_focus_chain([ self.vpaned ])
        self.hpaned_down.set_focus_chain([ builder.get_object('scrwInstruction') ])

        self.instructionView = builder.get_object('instructionView')
        self.instructionController = InstructionViewController(self, self.instructionView)

        self.functionView = builder.get_object('functionView')
        self.functionController = FunctionViewController(self, self.functionView,
            self.instructionController)

        #Load settings
        self.loading = True
        builder.get_object('actionShowInformationBox').set_active(config.showInstructionInformationBox)
        builder.get_object('actionShowSource').set_active(config.showHighLevelCode)
        self.setInformationBoxVisible()
        self.loading = False
        if inFile:
            self.loadFile(inFile)

        self.root.resize(config.width, config.height)

    def on_navBack_activate(self, widget):
        try:
            val = history.back()
        except:
            return
        return self.lookupFunction(val)

    def on_navForward_activate(self, widget):
        try:
            val = history.forward()
        except:
            return
        return self.lookupFunction(val)

    def on_quit(self, widget=None):
        config.width, config.height = self.root.get_size()
        config.save()
        gtk.main_quit()
        return

    def on_mnuSave_activate(self, widget):
        #TODO - refactor file handling
        pass

    def on_mnuOpen_activate(self, widget):
        file_open_dialog = gtk.FileChooserDialog(title="Open object file",
            action=gtk.FILE_CHOOSER_ACTION_OPEN,
                buttons=(gtk.STOCK_CANCEL,
                        gtk.RESPONSE_CANCEL,
                        gtk.STOCK_OPEN,
                        gtk.RESPONSE_OK))
        filter = gtk.FileFilter()
        filter.set_name("Object files")
        filter.add_pattern("*.o")
        filter.add_mime_type("application/x-object")
        filter.add_mime_type("application/x-executable")
        file_open_dialog.add_filter(filter)

        filter = gtk.FileFilter()
        filter.set_name("Dissy exported files (*.dissy)")
        filter.add_pattern("*.dissy")
        file_open_dialog.add_filter(filter)

        filter = gtk.FileFilter()
        filter.set_name("All files")
        filter.add_pattern("*")
        file_open_dialog.add_filter(filter)

        if file_open_dialog.run() == gtk.RESPONSE_OK:
            filename = file_open_dialog.get_filename()
            self.loadFile(filename)
            file_open_dialog.destroy()
        else:
            file_open_dialog.destroy()

    def on_mnuReload_activate(self, widget):
        self.loadFile()

    def loadFile(self, filename=None):
        global enable_value_analysis
        if filename == None:
            if not self.fileContainer:
                return
            filename = self.fileContainer.filename

        self.root.set_title("%s - %s" % (PROGRAM_NAME, filename))
        if filename.endswith('.dissy'):
            self.fileContainer = ExportedFile(filename, baseAddress=baseAddress)
        else:
            self.fileContainer = File(filename, baseAddress=baseAddress)
        self.functionView.set_model( FunctionModel.InfoModel(self.fileContainer).getModel() )
        if enable_value_analysis:
            self.valueAnalysisInfoProvider = ValueAnalysisInfoProvider(self.fileContainer)

    def on_mnuSaveAs_activate(self, widget):
        file_save_dialog = gtk.FileChooserDialog(title="Export file",
            action=gtk.FILE_CHOOSER_ACTION_SAVE,
                buttons=(gtk.STOCK_CANCEL,
                        gtk.RESPONSE_CANCEL,
                        gtk.STOCK_SAVE,
                        gtk.RESPONSE_OK))
        filter = gtk.FileFilter()
        filter.set_name("Dissy exported files (*.dissy)")
        filter.add_pattern("*.dissy")
        file_save_dialog.add_filter(filter)

        if file_save_dialog.run() == gtk.RESPONSE_OK:
            filename = file_save_dialog.get_filename()
            file_save_dialog.destroy()
        else:
            file_save_dialog.destroy()
            return

        if not filename.endswith('.dissy'):
            filename += '.dissy'

        f = open(filename, 'wb')
        expfile = self.fileContainer.toExportedFile()
        expfile.saveTo(f)
        f.close()

    def on_mnuAbout_activate(self, w=None):
        "Display the about dialogue"
        about = gtk.AboutDialog()
        about.set_name(PROGRAM_NAME)
        about.set_version("v%s" % (PROGRAM_VERSION) )
        about.set_copyright("(C) Simon Kagstrom, 2006-2011")
        about.set_website(PROGRAM_URL)
        about.run()
        about.hide()

    def redisplayFunction(self):
        try:
            fnCursor = self.functionView.get_cursor()[0]
            curFunction = self.functionView.get_model()[fnCursor][3]
        except TypeError:
            # There is no function currently being shown
            return
        insnCursor = self.instructionView.get_cursor()[0]
        curInstruction = None
        if insnCursor:
            curInstruction = self.instructionView.get_model()[insnCursor][self.instructionModel.COLUMN_INSTRUCTION]

        self.instructionModel.setCurInstruction(curInstruction)
        self.instructionModel.refreshModel()
        try:
            self.instructionView.set_cursor(insnCursor)
            self.instructionView.scroll_to_cell(insnCursor)
        except: # If nothing is selected this will fail
            pass

    def on_actionShowSource_toggled(self, action):
        if self.loading:
            return
        config.showHighLevelCode = not config.showHighLevelCode
        config.save()
        try:
            fnCursor = self.functionView.get_cursor()[0]
            curFunction = self.functionView.get_model()[fnCursor][3]
        except TypeError:
            # There is no function currently being shown
            return
        model = self.instructionView.get_model()
        curpath = self.instructionView.get_cursor()[0]
        try:
            curaddr = model[curpath][self.instructionModel.COLUMN_INSTRUCTION].address
        except:
            curaddr = None
        self.instructionView.set_model( InstructionModel.InfoModel(curFunction).getModel() )
        if curaddr:
            self.lookupFunction(curaddr)

    def on_actionShowInformationBox_toggled(self, widget):
        if self.loading:
            return
        config.showInstructionInformationBox = not config.showInstructionInformationBox
        config.save()
        self.setInformationBoxVisible()

    def setInformationBoxVisible(self):
        if config.showInstructionInformationBox:
            self.vboxInfoBox.show()
        else:
            self.vboxInfoBox.hide()

    def on_mnuPreferences_activate(self, widget):
        pd = PreferencesDialogue(self)

    def enableValueAnalysis(self, builder):
        global enable_value_analysis
        if not enable_value_analysis:
            mnuValueAnalysis = builder.get_object('mnuValueAnalysis')
            mnuValueAnalysis.set_sensitive(False)

    def on_mnuValueAnalysis_activate(self, widget):
        print 'Doing value analysis'
        self.valueAnalysisInfoProvider.analyse()
        print 'Done'

    def patternMatchBinCallback(self, entry, comboBox):
        markPattern = entry.get_text()
        self.searchwordHighlighter.setSearchPattern(markPattern)
        self.redisplayFunction()
        self.instructionView.grab_focus()

    def lookupFunction(self, val):
        function = self.fileContainer.lookup(val)
        history.disable()

        if function != None:
            model = self.functionView.get_model()
            self.functionView.set_cursor_on_cell(model.get_path(function.iter))
            self.functionView.row_activated(model.get_path(function.iter), self.functionController.viewColumns[0])

            # Return if this was just a label lookup
            if isinstance(val, str):
                history.enable()
                return True
            insn = function.lookup(val)

            if insn != None:
                model = self.instructionView.get_model()
                self.instructionView.set_cursor_on_cell(model.get_path(insn.iter))
            history.enable()
            return True

        history.enable()
        return False

    def addPasteBinEntry(self, comboBox, txt):
        """Add an entry to the pasteBin ComboBox. If the same entry is already
there, promote it to the top of the list"""
        count = 0
        found = -1
        for item in iter(comboBox.get_model()):
            if item[0] == txt:
                found = count
            count = count + 1

        # If it's found, remove it from the list
        if found != -1:
            comboBox.remove_text(found)

        # ... and add it again
        comboBox.prepend_text(txt)

    def lookupWord(self, word):
        """Lookup a word /address"""
        # Special-case Linux crashes
        r = linuxKernelCrashRegexp.match(word)
        if r != None:
            fn_name = r.group(1)
            fn_addend = r.group(2)
            function = self.fileContainer.lookup(fn_name)

            if function != None:
                val=function.address + long(fn_addend, 16)
                if self.lookupFunction( val ):
                    history.add(val)
            return

        try:
            # Try to convert to a number (handle some common cases)
            word = word.strip("[]+():-/|><")
            val = long(word, 16)
        except:
            val = word
        if self.lookupFunction(val):
            history.add(val)

    def pasteBinEntryActivated(self, entry, comboBox):
        """
        Called to lookup a symbol / address. Looks up a label or an
        address.
        """
        txt = entry.get_text()

        self.addPasteBinEntry(comboBox, txt)

        # Split the input in words and navigate to them
        for word in txt.split():
            self.lookupWord(word)

    def pasteBinChanged(self, comboBox, usr):
        """
        Called when an entry in the paste bin list is selected
        """
        which = comboBox.get_active()
        if which != -1:
            txt = comboBox.get_model()[which][0]
            self.addPasteBinEntry(comboBox, txt)
            for word in txt.split():
                self.lookupWord(word)

    def clearInstructionInfo(self):
        self.infoBox.set_markup('<b>Not available</b>')

    def updateInstructionInfo(self, instruction):
        # Check if this architecture supports instruction info
        if not hasattr(self.fileContainer.getArch(), "getInstructionInfo"):
            self.clearInstructionInfo()
            return
        instrInfo = self.fileContainer.getArch().getInstructionInfo(instruction)

        valueAnalysisInfo = ""
        try:
            valueAnalysisInfo = self.valueAnalysisInfoProvider.getInstructionInfo(instruction)
        except AttributeError, e:
            # Set above
            pass

        self.infoBox.set_markup(
            '<b>' + instrInfo.get('shortinfo', '') + "</b><br />\n\n" +
            instrInfo.get('description', '').replace('\n', '<br />\n') + "\n" +
            '<p>' + valueAnalysisInfo + '</p>'
            )

    def searchCommon(self, entity, key):
        key = key.lower()
        comp1 = ("0x%08x" % entity.address).lower()
        comp2 = entity.getLabel().lower()
        if isinstance(entity, Instruction):
            comp3 = entity.getOpcode() + entity.getArgs()
        else:
            comp3 = ""

        # Lookup either the address or the label when doing an interactive
        # search
        if comp1.find(key) != -1 or comp2.find(key) != -1 or comp3.find(key) != -1:
            return False
        return True

    def run(self):
        """ run is called to set off the GTK mainloop """
        gtk.main()
        return

class FunctionViewController:
    """ Controls the Function TreeView """

    def __init__(self, controller, treeview, instructionController):
        self.controller = controller
        self.functionView = treeview
        self.instructionController = instructionController

        self.setupFunctionView()

    def setupFunctionView(self):
        """ Sets up the columns to be displayed """

        # setup the cell renderers
        self.functionRenderer = gtk.CellRendererText()
        self.functionRenderer.set_property("font", "Monospace")

        self.functionView.connect( 'row-activated', self.functionRowActivated)
        self.functionView.set_search_column(0)
        self.functionView.set_search_equal_func(self.functionSearchCallback)

        self.viewColumns = {}
        # Connect column0 of the display with column 0 in our list model
        # The renderer will then display whatever is in column 0 of
        # our model .
        self.viewColumns[0] = gtk.TreeViewColumn("Address", self.functionRenderer, markup=0)
        self.viewColumns[1] = gtk.TreeViewColumn("Size", self.functionRenderer, markup=1)
        self.viewColumns[2] = gtk.TreeViewColumn("Label", self.functionRenderer, markup=2)

        # The columns active state is attached to the second column
        # in the model.  So when the model says True then the button
        # will show as active e.g on.
        for col in self.viewColumns.values():
            self.functionView.append_column( col )

    def functionSearchCallback(self, model, column, key, iter):
        """
        Callback for interactive searches.
        """
        entity = model[iter][3]
        return self.controller.searchCommon(entity, key)

    def functionRowActivated( self, view, iter, path):
        """
        Run when one row is selected (double-click/space)
        """
        model = self.functionView.get_model()
        entity = model[iter][3]
        entity.link()
        history.add(entity.address)

        self.instructionController.setInstructionModel(entity)

class InstructionViewController:
    """ Controls the Instruction TreeView """

    def __init__(self, controller, treeview):
        self.controller = controller
        self.insnView = treeview

        self.setupInstructionView()

    def setupInstructionView(self):
        # setup the cell renderers
        link_renderer = gtk.CellRendererPixbuf()

        insnRenderer = gtk.CellRendererText()
        addressRenderer = gtk.CellRendererText()
        callDstRenderer = gtk.CellRendererText()

        addressRenderer.set_property("font", "Monospace")
        insnRenderer.set_property("font", "Monospace")
        insnRenderer.set_property("width", 500)
        link_renderer.set_property("width", 22)
        link_renderer.set_property("height", 22)
        insnRenderer.set_property("height", 22)
        insnRenderer.set_property("editable", True)
        insnRenderer.connect("edited", self.insnCommentEdited, self.insnView)
        insnRenderer.connect("editing-started", self.insnCommentEditStart, self.insnView)
        callDstRenderer.set_property("font", "Monospace")

        self.insnView.connect( 'row-activated', self.insnRowActivated)
        self.insnView.connect_after( 'move-cursor', self.insnMoveCursor)
        self.insnView.connect( 'cursor-changed', self.insnCursorChanged )
        self.insnView.connect( 'key-press-event', self.insnKeyPress )
        self.insnView.set_search_column(0)
        self.insnView.set_search_equal_func(self.insnSearchCallback)

        self.insnColumns = {}
        # Connect column0 of the display with column 0 in our list model
        # The renderer will then display whatever is in column 0 of
        # our model .
        self.insnColumns[0] = gtk.TreeViewColumn("Address", addressRenderer, markup=0)
        self.insnColumns[1] = gtk.TreeViewColumn("b0", link_renderer, pixbuf=1)
        self.insnColumns[2] = gtk.TreeViewColumn("b1", link_renderer, pixbuf=2)
        self.insnColumns[3] = gtk.TreeViewColumn("b2", link_renderer, pixbuf=3)
        self.insnColumns[4] = gtk.TreeViewColumn("Instruction", insnRenderer, markup=4)
        self.insnColumns[4].set_resizable(True)
        self.insnColumns[5] = gtk.TreeViewColumn("f0", link_renderer, pixbuf=5)
        self.insnColumns[6] = gtk.TreeViewColumn("f1", link_renderer, pixbuf=6)
        self.insnColumns[7] = gtk.TreeViewColumn("f2", link_renderer, pixbuf=7)
        self.insnColumns[8] = gtk.TreeViewColumn("Target", callDstRenderer, markup=8)

        # The columns active state is attached to the second column
        # in the model.  So when the model says True then the button
        # will show as active e.g on.
        for col in self.insnColumns.values():
            self.insnView.append_column( col )

    def setInstructionModel(self, functionEntity):
        """
        Change the instruction model.
        """
        self.controller.instructionModel = InstructionModel.InfoModel(functionEntity,
            highlighters=self.controller.highlighters)
        model = self.controller.instructionModel.getModel()
        self.insnView.set_model( model )

        self.insnView.grab_focus()
        self.insnView.set_cursor(0)

    def insnSearchCallback(self, model, column, key, iter):
        """
        Callback for interactive searches.
        """
        entity = model[iter][self.controller.instructionModel.COLUMN_INSTRUCTION]
        if isinstance(entity, StrEntity):
            return True
        return self.controller.searchCommon(entity, key)

    def insnKeyPress(self, view, event):
        if event.string == ';':
            view.set_cursor(view.get_cursor()[0], focus_column=self.insnColumns[4], start_editing=True)
            return True

    def insnCommentEditStart(self, insnRenderer, editable, path, insnView):
        model = insnView.get_model()
        iter = model.get_iter_from_string(path)
        insn = model[iter][self.controller.instructionModel.COLUMN_INSTRUCTION]
        editable.set_text(insn.comment)

    def insnCommentEdited(self, insnRenderer, path, new_text, view):
        model = view.get_model()
        iter = model.get_iter_from_string(path)
        insn = model[iter][self.controller.instructionModel.COLUMN_INSTRUCTION]
        insn.comment = new_text
        view.set_cursor(view.get_cursor()[0], focus_column=self.insnColumns[0], start_editing=False)

    def insnCursorChanged(self, view):
        model = view.get_model()
        cur = model[view.get_cursor()[0]][self.controller.instructionModel.COLUMN_INSTRUCTION]
        if isinstance(cur, Instruction):
            self.controller.updateInstructionInfo(cur)
            self.controller.instructionModel.setCurInstruction(cur)
            self.controller.instructionModel.refreshModel()
        else:
            self.controller.clearInstructionInfo()

    def insnMoveCursor(self, view, step, count):
        model = view.get_model()
        if not model:
            return True
        curpath = view.get_cursor()[0][0]
        origpath = view.get_cursor()[0][0]

        if step == gtk.MOVEMENT_DISPLAY_LINES:
            #Find next Instruction
            try:
                while not isinstance(model[curpath][self.controller.instructionModel.COLUMN_INSTRUCTION],
                    Instruction):
                    curpath = curpath + count
            except IndexError:
                #Reverse move
                curpath = origpath - count
            if curpath < 0:
                #Reverse move
                curpath = origpath - count
            view.set_cursor(curpath)
        return True

    def insnRowActivated( self, view, iter, path):
        """
        Run when one row is selected (double-click/space)
        """
        model = view.get_model()
        functionModel = self.controller.functionView.get_model()
        try:
            entity = model[iter][self.controller.instructionModel.COLUMN_INSTRUCTION]
        except IndexError:
            # If the index is outside of the model
            return
        if isinstance(entity, Instruction) and entity.hasLink():
            link = entity.getOutLink()
            history.add(entity.address)
            if isinstance(link, Function):
                history.add(link.address)
                dst = link
                self.controller.functionView.set_cursor_on_cell(functionModel.get_path(dst.iter))
                self.controller.functionView.row_activated(functionModel.get_path(dst.iter), self.controller.functionController.viewColumns[0])
                view.set_cursor_on_cell(0)
            else:
                func = entity.getFunction()
                dst = func.lookup(link.address)
                if dst != None:
                    history.add(dst.address)
                    view.set_cursor(model.get_path(dst.iter))

def usage():
    print "Usage: %s -h [FILE]" % (PROGRAM_NAME.lower())
    print "Disassemble FILE and open in a graphical window.\n"
    print "  -t BASE_ADDRESS       Set the start address for the disassembled file (.text segment)"
    print "  -h                    Display this help and exit"
    sys.exit(1)

baseAddress = 0
if __name__ == "__main__":
    optlist, args = getopt.gnu_getopt(sys.argv[1:], "ht:")

    for opt, arg in optlist:
        if opt == "-h":
            usage()
        if opt == "-t":
            try:
                baseAddress = long(arg)
            except:
                try:
                    baseAddress = long(arg, 16)
                except:
                    raise
    if len(args) == 0:
        filename = None
    else:
        filename = args[0]

    myGUI = GUI_Controller(filename)
    myGUI.run()
