import envi
import envi.memory as e_mem
import envi.common as e_cmn
import envi.archs.i386.opconst as e_i386
import vivisect.impemu.monitor as vimp_mon

import logging
import collections
from pprint import pprint

logger = logging.getLogger(__name__)

# Malware Analysis Module
# OgreMemoryModule - allows for layering of memory cache, and identifying which layer the latest address was updated in

## AntiDebugging/AntiReversing Analysis:
# RDTSC
# AntiDebugging
# AntiVM check
# "Plateau" execution
# Branching into a modified bit of code
# Process Injection
# NTDLL access through PEB
# NtCreateFile with Exclusive Access to the EXE/DLL PE file(fails if Debugger Attached).  Filename, Permissions, Share==0
# CreateProcess*W  - Create the process SUSPENDED, write into it, resume the thread.
    # hook NtResumeThread and NtWriteVirtualMemory (kernel32_WriteProcessMemory+64?), if debugging




# do we want to make the Workspace Memory Map handle overwrites?  Or perhaps, in stages (at a plateau)?

class MalPlayMonitor(vimp_mon.EmulationMonitor):
    def __init__(self):
        self.pause=False
        self.eventlog = []
        self._branching = False
        self._branches = []
        self._plateaus = []
        self._initted = False

        self.ophist = collections.defaultdict(int)
        self.ophist_temp = collections.defaultdict(int)
        
    def doInit(self, emu):
        self.emu = emu
        emu.logread = True
        emu.logwrite = True
        emu.usecache = False
        self._wtrack = WriteTracker_1(self.emu)
        self._initted = True


    def doPause(self):
        print("PAUSING!")
        import envi.interactive as ei
        ei.dbg_interact(locals(), globals())
        self.resume()

    def pause(self):
        self.pause = True

    def resume(self):
        self.pause = False

    def checkWrittenMemory(self, va):
        self._wtrack.mergeFromWriteLog()
        written = self._wtrack.writtenSinceLast(va)
        if written:
            # track when we do this...
            self.logEvent("Executing Overwritten Code: %x" % va, self._wtrack._mem)
            self.logEvent(" - ophist", self.ophist_temp)
            self._wtrack._mem.ophist = self.ophist_temp
            pprint(self.ophist_temp)
            self.ophist_temp.clear()

            newmc = self._wtrack.applyNewMemCache()
            input("PRESS ENTER TO CONTINUE")

    def logEvent(self, title, data=None, maxrepr=100):
        '''
        Log an event during analysis.
        This will grab and store the opcode count of the current run
        '''
        valist = self.emu.path[2]['valist']
        opidx = len(valist)
        lastva = valist[-1]

        if data is None:
            datarep = ''
        else:
            datarep = repr(data)[:maxrepr]

        print("==--==--== EVENT: (near idx:%d  0x%x):  %r  %r" % (opidx, lastva, title, datarep))
        self.eventlog.append((opidx, lastva, title, data))

    def prehook(self, emu, op, starteip):
        if not self._initted:
            self.doInit(emu)

        # check if we're executing in modified bytes
        self.checkWrittenMemory(starteip)

        self.ophist[op.mnem] += 1
        self.ophist_temp[op.mnem] += 1

        if op.iflags & (envi.IF_CALL | envi.IF_BRANCH):
            self._branching = True
            logger.log(e_cmn.SHITE, "BRANCH: %x: %r" % (op.va, op))
            tgtva = op.opers[-1].getOperValue(op, emu=emu)
            self._branches.append([starteip, tgtva, None])

        if self.pause:
            self.doPause()

        # this is x86-specific, but too good to ignore.  perhaps need arch/plat-specific extension
        if op.opcode == e_i386.INS_RDTSC:
            self.logEvent("RDTSC!")

        if op.mnem.startswith('vm'):
            self.logEvent("VM Instruction!", op)

        # CHECK for PEB/TEB accesses
        # CHEKC for NTDLL accessess



    def posthook(self, emu, op, endeip):
        if self._branching:
            if self._branches[-1][2] is not None:
                print("Branch Landing not ok. %r\n  Adding new entry." % repr(self._branches[-1]))
                self._branches.append([None, None, None])

            self._branching = False
            self._branches[-1][2] = endeip

    def apicall(self, emu, op, pc, api, argv):
        pass

