clean that tool for windows HEAPs

scripts/haystack-find-heap.py

@@ -4,14 +4,17 @@
 import logging
 import os
 import sys
+import argparse
 
 from functools import partial
 
 import haystack
 from haystack import dump_loader
 from haystack import constraints
+from haystack import argparse_utils
 from haystack.search import searcher
 from haystack.outputters import text
+#from haystack.outputters import python
 
 """
 Search for HEAP.
@@ -31,6 +34,9 @@ def _init_module_name(self, memory_handler):
     def _init_constraints_filename(self, heap_module):
         raise NotImplementedError('_init_constraints_filename')
 
+    def _init_heap_record_name(self):
+        raise NotImplementedError('_init_heap_record_name')
+
     def search_heap(self, memdumpname):
         # we need a memory dump loader
         self.memory_handler = dump_loader.load(memdumpname)
@@ -50,12 +56,13 @@ def search_heap(self, memdumpname):
         # DEBUG PEB search
         #peb = my_model.import_module('haystack.structures.win32.winxp_32_peb')
         ##DEBUG
-
+        heap_record_name = self._init_heap_record_name()
+        heap_struct = getattr(heap_module, heap_record_name)
         # on ly return first results in each mapping
         results = []
         for mapping in self.memory_handler.get_mappings():
             log.debug("looking at %s", mapping)
-            res = my_searcher._search_in(mapping, heap_module.HEAP, nb=1, align=0x1000)
+            res = my_searcher._search_in(mapping, heap_struct, nb=1, align=0x1000)
             # DEBUG PEB search
             #res = my_searcher._search_in(mapping, peb.struct__PEB, nb=1, align=0x1000)
             if res:
@@ -81,7 +88,9 @@ def search_heap_direct(self, memdumpname, start_address_mapping):
                                                        my_constraints,
                                                        [m],
                                                        update_cb=partial(self.print_cb, self.memory_handler))
-        results = my_searcher._load_at(m, start_address_mapping, heap_module.HEAP, depth=5)
+        heap_record_name = self._init_heap_record_name()
+        heap_struct = getattr(heap_module, heap_record_name)
+        results = my_searcher._load_at(m, start_address_mapping, heap_struct, depth=5)
         #print haystack.output_to_python(memory_handler, [results])[0][0].toString()
         return results
 
@@ -115,6 +124,9 @@ def _init_module_name(self, memory_handler):
     def _init_constraints_filename(self, heap_module):
         return os.path.join(os.path.dirname(heap_module.__file__), 'win7heap.constraints')
 
+    def _init_heap_record_name(self):
+        return 'HEAP'
+
 class WinXPHeapFinder(HeapFinder):
     def _init_module_name(self, memory_handler):
         if 64 == memory_handler.get_target_platform().get_cpu_bits():
@@ -126,56 +138,58 @@ def _init_module_name(self, memory_handler):
     def _init_constraints_filename(self, heap_module):
         return os.path.join(os.path.dirname(heap_module.__file__), 'winxpheap.constraints')
 
-'''
-zeus.vmem.856.dump
-
-0x00090000-0x00190000
-0x00190000-0x001a0000
-0x001a0000-0x001b0000
-0x00350000-0x00360000
-0x003b0000-0x003c0000
-0x00c30000-0x00cb0000
-0x00d60000-0x00d70000
-0x00e20000-0x00e30000
-0x00e80000-0x00e90000
-0x7f6f0000-0x7f7f0000
-
-
+    def _init_heap_record_name(self):
+        return 'HEAP'
 
-DEBUG:basicmodel:ptr: UnusedUnCommittedRanges <class 'haystack.types.LP_4_struct__HEAP_UNCOMMMTTED_RANGE'>
-LP_4_struct__HEAP_UNCOMMMTTED_RANGE(3160606104) 0xbc630598 INVALID
-for f in `ls /home/jal/outputs/vol/zeus.vmem.1668.dump` ; do echo $f; xxd /home/jal/outputs/vol/zeus.vmem.1668.dump/$f | head | grep -c "ffee ffee" ; done | grep -B1 "1$"
-
-DEBUG:utils:obj._sub_addr_: 0xbc630598
+class LibcHeapFinder(HeapFinder):
+    def _init_module_name(self, memory_handler):
+        module_name = 'haystack.structures.libc.ctypes_malloc'
+        log.error("this doesn't not work on libc heap")
+        return module_name
 
-'''
+    def _init_constraints_filename(self, heap_module):
+        return os.path.join(os.path.dirname(heap_module.__file__), 'libcheap.constraints')
 
+    def _init_heap_record_name(self):
+        return 'malloc_chunk'
 
 
 def main(argv):
-    for f in [
-              #Win7HeapFinder(),
-              WinXPHeapFinder()
-             ]:
-        if len(argv) == 2:
-            memdumpname = argv[1]
-            if f.search_heap(memdumpname) is not None:
-                break
-        elif len(argv) == 3:
-            memdumpname, address = argv[1], int(argv[2], 16)
-            ret = f.search_heap_direct(memdumpname, address)
-            out = text.RecursiveTextOutputter(f.memory_handler)
-            print out.parse(ret[0]), ret[1]
-            break
-
-    #f.search_heap(memdumpname)
-    #f = Win7HeapFinder()
-    #f.search_heap_direct(memdumpname, address)
-
+    parser = argparse.ArgumentParser(prog='haystack-find-heap',
+                                          description="Find heaps in a dumpfile")
+    parser.add_argument('--host', action='store', default='winxp', help='winxp,win7')
+    parser.add_argument('dumpname', type=argparse_utils.readable, help='process memory dump name')
+    parser.add_argument('address', nargs='?', type=argparse_utils.int16, default=None, help='Load Heap from address (hex)')
+
+    opts = parser.parse_args(argv)
+
+    #if 'libc' == opts.host:
+    #    my_finder = LibcHeapFinder()
+    #el
+    if 'winxp' == opts.host:
+        my_finder = WinXPHeapFinder()
+    elif 'win7' == opts.host:
+        my_finder = Win7HeapFinder()
+    else:
+        raise ValueError('not such heap finder for %s' % opts.host)
+
+    memdumpname = opts.dumpname
+    if opts.address is None:
+        if my_finder.search_heap(memdumpname) is not None:
+            return
+    else:
+        address = opts.address
+        # just return the heap
+        ret = my_finder.search_heap_direct(memdumpname, address)
+        out = text.RecursiveTextOutputter(my_finder.memory_handler)
+        #out = python.PythonOutputter(my_finder.memory_handler)
+        print out.parse(ret[0], depth=4)
+        print 'Valid=', ret[1]
+        return
 
 
 if __name__ == '__main__':
     logging.basicConfig(level=logging.INFO)
     #logging.basicConfig(level=logging.DEBUG)
-    main(sys.argv)
+    main(sys.argv[1:])