Add print all chunks in vis_heap_chunks (#1604)

* Changing the arguments to vis_heap_chunks to be clearer

1. --native to --beyond_top
2. --display_all to --no_truncate

* Add print all chunks to vis_heap_chunks

* Preventing the use of the all_chunks argument together with the count argument in vis_heap_chunks

* Use linting for heap.py

* Fix test_vis_heap_chunks.py
According to cdd71a1d82 --display_all/-d moved to --no_truncate/-n

---------

Co-authored-by: Nerya Zadkani <nerya@tokagroup.com>

pwndbg/commands/heap.py

@@ -697,7 +697,8 @@ parser = argparse.ArgumentParser(
 
 Default to the current arena's active heap.""",
 )
-parser.add_argument(
+group = parser.add_mutually_exclusive_group()
+group.add_argument(
     "count",
     nargs="?",
     type=lambda n: max(int(n, 0), 1),
@@ -706,26 +707,35 @@ parser.add_argument(
 )
 parser.add_argument("addr", nargs="?", default=None, help="Address of the first chunk.")
 parser.add_argument(
-    "--naive",
-    "-n",
+    "--beyond_top",
+    "-b",
     action="store_true",
     default=False,
     help="Attempt to keep printing beyond the top chunk.",
 )
 parser.add_argument(
-    "--display_all",
-    "-a",
+    "--no_truncate",
+    "-n",
     action="store_true",
     default=False,
     help="Display all the chunk contents (Ignore the `max-visualize-chunk-size` configuration).",
 )
+group.add_argument(
+    "--all_chunks",
+    "-a",
+    action="store_true",
+    default=False,
+    help=" Display all chunks (Ignore the default-visualize-chunk-number configuration).",
+)
 
 
 @pwndbg.commands.ArgparsedCommand(parser, category=CommandCategory.HEAP)
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.OnlyWithResolvedHeapSyms
 @pwndbg.commands.OnlyWhenHeapIsInitialized
-def vis_heap_chunks(addr=None, count=None, naive=None, display_all=None) -> None:
+def vis_heap_chunks(
+    addr=None, count=None, beyond_top=None, no_truncate=None, all_chunks=None
+) -> None:
     """Visualize chunks on a heap, default to the current arena's active heap."""
     allocator = pwndbg.heap.current
 
@@ -749,8 +759,12 @@ def vis_heap_chunks(addr=None, count=None, naive=None, display_all=None) -> None
     cursor_backup = cursor
     chunk = Chunk(cursor)
 
-    for _ in range(count + 1):
-        # Don't read beyond the heap mapping if --naive or corrupted heap.
+    chunk_id = 0
+    while True:
+        if not all_chunks and chunk_id == count + 1:
+            break
+
+        # Don't read beyond the heap mapping if --beyond_top or corrupted heap.
         if cursor not in heap_region:
             chunk_delims.append(heap_region.end)
             break
@@ -764,12 +778,13 @@ def vis_heap_chunks(addr=None, count=None, naive=None, display_all=None) -> None
         else:
             chunk_delims.append(cursor)
 
-        if (chunk.is_top_chunk and not naive) or (cursor == heap_region.end - ptr_size * 2):
+        if (chunk.is_top_chunk and not beyond_top) or (cursor == heap_region.end - ptr_size * 2):
             chunk_delims.append(cursor + ptr_size * 2)
             break
 
         cursor += chunk.real_size
         chunk = Chunk(cursor)
+        chunk_id += 1
 
     # Build the output buffer, changing color at each chunk delimiter.
     # TODO: maybe print free chunks in bold or underlined
@@ -821,7 +836,7 @@ def vis_heap_chunks(addr=None, count=None, naive=None, display_all=None) -> None
         while cursor != stop:
             # skip the middle part of a huge chunk
             if (
-                not display_all
+                not no_truncate
                 and half_max_size > 0
                 and begin_addr + half_max_size <= cursor < end_addr - half_max_size
             ):

tests/gdb-tests/tests/heap/test_vis_heap_chunks.py

@@ -216,12 +216,12 @@ def test_vis_heap_chunk_command(start_binary):
     for omitted_line in omitted_result:
         assert omitted_line in default_result or set(omitted_line) == {"."}
 
-    display_all_result = gdb.execute("vis_heap_chunk -a", to_string=True).splitlines()
-    assert display_all_result == default_result
+    no_truncate_result = gdb.execute("vis_heap_chunk -n", to_string=True).splitlines()
+    assert no_truncate_result == default_result
 
     del default_result
     del omitted_result
-    del display_all_result
+    del no_truncate_result
 
     # Continue, mock overflow changing the chunk size
     gdb.execute("continue")