class WriteTracker_1:
    '''
    Track and access the writelog from the current WorkspaceEmulator.

    On checking the writelog, if it exceeds a certain max size, it's sorted into pages for faster
    lookups.

    Also, allows bookmarks for tracking of events.  This allows us to check if memory has been 
    written to *since* a particular bookmark

    ---
    a couple approaches:
    # 1
    * build layers of MemoryCaches, never apply writes
    * this allows rolling back/tracking which phase established which thing
    * check dirty cache page, followed by writes since bookmark
    '''
    def __init__(self, emu, maxlogsz=1000):
        self.emu = emu
        self.maxlogsz = maxlogsz
        self.bookmarks = [0]
        # last_bookmark is last time we created 
        self.last_bookmark = 0

        # _last_merged tracks last time the MemoryCache has been updated
        self._last_merged = 0   # from writelog into current MemoryCache
        self._mem = e_mem.MemoryCache(emu)
        self._mem_cache_bkmks = [(0, self._mem)]

        self.wlog = emu.path[2].get('writelog')


    def setBookmark(self):
        bkmk = len(self.wlog)
        self.bookmarks.append(bkmk)
        self.last_bookmark = bkmk       # current index in write stream
        return bkmk


    def mergeFromWriteLog(self):
        logger.log(e_cmn.MIRE, "mergeFromWriteLog")
        '''
        update the MemoryCache with recent writes
        '''
        bkmk = len(self.wlog)
        if bkmk == self._last_merged:
            return

        # apply writes to the memcache
        for idx in range(self._last_merged, bkmk):
            memwrite = self.wlog[idx]
            cva, tva, wdata = memwrite
            self._mem.writeMemory(tva, wdata)
            logger.log(e_cmn.SHITE, " - writing to memcache: %x" % tva)

        logger.log(e_cmn.MIRE, 'UPDATING _last_merged: %r' % bkmk)
        self._last_merged = bkmk
        return bkmk


    def writtenSinceLast(self, va):
        pageva = va & self._mem.pagemask

        # check if the page is marked dirty
        if not pageva in self._mem.pagedirty:
            return False

        # search through the writes to see if this location has been written to
        for idx in range(self.last_bookmark, len(self.wlog)):
            cva, tva, wdata = self.wlog[idx]
            if tva > va:
                continue
            if va >= (tva+len(wdata)):
                continue

            # if we reach here, va is contained in this write
            print("writtenSinceLast: va: %x  writelog: %x->%x->%r" % (va, cva, tva, wdata))
            return True

        return False


    def applyNewMemCache(self):
        newmc = e_mem.MemoryCache(self._mem)
        bkmk = self.setBookmark()
        newmc.bkmk = bkmk
        self._mem_cache_bkmks.append((bkmk, newmc))
        self._mem = newmc
        return newmc

    def _merge(self):
        newmc = self.applyNewMemCache()




class WriteTracker_2:
    '''
    Track and access the writelog from the current WorkspaceEmulator.

    On checking the writelog, if it exceeds a certain max size, it's sorted into pages for faster
    lookups.

    Also, allows bookmarks for tracking of events.  This allows us to check if memory has been 
    written to *since* a particular bookmark

    ---
    # 2
    * use MemoryCache, apply writes at each point where execution enters a modified byte.
    * check dirty cache page, followed by the true check of all writes since last bookmark
    '''
    def __init__(self, emu, maxlogsz=1000):
        raise Exception("NOT YET IMPLEMENTED")
        self.emu = emu
        self.maxlogsz = maxlogsz
        self.bookmarks = [0]
        self.last_bookmark = 0

        self._last_merged = 0
        self._mem = e_mem.MemoryCache(emu)
        self._mem_cache_bkmks = [(0, self._mem)]

    def setBookmark(self):
        bkmk = len(self.emu.path[0].get('writelog'))
        self.bookmarks.append(bkmk)
        self.last_bookmark = bkmk       # current index in write stream

        # apply writes to the memcache
        for idx in range(bkmk - self._last_merged):
            memwrite = self.emu.path[2]['writelog'][idx]
            cva, tva, wdata = memwrite
            self._mem.writeMemory(tva, wdata)

        return bkmk


    def writtenSinceBookmark(self, va, bkmkidx=-1):
        # check if writelog exceeds maxlogsz and 
        pageva = va & self._mem.pagemask

        # check if the page is marked dirty
        if not pageva in self._mem.pagedirty:
            return False


    def applyNewMemCache(self):
        newmc = e_mem.MemoryCache(self._mem)
        bkmk = self.setBookmark()
        newmc.bkmk = bkmk
        self._mem_cache_bkmks.append((bkmk, newmc))
        self._mem = newmc
        return newmc

    def _merge(self):
        newmc = self.applyNewMemCache()

class WriteTracker_3:
    '''
    Track and access the writelog from the current WorkspaceEmulator.

    On checking the writelog, if it exceeds a certain max size, it's sorted into pages for faster
    lookups.

    Also, allows bookmarks for tracking of events.  This allows us to check if memory has been 
    written to *since* a particular bookmark




    ---
    a couple approaches:
    # 1
    * build layers of MemoryCaches, never apply writes
    * this allows rolling back/tracking which phase established which thing
    * check dirty cache page, followed by writes since bookmark

    # 2
    * use MemoryCache, apply writes at each point where execution enters a modified byte.
    * check dirty cache page, followed by the true check of all writes since last bookmark

    # 3
    * just use the memwrite log with bookmarks.  this is simple
    * check all writes since last bookmark.  bookmark at each execution of modified byte.
    '''
    def __init__(self, emu, maxlogsz=1000):
        raise Exception("NOT YET IMPLEMENTED")
        self.emu = emu
        self.maxlogsz = maxlogsz
        self.bookmarks = [0]
        self.last_bookmark = 0

        self._last_merged = 0
        self._mem = e_mem.MemoryCache(emu)
        self._mem_cache_bkmks = [(0, self._mem)]

    def setBookmark(self):
        bkmk = len(self.emu.path[0].get('writelog'))
        self.bookmarks.append(bkmk)
        self.last_bookmark = bkmk       # current index in write stream

        # apply writes to the memcache
        for idx in range(bkmk - self._last_merged):
            memwrite = self.emu.path[2]['writelog'][idx]
            cva, tva, wdata = memwrite
            self._mem.writeMemory(tva, wdata)

        return bkmk


    def writtenSinceBookmark(self, va, bkmkidx=-1):
        # check if writelog exceeds maxlogsz and 
        pageva = va & self._mem.pagemask

        # check if the page is marked dirty
        if not pageva in self._mem.pagedirty:
            return False


    def applyNewMemCache(self):
        newmc = e_mem.MemoryCache(self._mem)
        bkmk = self.setBookmark()
        newmc.bkmk = bkmk
        self._mem_cache_bkmks.append((bkmk, newmc))
        self._mem = newmc
        return newmc

    def _merge(self):
        newmc = self.applyNewMemCache()

