#!/usr/bin/env 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

from dissy.Config import *
from dissy.File import File
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.FileDialogue import FileDialogue
from dissy.history import History
from dissy import FunctionModel
from dissy import InstructionModel
from dissy import InstructionModelHighlighter

NUM_JUMP_COLUMNS=3

# Navigation history
history = History()

def lookupFile(name):
    pathsToSearch = ['.', '/usr/local/share/%s' % (PROGRAM_NAME).lower(),
                     '/usr/share/%s' % (PROGRAM_NAME).lower()]
    for path in pathsToSearch:
        fullPath = "%s/%s" % (path, name)

        try:
            st = os.lstat(fullPath)
            return fullPath
        except:
            pass
    return None

def loadFile(fileName):
    try:
        f = open(lookupFile(fileName))
        out = f.read()
        f.close()
        return out
    except:
        return None

# Taken from the cellrenderer.py example
class GUI_Controller:
    """ The GUI class is the controller for Dissy """

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

        self.searchwordHighlighter = InstructionModelHighlighter.SearchwordHighlighter()
        self.conditionFlagHighlighter = InstructionModelHighlighter.ConditionFlagHighlighter()
        self.highlighters = [
            self.searchwordHighlighter,
            self.conditionFlagHighlighter,
        ]

        functionModel = FunctionModel.InfoModel(self.fileContainer).getModel()
        self.instructionModel = InstructionModel.InfoModel(None, highlighters=self.highlighters)
        insnModel = self.instructionModel.getModel()

        self.display = DisplayModel(self)

        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)

        # setup the main window
        self.root = gtk.Window(type=gtk.WINDOW_TOPLEVEL)
        self.root.set_title("%s - %s" % (PROGRAM_NAME, inFile))
        self.root.set_icon(icon)
        self.root.connect("destroy", self.destroy_cb)
        self.root.set_default_size(900, 600)

        # Boxes for the widgets
        vbox = gtk.VBox()
        hbox = gtk.HBox()

        # menubar
        self.uimgr = gtk.UIManager()
        self.accelgroup = self.uimgr.get_accel_group()
        self.root.add_accel_group(self.accelgroup)

        # Create an ActionGroup
        self.actiongroup = gtk.ActionGroup('UIManagerExample')

        # Create actions
        self.actiongroup.add_actions([('Quit', gtk.STOCK_QUIT, '_Quit', None,
                                       'Quit the Program', self.destroy_cb),
                                      ('Open', gtk.STOCK_OPEN, '_Open', None,
                                       'Open a file', lambda w: FileDialogue(self)),
                                      ('Reload', None, '_Reload', '<Control>r',
                                       'Reload a file', lambda w: self.loadFile() ),
                                      ('File', None, '_File'),
                                      ('Options', None, '_Options'),
                                      ('Navigation', None, '_Navigation'),
                                      ('Forward', gtk.STOCK_GO_FORWARD, '_Forward', '<Alt>Right',
                                       'Navigate forwards in history', lambda w: self.forward()),
                                      ('Back', gtk.STOCK_GO_BACK, '_Back', '<Alt>Left',
                                       'Navigate backwards in history', lambda w: self.back()),
                                      ('Preferences', gtk.STOCK_PREFERENCES, '_Preferences', None,
                                       'Configure preferences for %s' % (PROGRAM_NAME), self.preferencesDialogue),
                                      ('Toggle source', None, '_Toggle source', None,
                                       'Toggle the showing of high-level source', self.toggleHighLevelCode),
                                      ('Toggle information box', None, 'Toggle _information box', None,
                                       'Toggle the showing of the instruction information box', self.toggleInformationBox),
                                      ('Help', None, '_Help'),
                                      ('About', gtk.STOCK_ABOUT, '_About', None,
                                       'About %s' % PROGRAM_NAME, self.about),
                                      ])
        # Add the actiongroup to the uimanager
        self.uimgr.insert_action_group(self.actiongroup, 0)

        self.uimgr.add_ui_from_string(loadFile("menubar.xml"))

        # Pastebin for quick lookup of symbols
        pasteBin = gtk.combo_box_entry_new_text()

        # Pattern matcher
        patternMatchBin = gtk.Entry()

        # Move to the pasteBin with Ctrl-l
        pasteBin.child.add_accelerator("grab-focus", self.accelgroup,
                                       ord('L'), gtk.gdk.CONTROL_MASK, gtk.ACCEL_VISIBLE)
        # Move to the pattern match bin with Ctrl-k
        patternMatchBin.add_accelerator("grab-focus", self.accelgroup,
                                        ord('K'), gtk.gdk.CONTROL_MASK, gtk.ACCEL_VISIBLE)

        pasteBin.child.connect("activate", self.pasteBinCallback, pasteBin)

        patternMatchBin.connect("activate", self.patternMatchBinCallback, patternMatchBin)

        tooltips = gtk.Tooltips()
        tooltips.set_tip(pasteBin.child, "Lookup an address or symbol (shortcut Ctrl-l)")
        tooltips.set_tip(patternMatchBin, "Enter a pattern to highlight (shortcut Ctrl-k)")

        hbox.pack_start(gtk.Label("Lookup"), expand=False, padding=2)
        hbox.pack_start(pasteBin)
        hbox.pack_start(gtk.Label("Highlight"), expand=False, padding=2)
        hbox.pack_start(patternMatchBin, expand=False, padding=2)

        vbox.pack_start(self.uimgr.get_widget("/MenuBar"), expand=False)
        vbox.pack_start(hbox, expand=False, padding=2)

        sw_up = gtk.ScrolledWindow()
        sw_up.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)
        self.sw_down = gtk.ScrolledWindow()
        self.sw_down.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)

        self.vpaned = gtk.VPaned()
        self.vpaned.set_position(650/3)
        vbox.pack_start(self.vpaned)

        # Get the model and attach it to the view
        self.functionView, self.instructionView = self.display.makeViews( functionModel, insnModel )

        # Add instruction info section
        self.hpaned_down = gtk.HPaned()
        self.hpaned_down.set_position(900/3*2)
        self.vbox_instr_info = gtk.VBox()

        self.instr_info_label = gtk.Label("<b>Instruction Information</b>")
        self.instr_info_label.set_use_markup(True)
        self.vbox_instr_info.pack_start(self.instr_info_label, expand=False)
        self.instr_info_description = gtk.TextView()
        self.instr_info_description.set_editable(False)

        self.instr_info_description.set_wrap_mode(gtk.WRAP_WORD)
        self.vbox_instr_info.pack_start(self.instr_info_description)

        self.sw_down.add(self.instructionView)
        self.hpaned_down.pack1(self.sw_down, resize=True)
        self.hpaned_down.pack2(self.vbox_instr_info, resize=False)

        # Add our view into the scrolled window
        sw_up.add(self.functionView)
        self.vpaned.add1(sw_up)
        self.vpaned.add2(self.hpaned_down)

        self.root.add(vbox)

        # Setup focus chains (no tab-to-focus in the information box)
        vbox.set_focus_chain([ self.vpaned ])
        self.hpaned_down.set_focus_chain([ self.sw_down ])

        self.root.show_all()

        # Make the info box visible or not depending on the configuration
        self.setInformationBox()

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

        self.root.set_title("%s - %s" % (PROGRAM_NAME, filename))
        self.fileContainer = File(filename, baseAddress=baseAddress)
        self.functionView.set_model( FunctionModel.InfoModel(self.fileContainer).getModel() )

    def loadTimeoutCallback(self, o):
        self.fileContainer, done = o.parse(10)
        self.functionView.set_model( FunctionModel.InfoModel( self.fileContainer ).getModel() )
        return done

    def about(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-2009")
        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][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 toggleHighLevelCode(self, widget):
        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
        self.instructionView.set_model( InstructionModel.InfoModel(curFunction).getModel() )
        self.redisplayFunction()

    def setInformationBox(self):
        if config.showInstructionInformationBox:
            self.vbox_instr_info.show()
        else:
            self.vbox_instr_info.hide()

    def toggleInformationBox(self, widget):
        config.showInstructionInformationBox = not config.showInstructionInformationBox
        config.save()
        self.setInformationBox()

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

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

    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.display.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.display.insnView.get_model()
                self.display.insnView.set_cursor_on_cell(model.get_path(insn.iter))
                self.display.insnView.row_activated(model.get_path(insn.iter), self.display.insnColumns[0])
            history.enable()
            return True

        history.enable()
        return False

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

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

            # 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.getAddress() + long(fn_addend, 16)
                    if self.lookupFunction( val ):
                        history.add(val)
                continue

            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 clearInstructionInfo(self):
        self.instr_info_description.get_buffer().set_text('Not available')

    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)

        self.instr_info_description.get_buffer().set_text(
            instrInfo.get('shortinfo', '') + "\n\n" +
            instrInfo.get('description', '')
            )

    def forward(self, filename=None):
        try:
            val = history.forward()
        except:
            return
        return self.lookupFunction(val)

    def back(self, filename=None):
        try:
            val = history.back()
        except:
            return
        return self.lookupFunction(val)

    def destroy_cb(self, *kw):
        """ Destroy callback to shutdown the app """
        gtk.main_quit()
        return

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


