Basic implementation of heap commands (#36)

* Basic implementation of heap commands

* Optionally allow a user to specify the main_arena address

* Fix coloring, walk heap to find top_chunk

* Remove unnecessary casts

* Added docs and constructed values from types

pwndbg/chain.py

@@ -9,7 +9,7 @@ import pwndbg.vmmap
 
 LIMIT = 5
 
-def get(address, limit=LIMIT):
+def get(address, limit=LIMIT, offset=0):
     """
     Recursively dereferences an address.
 
@@ -24,15 +24,15 @@ def get(address, limit=LIMIT):
 
         result.append(address)
         try:
-            address = int(pwndbg.memory.poi(pwndbg.typeinfo.ppvoid, address))
+            address = int(pwndbg.memory.poi(pwndbg.typeinfo.ppvoid, address + offset))
         except gdb.MemoryError:
             break
 
     return result
 
 
-def format(value, limit=LIMIT, code=True):
-    chain = get(value, limit)
+def format(value, limit=LIMIT, code=True, offset=0):
+    chain = get(value, limit, offset)
 
     # Enhance the last entry
     # If there are no pointers (e.g. eax = 0x41414141), then enhance

pwndbg/color.py

@@ -55,7 +55,7 @@ def get(address, text = None):
     if text is None and isinstance(address, (long, int)) and address > 255:
         text = hex(int(address))
     if text is None:
-        text = int(address)
+        text = str(int(address))
 
     if color == NORMAL:
         return text
@@ -71,4 +71,3 @@ def legend():
         UNDERLINE + 'RWX' + NORMAL,
         'RODATA'
     ))
-

pwndbg/commands/defcon.py

@@ -10,7 +10,7 @@ from pwndbg.color import bold, blue, green, red
 
 @pwndbg.commands.Command
 @pwndbg.commands.OnlyWhenRunning
-def heap(addr=0x2aaaaaad5000):
+def defcon_heap(addr=0x2aaaaaad5000):
 # def heap(addr=0x2aaaaaaaf000):
     free = []
 
@@ -59,7 +59,7 @@ def heap_freebins(addr=0x0602558):
 
         print('    %#x %#x %s' % (addr, size, '*' if in_use else ''))
         addr = bk
-    
+
     print()
     return free
 

pwndbg/commands/heap.py

@@ -1,46 +1,171 @@
 #!/usr/bin/env python
-# -*- coding: utf-8 -*-
-"""
-Heap commands.
-"""
+
 from __future__ import print_function
-import argparse
 import gdb
+
 import pwndbg.commands
 
+from pwndbg.color import bold, yellow, red, underline
+
+PREV_INUSE = 1
+IS_MMAPED = 2
+NON_MAIN_ARENA = 4
+
+def value_from_type(type_name, addr):
+    gdb_type = pwndbg.typeinfo.load(type_name)
+    return gdb.Value(addr).cast(gdb_type.pointer()).dereference()
+
+def get_main_arena(addr=None):
+    if addr == None:
+        main_arena = gdb.lookup_symbol('main_arena')[0].value()
+    else:
+        main_arena = value_from_type('struct malloc_state', addr)
+
+    if main_arena == None:
+        print(red('Symbol \'main_arena\' not found. Try installing libc ' \
+                  'debugging symbols or specifying the main arena address ' \
+                  'and try again'))
+
+    return main_arena
+
+def get_heap_bounds():
+    page = None
+    for m in pwndbg.vmmap.get():
+        if m.objfile == '[heap]':
+            page = m
+            break
+
+    if m != None:
+        return (m.vaddr, m.vaddr + m.memsz)
+    else:
+        return (None, None)
+
+@pwndbg.commands.ParsedCommand
+@pwndbg.commands.OnlyWhenRunning
+def heap(addr=None):
+    """
+    Prints out all chunks in the main_arena, or the arena specified by `addr`.
+    """
+    main_arena = get_main_arena(addr)
+    if main_arena == None:
+        return
+
+    heap_base = get_heap_bounds()[0]
+    if heap_base == None:
+        print(red('Could not find the heap'))
+        return
+
+    top = main_arena['top']
+    last_remainder = main_arena['last_remainder']
+
+    print(bold('Top Chunk: ') + pwndbg.color.get(top))
+    print(bold('Last Remainder: ') + pwndbg.color.get(last_remainder))
+    print()
+
+    # Print out all chunks on the heap
+    # TODO: Add an option to print out only free or allocated chunks
+    addr = heap_base
+    while addr <= top:
+        chunk = malloc_chunk(addr)
+        size = int(chunk['size'])
+
+        # Clear the bottom 3 bits
+        size &= ~7
+        addr += size
+
 @pwndbg.commands.ParsedCommand
 @pwndbg.commands.OnlyWhenRunning
-def brk(n=0):
-    '''Get the address of brk(n=0)'''
-    gdb.execute('call brk(%i)' % n)
+def arena(addr=None):
+    """
+    Prints out the main arena or the arena at the specified by address.
+    """
+    main_arena = get_main_arena(addr)
+    if main_arena == None:
+        return
+
+    print(main_arena)
 
 @pwndbg.commands.ParsedCommand
 @pwndbg.commands.OnlyWhenRunning
-def sbrk(n=0):
-    '''Get the address of sbrk(n=0)'''
-    gdb.execute('call sbrk(%i)' % n)
+def bins(addr=None):
+    """
+    Prints out the contents of the fastbins of the main arena or the arena
+    at the specified address.
+    """
+    main_arena = get_main_arena(addr)
+    if main_arena == None:
+        return
+
+    fastbins = main_arena['fastbinsY']
+    bins = main_arena['bins']
 
+    size_t_size = pwndbg.typeinfo.load('size_t').sizeof
+    num_fastbins = int(fastbins.type.sizeof / fastbins.type.target().sizeof)
+    num_bins = int(bins.type.sizeof / bins.type.target().sizeof)
+    fd_field_offset = 2 * size_t_size
 
+    print(underline(yellow('fastbins')))
+    for i in range(num_fastbins):
+        size = 2 * size_t_size * (i + 1)
+        chain = pwndbg.chain.format(int(fastbins[i]), offset=fd_field_offset)
+        print((bold(size) + ': ').ljust(13) + chain)
 
+    # TODO: Print other bins
 
-p = argparse.ArgumentParser(prog='hheap')
+@pwndbg.commands.ParsedCommand
+@pwndbg.commands.OnlyWhenRunning
+def top_chunk(addr=None):
+    """
+    Prints out the address of the top chunk of the main arena, or of the arena
+    at the specified address.
+    """
+    main_arena = get_main_arena(addr)
+    if main_arena == None:
+        heap_start, heap_end = get_heap_bounds()
+        if heap_start == None:
+            print(red('Could not find the heap'))
+            return
+
+        # If we don't know where the main_arena struct is, just iterate
+        # through all the heap objects until we hit the last one
+        last_addr = None
+        addr = heap_start
+        while addr < heap_end:
+            chunk = value_from_type('struct malloc_chunk', addr)
+            size = int(chunk['size'])
+
+            # Clear the bottom 3 bits
+            size &= ~7
 
-p.add_argument('--size',
-               help='Heap size.  May be expressed as an integer or range (e.g. 32-64).')
-p.add_argument('--verbose', action='store_true',
-               help='Print more information')
-p.add_argument('--free', action='store_true',
-               help='Only show free slots')
-p.add_argument('address', type=int, default=0,
-               help='Heap allocation to display')
+            last_addr = addr
+            addr += size
+        print(pwndbg.color.get(last_addr))
+    else:
+        print(pwndbg.color.get(main_arena['top']))
 
-@pwndbg.commands.Command
+@pwndbg.commands.ParsedCommand
 @pwndbg.commands.OnlyWhenRunning
-def hheap(*a):
-    """Prints out heap information.
-    """ + p.format_help()
-    try:
-        args = p.parse_args(a)
-    except SystemExit:
-        return
+def malloc_chunk(addr):
+    """
+    Prints out the malloc_chunk at the specified address.
+    """
+    if not isinstance(addr, int):
+        addr = int(addr)
+
+    chunk = value_from_type('struct malloc_chunk', addr)
+    size = int(chunk['size'])
+    prev_inuse = (size & PREV_INUSE) == 1
+    is_mmaped = (size & IS_MMAPED) == 1
+    non_main_arena = (size & NON_MAIN_ARENA) == 1
+
+    header = pwndbg.color.get(addr)
+    if prev_inuse:
+        header += yellow(' PREV_INUSE')
+    if is_mmaped:
+        header += yellow(' IS_MMAPED')
+    if non_main_arena:
+        header += yellow(' NON_MAIN_ARENA')
+    print(header)
+    print(chunk)
 
+    return chunk