#!/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 8348 2006-05-27 09:21:33Z ska $
##
######################################################################
import pygtk, pango, getopt, sys, os, cgi

sys.path.append(".")

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

from dissy.Config import *
from dissy.Objdump import *
from dissy.File import *
from dissy.Entity import Entity
from dissy.Instruction import Instruction
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()
	    inFile = ""
	else:
	    o = Objdump(inFile)
	    self.fileContainer, done = o.parse()

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

	self.display = DisplayModel()

	# 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)),
				      ('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()

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

	pasteBin.child.connect("activate", self.pasteBinCallback, pasteBin)
	hbox.pack_start(gtk.Label("Lookup"), expand=False, padding=2)
	hbox.pack_start(pasteBin)

	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 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("v1")
	about.set_copyright("(C) Simon Kagstrom, 2006")
	about.set_website(PROGRAM_URL)
	about.show()

    def toggleHighLevelCode(self, widget):
	config.showHighLevelCode = not config.showHighLevelCode
	try:
	    curFunction = self.functionView.get_model()[self.functionView.get_cursor()[0]][3]
	except TypeError:
	    # There is no function
	    return
	self.instructionView.set_model( InstructionModel.InfoModel( curFunction ).getModel() )

    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 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):
	comp1 = ("0x%08x" % entity.getAddress())
	comp2 = entity.getLabel()

	# Lookup either the address or the label when doing an interactive
	# search
	if comp1.find(key) != -1 or comp2.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).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("Call dst", 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()
 	cur = model[view.get_cursor()[0]][9]
	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 "  -h    Display this help and exit"
    sys.exit(1)

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

    for opt, arg in optlist:
	if opt == "-h":
	    usage()
    if len(args) == 0:
	filename = None
    else:
	filename = args[0]

    myGUI = GUI_Controller(filename)
    myGUI.run()
