metasploit-framework/lib/metasm/samples/dbg-plugins/heapscan.rb

#    This file is part of Metasm, the Ruby assembly manipulation suite
#    Copyright (C) 2006-2009 Yoann GUILLOT
#
#    Licence is LGPL, see LICENCE in the top-level directory


# metasm debugger plugin
# adds some heap_* functions to interract with the target heap chunks
# functions:
#  heap_scan, scan for malloc chunks in the heaps and xrefs between them
#  heap_scanstruct, scan for arrays/linkedlists in the chunk graph
#  heap_chunk [addr], display a chunk
#  heap_array [addr], display an array of chunks from their root
#  heap_list [addr], display a linkedlist
#  heap_strscan [str], scan the memory for a raw string, display chunks xrefs
#  heap_snap, make a snapshot of the currently displayed structure, hilight fields change


# use precompiled native version when available
$heapscan_dir = File.join(File.dirname(plugin_filename).gsub('\\', '/'), 'heapscan')
require File.join($heapscan_dir, 'heapscan')

fname = case OS.current.shortname
when 'linos'
  'compiled_heapscan_lin'
when 'winos'
  case OS.current.version[0]
  when 5; 'compiled_heapscan_win'
  when 6; 'compiled_heapscan_win7'
  end
end
fname = File.join($heapscan_dir, fname)
if not File.exist?(fname + '.so') and File.exist?(fname + '.c')
  puts "compiling native scanner..."
  exe = DynLdr.host_exe.compile_c_file(DynLdr.host_cpu, fname + '.c')
  DynLdr.compile_binary_module_hack(exe)
  exe.encode_file(fname + '.so', :lib)
end
require fname if File.exist?(fname + '.so')

def heapscan_time(s='')
  @heapscan_time ||= nil
  t = Time.now
  log s + ' %.2fs' % (t-@heapscan_time) if @heapscan_time and s != ''
  @heapscan_time = t
  Gui.main_iter if gui
end

def heap; @heap ; end
def heap=(h) ; @heap = h ; end

def heapscan_scan(xr=true)
  heaps = []
  mmaps = []
  libc = nil
  pr = os_process
  pr.mappings.each { |a, l, p, f|
    case f.to_s
    when /heap/
      heaps << [a, l]
    when /libc[^a-zA-Z]/
      libc ||= a if p == 'r-xp'
    when ''
      mmaps << [a, l]
    end
  }

  heapscan_time ''
  @disassembler.parse_c ''
  if pr and OS.current.name =~ /winos/i
    if OS.current.version[0] == 5
      @heap = WindowsHeap.new(self)
      @heap.cp = @disassembler.c_parser
      @heap.cp.parse_file File.join($heapscan_dir, 'winheap.h') unless @heap.cp.toplevel.struct['_HEAP']
    else
      @heap = Windows7Heap.new(self)
      @heap.cp = @disassembler.c_parser
      @heap.cp.parse_file File.join($heapscan_dir, 'winheap7.h') unless @heap.cp.toplevel.struct['_HEAP']
    end
    @heap.heaps = heaps
  else
    @heap = LinuxHeap.new(self)
    @heap.cp = @disassembler.c_parser
    @heap.mmaps = mmaps
    @heap.scan_libc(libc)
    heapscan_time "libc!main_arena #{'%x' % @heap.main_arena_ptr}"
  end

  hsz = 0
  (heaps + mmaps).each { |a, l|
    hsz += l
    @heap.range.update a => l
  }

  log "#{hsz/1024/1024}M heap"

  @heap.scan_chunks
  heapscan_time "#{@heap.chunks.length} chunks"
  return if not xr

  @heap.scan_chunks_xr
  heapscan_time "#{@heap.xrchunksto.length} src, #{@heap.xrchunksfrom.length} dst"
end

def heapscan_structs
  heapscan_time
  @heap.bucketize
  heapscan_time "#{@heap.buckets.length} buckets"

  @heap.find_arrays
  heapscan_time "#{@heap.allarrays.length} arrays (#{@heap.allarrays.flatten.length} elems)"

  @heap.find_linkedlists
  heapscan_time "#{@heap.alllists.length} lists (#{@heap.alllists.flatten.length} elems)"
end

def heapscan_kernels
  heapscan_time
  @heap.find_kernels
  heapscan_time "#{@heap.kernels.length} kernels"
end

def heapscan_roots
  heapscan_time
  @heap.find_roots
  heapscan_time "#{@heap.roots.length} roots"
end

def heapscan_graph
  heapscan_time
  @heap.dump_graph
  heapscan_time 'graph.gv'
end

def gui_show_list(addr)
  a = resolve(addr)
  #@heap.cp.parse("struct ptr { void *ptr; };") if not @heap.cp.toplevel.struct['ptr']
  h = @heap.linkedlists[a]
  off = h.keys.first
  lst = h[off]

  if not st = lst.map { |l| @heap.chunk_struct[l] }.compact.first
    st = Metasm::C::Struct.new
    st.name = "list_#{'%x' % lst.first}"
    st.members = []
    (@heap.chunks[lst.first] / 4).times { |i|
      n = "u#{i}"
      t = Metasm::C::BaseType.new(:int)
      if i == off/4
        n = "next"
        t = Metasm::C::Pointer.new(st)
      end
      st.members << Metasm::C::Variable.new(n, t)
    }
    @heap.cp.toplevel.struct[st.name] = st
  end
  lst.each { |l| @heap.chunk_struct[l] = st }

  $ghw.addr_struct = {}
  lst.each { |aa|
    $ghw.addr_struct[aa] = @heap.cp.decode_c_struct(st.name, @memory, aa)
  }
  gui.parent_widget.mem.focus_addr(lst.first, :graphheap)
