Add get_memory_breakpoint_at to debugger + fix token_user in process + add prot param to virtual_alloc

windows/debug/debugger.py

@@ -406,24 +406,24 @@ class Debugger(object):
 
         fault_page = (fault_addr >> 12) << 12
 
-        #print("FAULT AT {0:#x} ({1})".format(fault_addr, fault_type))
-        if fault_page not in self._watched_pages[self.current_process.pid]:
+        mem_bp = self.get_memory_breakpoint_at(fault_addr, self.current_process)
+        if mem_bp is False: # No BP on this page
             return self.on_exception(exception)
-
-        for bp in self._watched_pages[self.current_process.pid][fault_page]:
-            if bp._addr <= fault_addr < bp._addr + bp.size:
-                # TODO: restore all page to real state :)
-                continue_flag = bp.trigger(self, exception)
-                self._explicit_single_step[self.current_thread.tid] = self.current_thread.context.EEFlags.TF
-                # If BP has not been removed in trigger, pas it
-                if bp in self._watched_pages[self.current_process.pid][fault_page]:
-                    self._pass_memory_breakpoint(bp, fault_page)
-                return continue_flag
-        else:
-            # If no BP on this page handle the fault address
+        if mem_bp is None: # Page as MEMBP but None handle this address
+            # This hack is bad, find a BP on the page to restore original access..
+            # TODO: stock original page protection elsewhere ?
+            bp = self._watched_pages[self.current_process.pid][fault_page][0]
             self._pass_memory_breakpoint(bp, fault_page)
             return DBG_CONTINUE
 
+        continue_flag = mem_bp.trigger(self, exception)
+        self._explicit_single_step[self.current_thread.tid] = self.current_thread.context.EEFlags.TF
+        # If BP has not been removed in trigger, pas it
+        if mem_bp in self._watched_pages[self.current_process.pid][fault_page]:
+            self._pass_memory_breakpoint(mem_bp, fault_page)
+        return continue_flag
+
+
         #for bp, vprot_begin, vprot_end, original_prot in self._watched_memory:
         #    if vprot_begin <= fault_addr < vprot_end:
         #        # It's the page for this MEMBP that triggeed the BP
@@ -651,6 +651,26 @@ class Debugger(object):
         ctx.EEFlags.TF = 1
         t.set_context(ctx)
 
+    ## Memory Breakpoint helper
+    def get_memory_breakpoint_at(self, addr, process=None):
+        """Get the memory breakpoint the handle `addr`
+        Return values are:
+            * ``False`` if the page as no memory breakpoint (real fault)
+            * ``None`` if the page as memBP but None handle ``addr``
+            * ``bp`` the MemBP that handle ``addr``
+        """
+        if process is None:
+            process = self.current_process
+
+        fault_page = (addr >> 12) << 12
+        if fault_page not in self._watched_pages[process.pid]:
+            return False
+
+        for bp in self._watched_pages[process.pid][fault_page]:
+            if bp._addr <= addr < bp._addr + bp.size:
+                return bp
+        return None
+
     # Public callback
     def on_exception(self, exception):
         """Called on exception event other that known breakpoint or requested single step. ``exception`` is one of the following type:

windows/winobject/process.py

@@ -675,13 +675,13 @@ class CurrentProcess(Process):
         bits = platform.architecture()[0]
         return int(bits[:2])
 
-    def virtual_alloc(self, size):
+    def virtual_alloc(self, size, prot=PAGE_EXECUTE_READWRITE):
         """Allocate memory in the process
 
         :return: The address of the allocated memory
         :rtype: :class:`int`
 		"""
-        return winproxy.VirtualAlloc(dwSize=size)
+        return winproxy.VirtualAlloc(dwSize=size, flProtect=prot)
 
     def virtual_free(self, addr):
         """Free memory in the process by virtual_alloc"""
@@ -798,13 +798,13 @@ class WinProcess(Process):
             pass
         return '<{0} "{1}" pid {2} at {3}>'.format(self.__class__.__name__, self.name, self.pid, hex(id(self)))
 
-    def virtual_alloc(self, size):
+    def virtual_alloc(self, size, prot=PAGE_EXECUTE_READWRITE):
         """Allocate memory in the process
 
         :return: The address of the allocated memory
         :rtype: :class:`int`
 		"""
-        return winproxy.VirtualAllocEx(self.handle, dwSize=size)
+        return winproxy.VirtualAllocEx(self.handle, dwSize=size, flProtect=prot)
 
     def virtual_free(self, addr):
         """Free memory in the process by virtual_alloc"""
@@ -1004,7 +1004,7 @@ class Token(AutoHandle):
         buffer_size = self.get_required_information_size(TokenUser)
         buffer = ctypes.c_buffer(buffer_size)
         self.get_informations(TokenUser, buffer)
-        return ctypes.cast(ctypes.byref(buffer), POINTER(TOKEN_USER))[0]
+        return ctypes.cast(buffer, POINTER(TOKEN_USER))[0]
 
     @property
     def computername(self):
@@ -1024,7 +1024,7 @@ class Token(AutoHandle):
         username = ctypes.c_buffer(usernamesize.value)
         computername = ctypes.c_buffer(computernamesize.value)
         peUse = SID_NAME_USE()
-        winproxy.LookupAccountSidA(None, sid, username, byref(usernamesize), computername, byref(computernamesize), peUse)
+        winproxy.LookupAccountSidA(None, sid, username, usernamesize, computername, computernamesize, peUse)
         return username[:usernamesize.value], computername[:computernamesize.value]
 
     def get_informations(self, info_type, data):