heapinfo 0.0.0

data/lib/heapinfo/chunks.rb

@@ -0,0 +1,20 @@
+module HeapInfo
+  # Self-defined array for collecting chunk(s)
+  class Chunks
+    # Instantiate a <tt>HeapInfo::Chunks</tt> object
+    def initialize
+      @chunks = []
+    end
+
+    def method_missing(method_sym, *arguments, &block)
+      return super unless @chunks.respond_to? method_sym
+      @chunks.send(method_sym, *arguments, &block)
+    end
+
+    # Hook <tt>#to_s</tt> for pretty printing.
+    # @return [String]
+    def to_s
+      @chunks.map(&:to_s).join("\n")
+    end
+  end
+end

data/lib/heapinfo/dumper.rb

@@ -0,0 +1,204 @@
+module HeapInfo
+  # Class for memory dump relation works
+  class Dumper
+    # Default dump length
+    DUMP_BYTES = 8
+
+    # Instantiate a <tt>HeapInfo::Dumper</tt> object
+    #
+    # @param [Hash] segments With values <tt>HeapInfo::Segment</tt>
+    # @param [String] mem_filename The filename that can be access for dump. Should be <tt>/proc/[pid]/mem</tt>.
+    def initialize(segments, mem_filename)
+      @segments, @filename = segments, mem_filename
+      need_permission unless dumpable?
+    end
+
+    # A helper for <tt>HeapInfo::Process</tt> to dump memory.
+    # @param [Mixed] args The use input commands, see examples of <tt>HeapInfo::Process#dump</tt>.
+    # @return [String, NilClass] Dump results. If error happend, <tt>nil</tt> is returned.
+    # @example
+    #   p dump(:elf, 4)
+    #   # => "\x7fELF"
+    def dump(*args)
+      return need_permission unless dumpable?
+      base, offset, len = Dumper.parse_cmd(args)
+      if base.instance_of?(Symbol) and @segments[base].instance_of?(Segment)
+        addr = @segments[base].base
+      elsif base.is_a? Integer
+        addr = base
+      else
+        fail # invalid usage
+      end
+      file = mem_f
+      file.pos = addr + offset
+      mem = file.read len
+      file.close
+      mem
+    rescue
+      nil 
+    end
+
+    # Return the dump result as chunks.
+    # see <tt>HeapInfo::Chunks</tt> and <tt>HeapInfo::Chunk</tt> for more information.
+    #
+    # Note: Same as <tt>dump</tt>, need permission of attaching another process.
+    # @return [HeapInfo::Chunks] An array of chunk(s).
+    # @param [Mixed] args Same as arguments of <tt>#dump</tt>
+    def dump_chunks(*args)
+      base = base_of(*args)
+      dump(*args).to_chunks(bits: @segments[:bits], base: base)
+    end
+
+    # Show dump results like in gdb's command <tt>x</tt>.
+    #
+    # Details are in <tt>HeapInfo:Process#x</tt>
+    # @param [Integer] count The number of result need to dump.
+    # @param [Mixed] commands Same format as <tt>#dump(*args)</tt>.
+    # @param [IO] io <tt>IO</tt> that use for printing.
+    # @return [NilClass] The return value of <tt>io.puts</tt>.
+    # @example
+    #   x 3, 0x400000
+    #   # 0x400000:       0x00010102464c457f      0x0000000000000000
+    #   # 0x400010:       0x00000001003e0002
+    def x(count, *commands, io: $stdout)
+      commands = commands + [count * size_t]
+      base = base_of(*commands)
+      res = dump(*commands).unpack(size_t == 4 ? "L*" : "Q*")
+      str = res.group_by.with_index{|_, i| i / (16 / size_t) }.map do |round, values|
+        "%#x:\t" % (base + round * 16) + values.map{|v| Helper.color "0x%0#{size_t * 2}x" % v}.join("\t")
+      end.join("\n")
+      io.puts str
+    end
+
+    # Search a specific value/string/regexp in memory.
+    # <tt>#find</tt> only return the first matched address.
+    # @param [Integer, String, Regexp] pattern The desired search pattern, can be value(<tt>Integer</tt>), string, or regular expression.
+    # @param [Integer, String, Symbol] from Start address for searching, can be segment(<tt>Symbol</tt>) or segments with offset. See examples for more information.
+    # @param [Integer] length The length limit for searching.
+    # @return [Integer, NilClass] The first matched address, <tt>nil</tt> is returned when no such pattern found.
+    # @example
+    #   find(/E.F/, :elf)
+    #   # => 4194305
+    #   find(0x4141414141414141, 'heap+0x10', 0x1000)
+    #   # => 6291472
+    #   find('/bin/sh', :libc)
+    #   # => 140662379588827
+    def find(pattern, from, length)
+      from = base_of(from)
+      length = 1 << 40 if length.is_a? Symbol
+      case pattern
+      when Integer; find_integer(pattern, from, length)
+      when String; find_string(pattern, from, length)
+      when Regexp; find_regexp(pattern, from, length)
+      else; nil
+      end
+    end
+
+
+    # Parse the dump command into <tt>[base, offset, length]</tt>
+    # @param [Array] args The command, see examples for more information
+    # @return [Array<Symbol, Integer>] <tt>[base, offset, length]</tt>, while <tt>base</tt> can be a [Symbol] or an [Integer]. <tt>length</tt> has default value equal to <tt>8</tt>.
+    # @example
+    #   HeapInfo::Dumper.parse_cmd([:heap, 32, 10])
+    #   # [:heap, 32, 10]
+    #   HeapInfo::Dumper.parse_cmd(['heap+0x10, 10'])
+    #   # [:heap, 16, 10]
+    #   HeapInfo::Dumper.parse_cmd(['heap+0x10'])
+    #   # [:heap, 16, 8]
+    #   HeapInfo::Dumper.parse_cmd([0x400000, 4])
+    #   # [0x400000, 0, 4]
+    def self.parse_cmd(args)
+      args = split_cmd args
+      return :fail unless args.size == 3
+      len = args[2].nil? ? DUMP_BYTES : Integer(args[2])
+      offset = Integer(args[1])
+      base = args[0]
+      base = Helper.integer?(base) ? Integer(base) : base.delete(':').to_sym
+      [base, offset, len]
+    end
+
+    # Helper for <tt>#parse_cmd</tt>.
+    #
+    # Split commands to exactly three parts: <tt>[base, offset, length]</tt>
+    # <tt>length</tt> is <tt>nil</tt> if not present.
+    # @param [Array] args
+    # @return [Array<String>] <tt>[base, offset, length]</tt> in string expression.
+    # @example
+    #   HeapInfo::Dumper.split_cmd([:heap, 32, 10])
+    #   # ['heap', '32', '10']
+    #   HeapInfo::Dumper.split_cmd([':heap+0x10, 10'])
+    #   # [':heap', '0x10', '10']
+    #   HeapInfo::Dumper.split_cmd([':heap+0x10'])
+    #   # [':heap', '0x10', nil]
+    #   HeapInfo::Dumper.split_cmd([0x400000, 4])
+    #   # ['4194304', 0, '4']
+    def self.split_cmd(args)
+      args = args.join(',').delete(' ').split(',').reject(&:empty?) # 'heap, 100', 32 => 'heap', '100', '32'
+      return [] if args.empty?
+      if args[0].include? '+' # 'heap+0x1'
+        args.unshift(*args.shift.split('+', 2))
+      elsif args.size <= 2 # no offset given
+        args.insert(1, 0)
+      end
+      args << nil if args.size <= 2 # no length given
+      args[0, 3]
+    end
+
+  private
+    def need_permission
+      puts Helper.color(%q(Could not attach to process. Check the setting of /proc/sys/kernel/yama/ptrace_scope, or try again as the root user.  For more details, see /etc/sysctl.d/10-ptrace.conf), sev: :fatal)
+    end
+
+    # use /proc/[pid]/mem for memory dump, must sure have permission
+    def dumpable?
+      mem_f.close
+      true
+    rescue => e
+      if e.is_a? Errno::EACCES
+        false
+      else
+        throw e
+      end
+    end
+
+    def mem_f
+      File.open(@filename)
+    end
+
+    def base_of(*args)
+      base, offset, _ = Dumper.parse_cmd(args)
+      base = @segments[base].base if @segments[base].is_a? Segment
+      base + offset
+    end
+
+    def find_integer(value, from, length)
+      find_string([value].pack(size_t == 4 ? "L*" : "Q*"), from, length)
+    end
+
+    def find_string(string, from ,length)
+      batch_dumper(from, length) {|str| str.index string}
+    end
+
+    def find_regexp(pattern, from ,length)
+      batch_dumper(from, length) {|str| str =~ pattern}
+    end
+
+    def batch_dumper(from, remain_size)
+      page_size = 0x1000
+      while remain_size > 0
+        dump_size = [remain_size, page_size].min
+        str = dump(from, dump_size)
+        break if str.nil? # unreadable
+        break unless (idx = yield(str)).nil?
+        break if str.length < dump_size # remain is unreadable
+        remain_size -= str.length
+        from += str.length
+      end
+      return if idx.nil?
+      from + idx
+    end
+    def size_t
+      @segments[:bits] / 8
+    end
+  end
+end

