
'''
The vdb command / GUI extensions from vivisect!
'''
import time
import traceback

import vdb
import vtrace
import vdb.qt.main as vdb_qt_main

import envi.cli as e_cli
import vivisect.colormap as viv_color

from envi.threads import firethread
from vqt.main import idlethread, idlethreadsync

class VivNotif(vtrace.Notifier):

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

    @idlethreadsync
    def notify(self, event, trace):
        try:

            if event == vtrace.NOTIFY_BREAK and not trace.shouldRunAgain():

                pc = trace.getProgramCounter()
                if self.vivgui.vw.isValidPointer(pc):
                    self.vivgui.vivNavSignal.emit('0x%.8x' % pc)

                # If there are stalker hits, lets make a color map for them...
                shits = trace.getMeta('StalkerHits', None)
                if shits != None:
                    self.vivgui.vw.vprint('vdb ext found stalker hits!')
                    cmap = viv_color.VivColorMap(self.vivgui.vw)
                    for hit in shits:
                        try:
                            cmap.colorBlock(hit, 'yellow')
                        except Exception, e:
                            self.vivgui.vw.vprint('WARNING: stalker color map: %s' % e)

                    #self.vivgui.setColorMap(cmap.getColorDict())
                    #cmap.saveAs('stalker')
                    cmap.setGuiMap()

        except Exception, e:

            traceback.print_exc()


@firethread
def doLoad(db, trace, base):
    db.vw.vprint('Loading from memory 0x%.8x...' % base)
    db.vw.loadFromMemory(trace, base)
    db.vw.vprint('...loaded.')
    db.vw.analyze()

def vivimport(db, line):
    '''
    Import either a library or memory map from the target
    process into the current vivisect workspace.
    '''
    vw = db.vw
    trace = db.getTrace()

    argv = e_cli.splitargs(line)
    if len(argv) != 1:
        return db.do_help("vivimport")

    try:
        base = db.parseExpression(line)
    except Exception, e:
        db.vprint("Invalid Address Expression: %s" % argv[0])
        return

    # Init Arch/Format if they are not...
    vw.initMeta('Architecture', trace.getMeta('Architecture'))
    vw.initMeta('Format', trace.getMeta('Format'))

    doLoad(db, trace, base)

def extendVdb(db, vivgui):
    db.vw = vivgui.vw # Store a ref to the workspace...
    db.vivgui = vivgui # Also a reference to the GUI
    db.registerNotifier(vtrace.NOTIFY_ALL, VivNotif(vivgui))
    db.registerCmdExtension(vivimport)

@firethread
def fireattach(trace, pid):
    trace.attach(pid)

@idlethread
def runVdb(vivgui, pid=None):
    try:
        db = vdb.Vdb()
        extendVdb(db, vivgui)
        vgui = vdb_qt_main.VdbWindow(db)
        vgui.show()
        if pid != None:
            fireattach(db.trace, pid)
    except Exception, e:
        vivgui.vw.vprint('Error Running VDB: %s' % e)