end

def gui_show_array(addr)
  head = resolve(addr)
  e = @heap.xrchunksto[head].to_a.find { |ee| @heap.arrays[ee] and @heap.arrays[ee][head] }
  return if not e
  lst = @heap.arrays[e][head]

  if not st = @heap.chunk_struct[head]
    st = Metasm::C::Struct.new
    st.name = "array_#{'%x' % head}"
    st.members = []
    (@heap.chunks[head] / 4).times { |i|
      n = "u#{i}"
      v = @memory[head+4*i, 4].unpack('L').first
      if @heap.chunks[v]
        t = Metasm::C::Pointer.new(Metasm::C::BaseType.new(:void))
      else
        t = Metasm::C::BaseType.new(:int)
      end
      st.members << Metasm::C::Variable.new(n, t)
    }
    @heap.cp.toplevel.struct[st.name] ||= st
  end
  @heap.chunk_struct[head] = st

  $ghw.addr_struct = { head => @heap.cp.decode_c_struct(st.name, @memory, head) }

  if not st = lst.map { |l| @heap.chunk_struct[l] }.compact.first
    e = lst.first
    st = Metasm::C::Struct.new
    st.name = "elem_#{'%x' % head}"
    st.members = []
    (@heap.chunks[e] / 4).times { |i|
      n = "u#{i}"
      v = @memory[e+4*i, 4].unpack('L').first
      if @heap.chunks[v]
        t = Metasm::C::Pointer.new(Metasm::C::BaseType.new(:void))
      else
        t = Metasm::C::BaseType.new(:int)
      end
      st.members << Metasm::C::Variable.new(n, t)
    }
    @heap.cp.toplevel.struct[st.name] ||= st
  end
  lst.each { |l| @heap.chunk_struct[l] = st }

  lst.each { |aa|
    $ghw.addr_struct[aa] = @heap.cp.decode_c_struct(st.name, @memory, aa)
  }
  gui.parent_widget.mem.focus_addr(head, :graphheap)
end


if gui
  require File.join($heapscan_dir, 'graphheap')
  $ghw = Metasm::Gui::GraphHeapWidget.new(@disassembler, gui.parent_widget.mem)
  gui.parent_widget.mem.addview :graphheap, $ghw
  $ghw.show if $ghw.respond_to?(:show)

  gui.new_command('heap_scan', 'scan the heap(s)') { |*a| heapscan_scan ; $ghw.heap = @heap }
  gui.new_command('heap_scan_noxr', 'scan the heap(s), no xrefs') { |*a| heapscan_scan(false) ; $ghw.heap = @heap }
  gui.new_command('heap_scan_xronly', 'scan the heap(s) for xrefs') { |*a| $ghw.heap.scan_chunks_xr }
  gui.new_command('heap_scanstructs', 'scan the heap for arrays/lists') { |*a| heapscan_structs }
  gui.new_command('heap_list', 'show a linked list') { |a|
    if a.to_s != ''
      gui_show_list(a)
    else
      l = [['addr', 'len']]
      @heap.alllists.each { |al|
        l << [Expression[al.first], al.length]
      }
      gui.listwindow('lists', l) { |*aa| gui_show_list(aa[0][0]) }
    end
  }
  gui.new_command('heap_array', 'show an array') { |a|
    if a.to_s != ''
      gui_show_array(a)
    else
      l = [['addr', 'len']]
      @heap.allarrays.each { |al|
        l << [Expression[al.first], al.length]
      }
      gui.listwindow('arrays', l) { |*aa| gui_show_array(aa[0][0]) }
    end
  }
  gui.new_command('heap_chunk', 'show a chunk') { |a|
    a = resolve(a)
    gui.parent_widget.mem.focus_addr(a, :graphheap)
    $ghw.do_focus_addr(a)
  }
  gui.new_command('heap_strscan', 'scan a string') { |a|
    sa = pattern_scan(a)
    log "found #{sa.length} strings : #{sa.map { |aa| Expression[aa] }.join(' ')}"
    sa.each { |aa|
      next if not ck = @heap.find_chunk(aa)
      log "ptr #{Expression[aa]} in chunk #{Expression[ck]} (#{Expression[@heap.chunks[ck]]}) in list #{@heap.linkedlists && @heap.linkedlists[ck] && true} in array #{@heap.arrays[ck].map { |k, v| "#{Expression[k]} (#{v.length})" }.join(', ') if @heap.arrays and @heap.arrays[ck]}"
    }
  }
  gui.new_command('heap_ptrscan', 'scan a pointer') { |a|
    a = resolve(a)
    if @heap.chunks[a]
      pa = @heap.xrchunksfrom[a].to_a
    else
      pa = pattern_scan(Expression.encode_imm(a, @cpu.size/8, @cpu.endianness))
    end
    log "found #{pa.length} pointers : #{pa.map { |aa| Expression[aa] }.join(' ')}"
    pa.each { |aa|
      next if not ck = @heap.find_chunk(aa)
      log "ptr @#{Expression[aa]} in chunk #{Expression[ck]} (#{Expression[@heap.chunks[ck]]}) in list #{@heap.linkedlists && @heap.linkedlists[ck] && true} in array #{@heap.arrays[ck].map { |k, v| "#{Expression[k]} (#{v.length})" }.join(', ') if @heap.arrays and @heap.arrays[ck]}"
    }
  }

  gui.new_command('heap_snap', 'snapshot the current heap struct') { |a|
    $ghw.snap
  }
  gui.new_command('heap_snap_add', 'snapshot, ignore fields changed between now and last snap') { |a|
    $ghw.snap_add
  }
end