class DisplayModel:
    """ Displays the Info_Model model in a view """

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

    def makeFunctionView( self, model ):
        """ Form a view for the Tree Model """
        self.functionView = gtk.TreeView( model )

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

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

        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 )
        return self.functionView

    def searchCommon(self, entity, key):
        key = key.lower()
        comp1 = ("0x%08x" % entity.getAddress()).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 functionSearchCallback(self, model, column, key, iter, unused):
        """
        Callback for interactive searches.
        """
        entity = model[iter][3]
        return self.searchCommon(entity, key)

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

    def functionRowActivated( self, view, iter, path, model ):
        """
        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.controller.instructionModel = InstructionModel.InfoModel(entity,
            highlighters=self.controller.highlighters)
        model = self.controller.instructionModel.getModel()
        self.insnView.set_model( model )
        self.insnView.connect( 'row-activated', self.insnRowActivated, model )

    def makeInstructionView(self, model):
        self.insnView = gtk.TreeView( model )

        # 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)
        callDstRenderer.set_property("font", "Monospace")

        self.insnView.connect( 'row-activated', self.insnRowActivated, model )
        self.insnView.connect( 'move-cursor', self.insnMoveCursor, None )
        self.insnView.connect( 'cursor-changed', self.insnCursorChanged )
        self.insnView.set_search_column(0)
        self.insnView.set_search_equal_func(self.insnSearchCallback, model)

        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[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 )
        return self.insnView

    def insnCursorChanged(self, view):
        model = view.get_model()
        cur = model[view.get_cursor()[0]][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, user):
        model = view.get_model()
        try:
            cur = model[view.get_cursor()[0]][InstructionModel.COLUMN_INSTRUCTION]
        except:
            # There is no model, just ignore
            return
        function = cur.getFunction()

        if step == gtk.MOVEMENT_DISPLAY_LINES:
            all = function.getAll()
            nextIdx = all.index(cur)
            try:
                while not isinstance(all[nextIdx + count], Instruction):
                    nextIdx = nextIdx + count
            except IndexError:
                return True
            if nextIdx < 0:
                return True
            view.set_cursor(model.get_path(all[nextIdx].iter))

        return True


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



    def makeViews( self, functionModel, insnModel ):

        functionView, instructionView = self.makeFunctionView( functionModel), self.makeInstructionView( insnModel )

        return functionView, instructionView

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()
