simplify code.

hugsy · Dec 16, 2023 · e9e2106 · e9e2106
1 parent 9200420
commit e9e2106
Showing 1 changed file with 13 additions and 12 deletions.
diff --git a/gef.py b/gef.py
@@ -6277,8 +6277,7 @@ def do_invoke(self, _: List[str], **kwargs: Any) -> None:
         return
 
 class GlibcHeapChunkSummary:
-    def __init__(self, desc):
-        self.desc = desc
+    def __init__(self):
         self.count = 0
         self.total_size = 0
 
@@ -6290,15 +6289,15 @@ class GlibcHeapArenaSummary:
     def __init__(self) -> None:
         self.size_distribution = {}
         self.flag_distribution = {
-            "PREV_INUSE": GlibcHeapChunkSummary("PREV_INUSE"),
-            "IS_MMAPPED": GlibcHeapChunkSummary("IS_MMAPPED"),
-            "NON_MAIN_ARENA": GlibcHeapChunkSummary("NON_MAIN_ARENA")
+            "PREV_INUSE": GlibcHeapChunkSummary(),
+            "IS_MMAPPED": GlibcHeapChunkSummary(),
+            "NON_MAIN_ARENA": GlibcHeapChunkSummary()
         }
 
     def process_chunk(self, chunk: GlibcChunk) -> None:
         per_size_summary = self.size_distribution.get(chunk.size, None)
         if per_size_summary is None:
-            per_size_summary = GlibcHeapChunkSummary(str(chunk.size))
+            per_size_summary = GlibcHeapChunkSummary()
             self.size_distribution[chunk.size] = per_size_summary
         per_size_summary.process_chunk(chunk)
 
@@ -6312,13 +6311,13 @@ def process_chunk(self, chunk: GlibcChunk) -> None:
     def print(self) -> None:
         gef_print(f"== Chunk distribution by size ==")
         gef_print("{:<10s}\t{:<10s}\t{:s}".format("ChunkSize", "Count", "TotalSize"))
-        for _, chunk_summary in sorted(self.size_distribution.items(), key=lambda x: x[1].total_size, reverse=True):
-            gef_print("{:<10s}\t{:<10d}\t{:<d}".format(chunk_summary.desc, chunk_summary.count, chunk_summary.total_size))
+        for chunk_size, chunk_summary in sorted(self.size_distribution.items(), key=lambda x: x[1].total_size, reverse=True):
+            gef_print("{:<10d}\t{:<10d}\t{:<d}".format(chunk_size, chunk_summary.count, chunk_summary.total_size))
 
         gef_print(f"\n== Chunk distribution by flag ==")
         gef_print("{:<15s}\t{:<10s}\t{:s}".format("Flag", "TotalCount", "TotalSize"))
-        for _, chunk_summary in self.flag_distribution.items():
-            gef_print("{:<15s}\t{:<10d}\t{:<d}".format(chunk_summary.desc, chunk_summary.count, chunk_summary.total_size))
+        for chunk_flag, chunk_summary in self.flag_distribution.items():
+            gef_print("{:<15s}\t{:<10d}\t{:<d}".format(chunk_flag, chunk_summary.count, chunk_summary.total_size))
 
         gef_print("")
 
@@ -6376,8 +6375,10 @@ def dump_chunks_heap(self, start: int, end: int, arena: GlibcArena, allow_unalig
         chunk_iterator = GlibcChunk(start, from_base=True, allow_unaligned=allow_unaligned)
         heap_summary = GlibcHeapArenaSummary()
         for chunk in chunk_iterator:
+            is_heap_corrupted = chunk.base_address > end
+
             if summary:
-                if chunk.base_address > end:
+                if is_heap_corrupted:
                     err("Corrupted heap, cannot continue.")
                     return False
 
@@ -6388,7 +6389,7 @@ def dump_chunks_heap(self, start: int, end: int, arena: GlibcArena, allow_unalig
                         f"{chunk!s} {LEFT_ARROW} {Color.greenify('top chunk')}")
                     break
 
-                if chunk.base_address > end:
+                if is_heap_corrupted:
                     err("Corrupted heap, cannot continue.")
                     return False