Fix breakpoint oneshot, delete access violation

Author: dshikashio

pybag/dbgeng/breakpoints.py

@@ -36,16 +36,19 @@ def __call__(self, *args):
         else:
             ret = DbgEng.DEBUG_STATUS_NO_CHANGE
 
+        if rawbp.GetFlags() & DbgEng.DEBUG_BREAKPOINT_ONE_SHOT:
+            self._remove_stale(bpid)
+
         count -= 1
         if count == 0:
-            # We can disable here and but not delete
-            rawbp.RemoveFlags(DbgEng.DEBUG_BREAKPOINT_ENABLED)
+            self._remove(rawbp, bpid)
+
         return ret
 
+
     def set(self, expr, handler=None, type=None, windbgcmd=None, oneshot=False,
                   passcount=None, threadid=None, size=None, access=None,
                   count=0xffffffff):
-
         existingid = self.find(expr)
         if existingid != -1:
             # Just enable
@@ -70,9 +73,7 @@ def set(self, expr, handler=None, type=None, windbgcmd=None, oneshot=False,
         if threadid:
             bp.SetMatchThreadId(threadid)
         if oneshot:
-            count = 1
-            # comtypes can't properly handle true one shot
-            #bp.AddFlags(DbgEng.DEBUG_BREAKPOINT_ONE_SHOT)
+            bp.AddFlags(DbgEng.DEBUG_BREAKPOINT_ONE_SHOT)
         self._bp[id] = (count,handler)
         if not bp.GetFlags() & DbgEng.DEBUG_BREAKPOINT_DEFERRED:
             self.enable(id)
@@ -121,7 +122,7 @@ def find_expr(self, expr):
             try:
                 bp = self._control.GetBreakpointById(bpid)
             except exception.E_NOINTERFACE_Error:
-                self.breakpoints._remove_stale(bpid)
+                self._remove_stale(bpid)
                 continue
 
             bpexpr = bp.GetOffsetExpression()
@@ -137,7 +138,7 @@ def find_offset(self, offset):
             try:
                 bp = self._control.GetBreakpointById(bpid)
             except exception.E_NOINTERFACE_Error:
-                self.breakpoints._remove_stale(bpid)
+                self._remove_stale(bpid)
                 continue
 
             bpoff = bp.GetOffset()

pybag/dbgeng/idebugbreakpoint.py

@@ -4,10 +4,15 @@
 from . import core as DbgEng
 from . import exception
 
+def fakeRelease(*args):
+    # we don't want to reference the object
+    return
+
 class DebugBreakpoint(object):
     def __init__(self, bp):
         self._bp = bp
         exception.wrap_comclass(self._bp)
+        self._bp.Release = fakeRelease
 
     # IDebugBreakpoint
 

pybag/pydbg.py

@@ -540,7 +540,7 @@ def be(self, id):
         self.breakpoints.enable(id)
 
     def bp(self, expr=None, handler=None, windbgcmd=None, oneshot=False,
-            passcount=None, threadid=None):
+            passcount=None, threadid=None, count=0xffffffff):
         """bp(expr,handler,windbgcmd) -> Breakpoint on expression"""
         if expr is None:
             expr = self.reg.get_pc()
@@ -554,7 +554,8 @@ def bp(self, expr=None, handler=None, windbgcmd=None, oneshot=False,
                                     windbgcmd,
                                     oneshot,
                                     passcount,
-                                    threadid)
+                                    threadid,
+                                    count=count)
 
     def ba(self, expr=None, handler=None, windbgcmd=None, oneshot=False,
             passcount=None, threadid=None, size=None, access=None):