data/lib/heapinfo/ext/string.rb

@@ -0,0 +1,29 @@
+module HeapInfo
+  module Ext
+    module String
+      # Methods to be mixed into String
+      module InstanceMethods
+        def to_chunk(bits: 64, base: 0)
+          size_t = bits / 8
+          dumper = lambda{|addr, len| self[addr-base, len]}
+          Chunk.new(size_t, base, dumper)
+        end
+
+        def to_chunks(bits: 64, base: 0)
+          size_t = bits / 8
+          chunks = Chunks.new
+          cur = 0
+          while cur + size_t * 2 <= self.length
+            now_chunk = self[cur, size_t * 2].to_chunk
+            sz = now_chunk.size
+            chunks << self[cur, sz + 1].to_chunk(bits: bits, base: base + cur) # +1 for dump prev_inuse
+            cur += sz
+          end
+          chunks
+        end
+      end
+    end
+  end
+end
+
+::String.send(:include, HeapInfo::Ext::String::InstanceMethods)

data/lib/heapinfo/helper.rb

@@ -0,0 +1,123 @@
+module HeapInfo
+  # Some helper functions
+  module Helper
+    # Get the process id of a process
+    # @param [String] prog The request process name
+    # @return [Fixnum] process id
+    def self.pidof(prog)
+      # plz, don't cmd injection your self :p
+      pid = `pidof #{prog}`.strip.to_i
+      return nil if pid == 0 # process not exists yet
+      throw "pidof #{prog} fail" unless pid.between?(2, 65535)
+      pid
+      #TODO: handle when multi processes exists
+    end
+
+    # Create read <tt>/proc/[pid]/*</tt> methods
+    %w(exe maps).each do |method|
+      self.define_singleton_method("#{method}_of".to_sym) do |pid|
+        begin
+          IO.binread("/proc/#{pid}/#{method}")
+        rescue
+          throw "reading /proc/#{pid}/#{method} error"
+        end
+      end
+    end
+    
+    # Parse the contents of <tt>/proc/[pid]/maps</tt>.
+    #
+    # @param [String] content The file content of <tt>/proc/[pid]/maps</tt>
+    # @return [Array] In form of <tt>[[start, end, permission, name], ...]</tt>. See examples.
+    # @example
+    #   HeapInfo::Helper.parse_maps(<<EOS
+    #     00400000-0040b000 r-xp 00000000 ca:01 271708                             /bin/cat
+    #     00bc4000-00be5000 rw-p 00000000 00:00 0                                  [heap]
+    #     7f2788315000-7f2788316000 r--p 00022000 ca:01 402319                     /lib/x86_64-linux-gnu/ld-2.19.so
+    #     EOS
+    #   )
+    #   # [[0x400000, 0x40b000, 'r-xp', '/bin/cat'], 
+    #   # [0xbc4000, 0xbe5000, 'rw-p', '[heap]'], 
+    #   # [0x7f2788315000, 0x7f2788316000, 'r--p', '/lib/x86_64-linux-gnu/ld-2.19.so']]
+    def self.parse_maps(content)
+      lines = content.split("\n")
+      lines.map do |line|
+        s = line.scan(/^([0-9a-f]+)-([0-9a-f]+)\s([rwxp-]{4})[^\/|\[]*([\/|\[].+)$/)[0]
+        next nil if s.nil?
+        s[0],s[1] = s[0,2].map{|h|h.to_i(16)}
+        s
+      end.compact
+    end
+
+    # Color codes for pretty print
+    COLOR_CODE = {
+      esc_m: "\e[0m",
+      normal_s: "\e[38;5;1m", # red
+      integer: "\e[38;5;12m", # light blue
+      fatal: "\e[38;5;197m", # dark red
+      bin: "\e[38;5;120m", # light green
+      klass: "\e[38;5;155m", # pry like
+      sym: "\e[38;5;229m", # pry like
+    }
+    # Wrapper color codes for for pretty inspect
+    # @param [String] s Contents for wrapper
+    # @param [Symbol?] sev Specific which kind of color want to use, valid symbols are defined in <tt>#COLOR_CODE</tt>.
+    # If this argument is not present, will detect according to the content of <tt>s</tt>
+    # @return [String] wrapper with color codes. 
+    def self.color(s, sev: nil)
+      s = s.to_s
+      color = ''
+      cc = COLOR_CODE
+      if cc.keys.include?(sev)
+        color = cc[sev]
+      elsif s =~ /^(0x)?[0-9a-f]+$/ # integers
+        color = cc[:integer]
+      else #normal string
+        color = cc[:normal_s]
+      end
+      "#{color}#{s.sub(cc[:esc_m], color)}#{cc[:esc_m]}"
+    end
+
+    # Unpack strings to integer.
+    #
+    # Like the <tt>p32</tt> and <tt>p64</tt> in pwntools.
+    #
+    # @param [Integer] size_t Either 4 or 8.
+    # @param [String] data String to be unpack.
+    # @return [Integer] Unpacked result.
+    # @example
+    #   HeapInfo::Helper.unpack(4, "\x12\x34\x56\x78")
+    #   # 0x78563412
+    #   HeapInfo::Helper.unpack(8, "\x12\x34\x56\x78\xfe\xeb\x90\x90")
+    #   # 0x9090ebfe78563412
+    def self.unpack(size_t, data)
+      data.unpack(size_t == 4 ? 'L*' : 'Q*')[0]
+    end
+
+    # Retrieve pure class name(without module) of an object
+    # @param [Object] obj Any instance
+    # @return [String] Class name of <tt>obj</tt>
+    # @example
+    #   # suppose obj is an instance of HeapInfo::Chunk
+    #   Helper.class_name(obj)
+    #   # => 'Chunk'
+    def self.class_name(obj)
+      obj.class.name.split('::').last || obj.class.name
+    end
+
+    # For checking a string is actually an integer
+    # @param [String] str String to be checked
+    # @return [Boolean] If <tt>str</tt> can be converted into integer
+    # @example
+    #   Helper.integer? '1234'
+    #   # => true
+    #   Helper.integer? '0x1234'
+    #   # => true
+    #   Helper.integer? '0xheapoverflow'
+    #   # => false
+    def self.integer?(str)
+      !!Integer(str)
+    rescue ArgumentError, TypeError
+      false
+    end
+  end
+end

data/lib/heapinfo/libc.rb

@@ -0,0 +1,46 @@
+module HeapInfo
+  class Libc < Segment
+    def main_arena_offset
+      return @_main_arena_offset if @_main_arena_offset
+      return nil unless exhaust_search :main_arena
+      @_main_arena_offset
+    end
+
+    def main_arena
+      return @_main_arena.reload if @_main_arena
+      off = main_arena_offset
+      return if off.nil?
+      @_main_arena = Arena.new(off + self.base, process.bits, lambda{|*args|process.send(:dumper).dump(*args)})
+    end
+
+    def self.find(maps, name, process)
+      obj = super(maps, name)
+      obj.send(:process=, process)
+      obj
+    end
+
+
+  private
+    attr_accessor :process
+    # only for searching offset of main_arena now
+    def exhaust_search(symbol)
+      return false if symbol != :main_arena
+      # TODO: read from cache
+      @_main_arena_offset = resolve_main_arena_offset
+      true
+    end
+
+    def resolve_main_arena_offset
+      tmp_elf = HeapInfo::TMP_DIR + "/get_arena"
+      libc_file = HeapInfo::TMP_DIR + "/libc.so.6"
+      ld_file = HeapInfo::TMP_DIR + "/ld.so"
+      flags = "-w #{@process.bits == 32 ? '-m32' : ''}"
+      %x(cp #{self.name} #{libc_file} && \
+         cp #{@process.ld.name} #{ld_file} && \
+         chdir #{File.expand_path('../tools', __FILE__)} && \
+         gcc #{flags} get_arena.c -o #{tmp_elf} 2>&1 > /dev/null && \
+         #{ld_file} --library-path #{HeapInfo::TMP_DIR} #{tmp_elf} && \
+         rm #{tmp_elf} #{libc_file} #{ld_file}).to_i(16)
+    end
+  end
+end

data/lib/heapinfo/nil.rb

@@ -0,0 +1,27 @@
+module HeapInfo
+  # Self defined <tt>nil</tt> like class.
+  #
+  # Be the return values of <tt>#dump</tt> or <tt>#dump_chunks</tt>, to prevent use the return value for calculating accidentally while exploiting remote.
+  class Nil
+    %i(nil? inspect to_s).each do |method_sym|
+      define_method(method_sym){|*args, &block| nil.send(method_sym, *args, &block)}
+    end
+
+    # Hook all missing methods
+    # @return [HeapInfo::Nil] return <tt>self</tt> so that it can be a <tt>nil</tt> chain.
+    # @example
+    #   # h.dump would return Nil when process not found
+    #   p h.dump(:heap)[8,8].unpack("Q*)
+    #   # => nil
+    def method_missing(method_sym, *args, &block)
+      return nil.send(method_sym, *args, &block) if nil.respond_to? method_sym
+      self
+    end
+
+    # To prevent error raised when using <tt>puts Nil.new</tt>
+    # @return [Array] Empty array
+    def to_ary
+      []
+    end
+  end
+end