# Copyright 2005 Cody Brocious.  Released under the terms of the GNU GPL v2

# Huge thanks to Kai Killebrew for the idea of using an infinite loop, leading to this idea.

import process, struct, time

class ProcessWatcher:
    def __init__(self, name=None, pid=None, allocate=True, code_address=None, data_address=None):
        self.p = process.Process(name=name, pid=pid)
        self.code_address = code_address
        self.data_address = data_address
        self.allocate = allocate
        if self.allocate:
            if not code_address or not data_address:
                allocated = self.p.allocatePages(4096)
                if not data_address:
                    self.data_address = allocated
                    allocated += 5
                if not code_address:
                    self.code_address = allocated
        self.watches = {}
        self.watchcallbacks = [None]
    def addEbpWatch(self, where, code_len, var, callback):
        self.watches[where] = (code_len, '\x8B\x8D' + struct.pack('<L',  var)) # mov ecx, [ebp+<var>]
        self.watchcallbacks.append(callback)
    def addEbpRefWatch(self, where, code_len, var, callback):
        self.watches[where] = (code_len, '\x8D\x8D' + struct.pack('<L', var)) # lea ecx, [ebp+<var>]
        self.watchcallbacks.append(callback)
    def addDataWatch(self, where, code_len, address, callback):
        self.watches[where] = (code_len, '\x8B\x0D' + struct.pack('<L', address)) # mov ecx, dword ptr ds:[<address>]
        self.watchcallbacks.append(callback)
    def doWatch(self):
        bp = False
        while not bp:
            while True:
                d = self.p.read(self.data_address, 1)
                if not d:
                    return False
                d = ord(d)
                if d == 0 or not d & 0x80L:
                    time.sleep(0.025)
                    continue
                else:
                    if not self.watchcallbacks[0xFF - d](struct.unpack('<L', self.p.read(self.data_address+1, 4))[0]):
                        bp = True
                    break
            self.p.write(self.data_address, '\x00')
    def injectWatches(self):
        i = 0
        self.p.write(self.data_address, '\x00')

        off = self.code_address
        jt = []
        for where in self.watches:
            i += 1
            
            old_len, code = self.watches[where]
            old = self.p.read(where, old_len)
            mc  = '\xE9' + struct.pack('<l', off - (where + 5))                 # jmp <code_address>
            mc += '\x90' * (old_len - 5)                                        # nop
            self.p.write(where, mc)

            mc  = '\x9C'                                                        # pushfd
            mc += '\x50'                                                        # push eax
            mc += '\x51'                                                        # push ecx
            mc += code                                                          # watch code
            mc += '\xB8' + struct.pack('<L', off + len(mc) + 17)                # mov eax, <REENTRY>
            mc += '\xC6\x05' + struct.pack('<L', self.data_address) + chr(i)    # mov byte ptr ds:[<data address>], <i>
            mc += '\xE9\xDE\xAD\xBE\xEF'                                        # jmp <main>

            mc += '\x58'                                                        # REENTRY: pop eax
            mc += '\x9D'                                                        # popfd
            mc += self.recalculateOffsets(code=old,
                                          old_loc=where,
                                          new_loc=off+len(mc))                  # old code with offsets recalculated if necessary
            off += len(mc) + 5
            mc += '\xE9' + struct.pack('<l', (where + old_len) - off)           # jmp <call>

            jt.append(mc)

        main = self.code_address + len(''.join(jt))

        off = self.code_address
        for mc in jt:
            mc = mc.replace('\xDE\xAD\xBE\xEF', struct.pack('<L', main - (off + mc.find('\xDE\xAD\xBE\xEF') + 4)))
            self.p.write(off, mc)
            off += len(mc)

        mc  = '\x89\x0D' + struct.pack('<L', self.data_address+1)               # mov dword ptr ds:[<data address+1>], ecx
        mc += '\x59'                                                            # pop ecx
        mc += '\x50'                                                            # push eax
        mc += '\xF6\x15' + struct.pack('<L', self.data_address)                 # not byte ptr ds:[<data address>]
        mc += '\xA0' + struct.pack('<L', self.data_address)                     # START: mov al, byte ptr ds:[<data address>]
        mc += '\x3C\x00'                                                        # cmp al, 00h
        mc += '\x74\x02'                                                        # je END
        mc += '\xEB\xF5'                                                        # jmp START
        mc += '\x58'                                                            # END: pop eax
            
        mc += '\xFF\xE0'                                                        # jmp eax
                                       
        self.p.write(main, mc)

    def recalculateOffsets(self, code, old_loc, new_loc):
        if len(code) == 5 and (code[0] == '\xE8' or code[0] == '\xE9'):
            code, offset = code[0], struct.unpack('<l', code[1:])[0]
        elif len(code) == 6 and struct.unpack('<H', code[0:2])[0] in (0x0F84, 0x0F85, 0x0F82, 0x0F86, 0x0F8C, 0x0F8E, 0x0F87, 0x0F83, 0x0F8F, 0x0F8D):
            code, offset = code[0:2], struct.unpack('<l', code[2:])[0]
        else:
            return code
        return code + struct.pack('<l', (old_loc + offset) - new_loc)