Improve probeleak command (#698)

* Feat: Enhance probeleak command

* fix: probeleak defalut stop_cnt

* fix: probe_leak arg's doc

* probeleak: add flags arg; change args name

pwndbg/commands/probeleak.py

@@ -5,6 +5,7 @@ from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import os
 import argparse
 import math
 
@@ -32,28 +33,51 @@ def find_module(addr, max_distance):
 
     return pages[-1]
 
+def satisfied_flags(require_flags, flags):
+    return (require_flags & ~(flags)) == 0
+
+def flags_str2int(flags_s):
+    flag_i = 0
+    if 'r' in flags_s:
+        flag_i |= os.R_OK
+    if 'w' in flags_s:
+        flag_i |= os.W_OK
+    if 'x' in flags_s:
+        flag_i |= os.X_OK
+    return flag_i
+
 parser = argparse.ArgumentParser(description='''
 Pointer scan for possible offset leaks.
 Examples:
     probeleak $rsp 0x64 - leaks 0x64 bytes starting at stack pointer and search for valid pointers
-    probeleak $rsp 0x64 0x10 - as above, but pointers may point 0x10 bytes outside of memory page
+    probeleak $rsp 0x64 --max-dist 0x10 - as above, but pointers may point 0x10 bytes outside of memory page
+    probeleak $rsp 0x64 --point-to libc --max-ptrs 1 --flags rwx - leaks 0x64 bytes starting at stack pointer and 
+        search for one valid pointer which points to a libc rwx page
 ''')
 parser.add_argument('address', nargs='?', default='$sp',
                     help='Leak memory address')
 parser.add_argument('count', nargs='?', default=0x40,
                     help='Leak size in bytes')
-parser.add_argument('max_distance', nargs='?', default=0x0,
+parser.add_argument('--max-distance', type=int, default=0x0,
                     help='Max acceptable distance between memory page boundry and leaked pointer')
+parser.add_argument('--point-to', type=str, default=None,
+                    help='Mapping name of the page that you want the pointers point to')
+parser.add_argument('--max-ptrs', type=int, default=0,
+                    help='Stop search after find n pointers, default 0')
+parser.add_argument('--flags', type=str, default=None,
+                    help='flags of the page that you want the pointers point to. [e.g. rwx]')
 
 @pwndbg.commands.ArgparsedCommand(parser)
 @pwndbg.commands.OnlyWhenRunning
-def probeleak(address=None, count=0x40, max_distance=0x0):
+def probeleak(address=None, count=0x40, max_distance=0x0, point_to=None, max_ptrs=0, flags=None):
 
     address = int(address)
     address &= pwndbg.arch.ptrmask
     ptrsize = pwndbg.arch.ptrsize
     count   = max(int(count), ptrsize)
     off_zeros = int(math.ceil(math.log(count,2)/4))
+    if flags != None:
+        require_flags = flags_str2int(flags)
 
     if count > address > 0x10000: # in case someone puts in an end address and not a count (smh)
         print(message.warn("Warning: you gave an end address, not a count. Substracting 0x%x from the count." % (address)))
@@ -70,10 +94,15 @@ def probeleak(address=None, count=0x40, max_distance=0x0):
         return
 
     found = False
+    find_cnt = 0
     for i in range(0, len(data) - ptrsize + 1):
         p = pwndbg.arch.unpack(data[i:i+ptrsize])
         page = find_module(p, max_distance)
         if page:
+            if point_to != None and point_to not in page.objfile:
+                continue
+            if flags != None and not satisfied_flags(require_flags, page.flags):
+                continue
             if not found:
                 print(M.legend())
                 found = True
@@ -96,8 +125,11 @@ def probeleak(address=None, count=0x40, max_distance=0x0):
             symbol = pwndbg.symbol.get(p)
             if symbol:
                 text += ' (%s)' % symbol
-
             print(text)
 
+            find_cnt += 1
+            if max_ptrs != 0 and find_cnt >= max_ptrs:
+                break
+
     if not found:
         print(message.hint('No leaks found at 0x%x-0x%x :(' % (address, address+count)))