#!/usr/bin/env python
######################################################################
##
## Copyright (C) 2006,  Blekinge Institute of Technology
##
## Filename:      dispy.py
## Author:        Simon Kagstrom <ska@bth.se>
## Description:   The main program
##
## $Id: dissy 12443 2006-11-25 10:14:44Z ska $
##
######################################################################
import pygtk, pango, getopt, sys, os, cgi, re

sys.path.append(".")

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

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

NUM_JUMP_COLUMNS=3

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

	try:
	    f = open(fullPath)
	    out = f.read()
	    f.close()
	    return out
	except:
	    pass
    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.markPattern = None

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

	self.display = DisplayModel(self.markPattern)

	# setup the main window
	self.root = gtk.Window(type=gtk.WINDOW_TOPLEVEL)
	self.root.set_title("%s - %s" % (PROGRAM_NAME, inFile))
	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'),
				      ('Preferences', gtk.STOCK_PREFERENCES, '_Preferences', None,
				       'Configure preferences for %s' % (PROGRAM_NAME), lambda w: PreferencesDialogue()),
				      ('Toggle source', None, '_Toggle source', None,
				       'Toggle the showing of high-level source', self.toggleHighLevelCode),
				      ('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)
	sw_down = gtk.ScrolledWindow()
	sw_down.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)

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

	vbox.set_focus_chain([ vpaned ])

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

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

	self.root.add(vbox)

	self.root.show_all()

    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")
	about.set_website(PROGRAM_URL)
	about.show()

    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]
	self.instructionView.set_model( InstructionModel.InfoModel(curFunction, self.markPattern).getModel() )
	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
	self.redisplayFunction()

    def patternMatchBinCallback(self, entry, comboBox):
	markPattern = entry.get_text()
	self.markPattern = re.compile(markPattern)

	self.display.markPattern = self.markPattern
	self.redisplayFunction()

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

	try:
	    # Try to convert to a number (handle some common cases)
	    text = entry.get_text().strip()
	    if not text.startswith("0x"):
		text = "0x%s" % (text)
	    if text.endswith(":"):
		text = text[:-1]
	    val = long(text, 16)
	except:
	    val = entry.get_text()
	function = self.fileContainer.lookup(val)

	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):
		return
	    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])

    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, markPattern):
	self.markPattern = markPattern

    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][9]
	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()
	model = InstructionModel.InfoModel(entity, self.markPattern).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.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 insnMoveCursor(self, view, step, count, user):
	model = view.get_model()
	try:
	    cur = model[view.get_cursor()[0]][9]
	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][9]
	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):
		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:
		    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 "  --text-address=ADDRESS    Set the start address for the file (code)"
    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()