heapinfo 0.1.0 → 1.0.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: bb558249510b9c160a5f0f699d702a17b1189624
-  data.tar.gz: dacc4780724a2254d8137a23d2f920734db9034d
+  metadata.gz: 22b8757e3a48390494b019943a52754a154051c3
+  data.tar.gz: 4b4af3bf4d67e0ca35b9f5f8feb537499565f56e
 SHA512:
-  metadata.gz: 8fce3e42275b5895a2bec83dbe31d216d1d5a7506905e4d0e527f449c4418f9af00d6f9133d8b1074bb293e7a33ac0557833539562d68dd7a7cfc22d04404885
-  data.tar.gz: 1d036055d53f75e3d9639af6274682b29f5d8f6b275d4644d2619fa29ce4bbed559871ea9caf4b72c66fe81a82dcf4bd0496de239fd88b807df8b50757790ef5
+  metadata.gz: 3872db0f32aaa2fbbedcca004163cb0b5181495cc81605a6587ef5c6dfcebf60f43209146aeda541c716cbdb6bca2bd2770927643cf3bfe50ade2cf7b709a282
+  data.tar.gz: 82a57d0254a650ba531b5f00db676acc3ef4635cd0dbcdc34877f5a09ea8b369667377d96456e466e0a20ad2a6fb3dcb3737cb05966c95cd2483315bb24b7a8c

data/README.md

@@ -1,5 +1,5 @@
 [![Build Status](https://travis-ci.org/david942j/heapinfo.svg?branch=master)](https://travis-ci.org/david942j/heapinfo)
-![](http://ruby-gem-downloads-badge.herokuapp.com/heapinfo?type=total)
+[![Downloads](http://ruby-gem-downloads-badge.herokuapp.com/heapinfo?type=total)](https://rubygems.org/gems/heapinfo)
 [![Code Climate](https://codeclimate.com/github/david942j/heapinfo/badges/gpa.svg)](https://codeclimate.com/github/david942j/heapinfo)
 [![Issue Count](https://codeclimate.com/github/david942j/heapinfo/badges/issue_count.svg)](https://codeclimate.com/github/david942j/heapinfo)
 [![Test Coverage](https://codeclimate.com/github/david942j/heapinfo/badges/coverage.svg)](https://codeclimate.com/github/david942j/heapinfo/coverage)
@@ -7,7 +7,7 @@
 [![MIT License](https://img.shields.io/badge/license-MIT-blue.svg)](http://choosealicense.com/licenses/mit/)
 
 ## HeapInfo
-As pwn lovers, while playing CTF with heap exploitation, we always need a debugger (e.g. gdb) for tracking memory layout. But we don't really need a debugger if we just want to see whether the heap layout same as our imagine or not. Hope this small tool helps us exploit easier ;).
+As pwn lovers, while playing CTF with heap exploitation, we always need a debugger (e.g. gdb) for tracking memory layout. But we don't really need gdb if we want to see whether the heap layout same as our imagine or not. Hope this small tool helps us exploit easier ;).
 
 Implement with ruby because I love ruby :P. But might also implement with Python (if no others did) in the future.
 
@@ -15,7 +15,7 @@ If you prefer [pwntools](https://github.com/Gallopsled/pwntools) for exploiting,
 
 Any suggestion of features or bug issues is welcome.
 
-Relation works are [pwntools-ruby](https://github.com/peter50216/pwntools-ruby) and [gdbpwn](https://github.com/scwuaptx/Pwngdb).
+Related works are [pwntools-ruby](https://github.com/peter50216/pwntools-ruby) and [gdbpwn](https://github.com/scwuaptx/Pwngdb).
 
 ## Install
 **HeapInfo** is still under developing for more features, so the version might change frequently :p
@@ -25,14 +25,18 @@ $ gem install heapinfo
 ```
 
 ## Features
-* Can use in your ruby exploit script or in irb/pry
-* **HeapInfo** works when the `victim` is being traced! i.e. you can use ltrace/strace/gdb and **HeapInfo** simultaneously!
-* `dump` - dump arbitrarily address memory.
-* `layouts` - show the current bin layouts, very useful for heap exploitation.
+* Can use in your ruby exploit script or in irb/pry.
+* **HeapInfo** works when `victim` is being traced! i.e. you can use ltrace/strace/gdb and **HeapInfo** simultaneously!
+* `dump` - Dump arbitrarily address memory.
+* `layouts` - Show the current bin layouts, very useful for heap exploitation.
+* `offset` - Show the offset between given address and segment. Very useful for calculating relative offset.
 * `x` - Provide gdb-like commands.
 * `find` - Provide gdb-like commands.
 * More features and details can be found in [RDoc](http://www.rubydoc.info/github/david942j/heapinfo/master/)
 
+### Under developing
+* `free` - Show what will happend when `free` an address.
+
 ## Usage
 
 #### Load
@@ -68,48 +72,60 @@ NOTICE: While the process is not found, most methods will return `nil`. One way
 h = heapinfo('remote')
 # Process not found
 h.pid # nil
-h.debug {
-  fail unless leak_libc_base == h.libc.base
-  # wrapper with `debug` so that no error will be raised when pwning remote service
-}
+
+# wrapper with `debug` so that no error will be raised when pwning remote service
+h.debug { fail unless leak_libc_base == h.libc.base }
 ```
 
 #### Dump
 Query content of specific address.
 
-NOTICE: you MUST have permission of attaching a program, otherwise dump will fail.
+NOTICE: You MUST have permission of attaching a program, otherwise dump will fail.
 
 i.e. `/proc/sys/kernel/yama/ptrace_scope` set to 0 or run as root.
 
 ```ruby
-h.debug {
+h.debug do
   p h.dump(:libc, 8)
   # => "\x7FELF\x02\x01\x01\x00"
   p h.dump(:heap, 16)
   # => "\x00\x00\x00\x00\x00\x00\x00\x00\x31\x00\x00\x00\x00\x00\x00\x00"
-  p h.dump('heap+0x30, 16') # support offset!
+  p h.dump('heap+0x30', 16) # support offset!
   # => "\x00\x00\x00\x00\x00\x00\x00\x00\x81\x00\x00\x00\x00\x00\x00\x00"
-  p h.dump(:elf, 8)
+  p h.dump('heap+0x30 * 3 + 0x8', 16) # and even complex formula!
+  # => "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+  p h.dump(:program, 8)
   # => "\x7FELF\x02\x01\x01\x00"
-  p h.dump(0x400000, 8) # or simply give addr
+  p h.dump(0x400000, 8) # or simply give address
   # => "\x7FELF\x02\x01\x01\x00"
-}
-# invalid examples:
+end
+
+# invalid example:
 # h.dump('meow') # no such segment
-# h.dump('heap-1, 64') # not support `-`
 ```
 
 #### layouts
 ```ruby
-h.layouts :fastbin
+h.layouts :fast
 ```
 ![fastbin layouts](https://github.com/david942j/heapinfo/blob/master/examples/fastbin_layouts.png?raw=true)
 
 ```ruby
-h.layouts :unsorted_bin, :smallbin
+h.layouts :unsorted, :small
 ```
 ![smallbin layouts](https://github.com/david942j/heapinfo/blob/master/examples/unsorted_smallbin_layouts.png?raw=true)
 
+#### offset
+```ruby
+h.offset(0x7fda86fe8670)
+# 0xf6670 after libc
+h.offset(0x1839cd0, :heap) # specific segment name
+# 0x20cd0 after heap
+h.offset(0x1839cd0)
+# 0x20cd0 after heap
+```
+![offset](https://github.com/david942j/heapinfo/blob/master/examples/offset.png?raw=true)
+
 #### x - gdb-like command
 ```ruby
 h.x 8, :heap
@@ -128,6 +144,12 @@ h.find(/E.F/, 0x400000, 4)
 # => 4194305 # 0x400001
 h.find(/E.F/, 0x400000, 3)
 # => nil
-sh_offset = h.find('/bin/sh', :libc) - h.libc.base
-# => 1559771 # 0x17ccdb
+h.offset(h.find('/bin/sh', :libc))
+# 0x18c177 after libc
 ```
+
+## Tests
+**HeapInfo** currently only run tests on ubuntu, followings are tested glibc versions:
+* libc-2.19
+* libc-2.23
+* libc-2.24

data/lib/heapinfo.rb

@@ -19,8 +19,8 @@ module HeapInfo
   # @param [String, Fixnum] prog
   #   The program name of victim. If a number is given, seem as pid (useful when multi-processes exist).
   # @param [Hash] options Give library's file name.
-  # @option options [String, Regexp] :libc file name of glibc, default is <tt>/libc[^\w]/</tt>
-  # @option options [String, Regexp] :ld file name of dynamic linker/loader, default is <tt>/\/ld-.+\.so/</tt>
+  # @option options [String, Regexp] :libc file name of glibc, default is +/bc[^a-z]*\.so/+.
+  # @option options [String, Regexp] :ld file name of dynamic linker/loader, default is +/\/ld-.+\.so/+.
   # @return [HeapInfo::Process] The object for further usage
   # @example
   #   h = heapinfo './victim'

data/lib/heapinfo/arena.rb

@@ -13,7 +13,7 @@ module HeapInfo
     attr_reader :smallbin
     # @return [Integer] The +system_mem+ in arena.
     attr_reader :system_mem
-    # attr_reader :largebin, :last_remainder
+    # attr_reader :largebin
 
     # Instantiate a {HeapInfo::Arena} object.
     #
@@ -60,12 +60,13 @@ module HeapInfo
     # @param [Symbol] args Bin type(s) you want to see.
     # @return [String] Bin layouts that wrapper with color codes.
     # @example
-    #   puts h.libc.main_arena.layouts :fastbin, :unsorted_bin, :smallbin
+    #   puts h.libc.main_arena.layouts(:fast, :unsorted, :small)
     def layouts(*args)
+      args = args.map(&:to_s).join('|')
       res = ''
-      res += fastbin.map(&:inspect).join if args.include? :fastbin
-      res += unsorted_bin.inspect if args.include? :unsorted_bin
-      res += smallbin.map(&:inspect).join if args.include? :smallbin
+      res += fastbin.map(&:inspect).join if args.include?('fast')
+      res += unsorted_bin.inspect if args.include?('unsort')
+      res += smallbin.map(&:inspect).join if args.include?('small')
       res
     end
 

data/lib/heapinfo/cache.rb

@@ -2,65 +2,73 @@ require 'digest'
 module HeapInfo
   # Self implment file-base cache manager.
   #
-  # Values are recorded in files based on <tt>Marshal</tt>.
+  # Values are recorded in files based on +Marshal+.
   module Cache
     # Directory for caching files.
-    # e.g. HeapInfo will record main_arena_offset for glibc(s)
+    # e.g. HeapInfo will record main_arena_offset for glibc(s).
     CACHE_DIR = File.join(ENV['HOME'], '.cache/heapinfo')
 
-    # Get the key for store libc offsets
-    #
-    # @param [String] libc_path The realpath to libc file
-    # @return [String] The key for cache read/write.
-    def self.key_libc_offset(libc_path)
-      File.join('libc', Digest::MD5.hexdigest(IO.binread(libc_path)), 'offset')
-    end
+    # Define class methods.
+    module ClassMethods
+      # Get the key for store libc offsets.
+      #
+      # @param [String] libc_path The realpath to libc file.
+      # @return [String] The key for cache read/write.
+      def key_libc_offset(libc_path)
+        File.join('libc', Digest::MD5.hexdigest(IO.binread(libc_path)), 'offset')
+      end
 
-    # Write cache to file.
-    #
-    # @param [String] key In file path format, only accept <tt>[\w\/]</tt> to prevent horrible things.
-    # @param [Object] value <tt>value</tt> will be stored with <tt>#Marshal::dump</tt>.
-    # @return [Boolean] true
-    def self.write(key, value)
-      filepath = realpath key
-      FileUtils.mkdir_p(File.dirname(filepath))
-      IO.binwrite(filepath, Marshal.dump(value))
-      true
-    end
+      # Write cache to file.
+      #
+      # @param [String] key In file path format, only accept +[\w/]+ to prevent horrible things.
+      # @param [Object] value +value+ will be stored with +Marshal#dump+.
+      # @return [Boolean] true
+      def write(key, value)
+        filepath = realpath(key)
+        FileUtils.mkdir_p(File.dirname(filepath))
+        IO.binwrite(filepath, Marshal.dump(value))
+        true
+      end
 
-    # Read cache from file.
-    #
-    # @param [String] key In file path format, only accept <tt>[\w\/]</tt> to prevent horrible things.
-    # @return [Object, NilClass] value that recorded, return <tt>nil</tt> when cache miss.
-    def self.read(key)
-      filepath = realpath key
-      return unless File.file?(filepath)
-      Marshal.load(IO.binread(filepath))
-    rescue
-      nil # handle if file content invalid
-    end
+      # Read cache from file.
+      #
+      # @param [String] key In file path format, only accept +[\w\/]+ to prevent horrible things.
+      # @return [Object, NilClass] Value that recorded, return +nil+ when cache miss.
+      def read(key)
+        filepath = realpath(key)
+        return unless File.file?(filepath)
+        Marshal.load(IO.binread(filepath))
+      rescue
+        nil # handle if file content invalid
+      end
 
-    # @param [String] key
-    # @return [String] Prepend with <tt>CACHE_DIR</tt>
-    def self.realpath(key)
-      raise ArgumentError, 'Invalid key(file path)' if key =~ %r{[^\w/]}
-      File.join(CACHE_DIR, key)
-    end
+      # Clear the cache directory.
+      # @return [void]
+      def clear_all
+        FileUtils.rm_rf(CACHE_DIR)
+      end
 
-    # @return [void]
-    def self.init
-      FileUtils.mkdir_p(CACHE_DIR)
-    rescue
-      # To prevent ~/ is not writable.
-      send(:remove_const, :CACHE_DIR)
-      const_set(:CACHE_DIR, File.join(HeapInfo::TMP_DIR, '.cache/heapinfo'))
-    end
+      private
+
+      # @return [void]
+      def init
+        FileUtils.mkdir_p(CACHE_DIR)
+      rescue
+        # To prevent ~/ is not writable.
+        send(:remove_const, :CACHE_DIR)
+        const_set(:CACHE_DIR, File.join(HeapInfo::TMP_DIR, '.cache/heapinfo'))
+        FileUtils.mkdir_p(CACHE_DIR)
+      end
 
-    # Clear the cache directory.
-    # @return [void]
-    def self.clear_all
-      FileUtils.rm_rf CACHE_DIR
+      # @param [String] key
+      # @return [String] Prepend with {HeapInfo::Cache::CACHE_DIR}.
+      def realpath(key)
+        raise ArgumentError, 'Invalid key(file path)' if key =~ %r{[^\w/]}
+        File.join(CACHE_DIR, key)
+      end
     end
+
+    extend ClassMethods
     init
   end
 end

data/lib/heapinfo/dumper.rb

@@ -48,22 +48,21 @@ module HeapInfo
     #
     # Details are in {HeapInfo:Process#x}.
     # @param [Integer] count The number of result need to dump.
-    # @param [Mixed] commands Same format as +#dump(*args)+.
-    # @param [IO] io +IO+ that use for printing.
-    # @return [NilClass] The return value of +io.puts+.
+    # @param [Symbol, String, Integer] address The base address to be dumped.
+    # @return [void]
     # @example
     #   x 3, 0x400000
     #   # 0x400000:       0x00010102464c457f      0x0000000000000000
     #   # 0x400010:       0x00000001003e0002
-    def x(count, *commands, io: $stdout)
-      commands += [count * size_t]
+    def x(count, address)
+      commands = [address, 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|
-        format("%#x:\t", (base + round * 16)) +
+        Helper.hex(base + round * 16) + ":\t" +
           values.map { |v| Helper.color(format("0x%0#{size_t * 2}x", v)) }.join("\t")
       end.join("\n")
-      io.puts str
+      puts str
     end
 
     # Search a specific value/string/regexp in memory.
@@ -92,57 +91,6 @@ module HeapInfo
       end
     end
 
-    # Parse the dump command into +[base, offset, length]+
-    # @param [Array] args The command, see examples for more information
-    # @return [Array<Symbol, Integer>]
-    #   +[base, offset, length]+, while +base+ can be a +Symbol+ or an +Integer+.
-    #   +length+ has default value equals to +8+.
-    # @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 +#parse_cmd+.
-    #
-    # Split commands to exactly three parts: +[base, offset, length]+.
-    # +length+ is +nil+ if not present.
-    # @param [Array] args
-    # @return [Array<String>] +[base, offset, length]+ 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
@@ -169,14 +117,28 @@ module HeapInfo
       File.open(@filename)
     end
 
-    def base_len_of(*args)
-      base, offset, len = Dumper.parse_cmd(args)
-      if HeapInfo::ProcessInfo::EXPORT.include?(base) && (segment = @info.send(base)).is_a?(Segment)
-        base = segment.base
-      elsif !base.is_a?(Integer)
-        raise ArgumentError, "Invalid base: #{base}" # invalid usage
-      end
-      [base + offset, len]
+    # Get the base address and length.
+    #
+    # @param [Integer, Symbol, String] arg The base address, see examples.
+    # @param [Integer] len An integer.
+    # @example
+    #   base_len_of(123, 321) #=> [123, 321]
+    #   base_len_of(123) #=> [123, DUMP_BYTES]
+    #   base_len_of(:heap, 10) #=> [0x603000, 10] # assume heap base @ 0x603000
+    #   base_len_of('heap+0x30', 10) #=> [0x603030, 10]
+    #   base_len_of('elf+0x3*2-1') #=> [0x400005, DUMP_BYTES]
+    def base_len_of(arg, len = DUMP_BYTES)
+      values = HeapInfo::ProcessInfo::EXPORT.map do |seg|
+        segment = @info.respond_to?(seg) && @info.send(seg)
+        [seg, segment.base] if segment.is_a?(Segment)
+      end.compact.to_h
+      base = case arg
+             when Integer then arg
+             when Symbol then values[arg]
+             when String then Helper.evaluate(arg, store: values)
+             end
+      raise ArgumentError, "Invalid base: #{arg.inspect}" unless base.is_a?(Integer) # invalid usage
+      [base, len]
     end
 
     def base_of(*args)

data/lib/heapinfo/glibc/free.rb

@@ -37,8 +37,8 @@ module HeapInfo
         int_free_fast(av, ptr, size)
       elsif !chunk_is_mmapped(ptr) # Though this has been checked in #libc_free
         int_free_small(av, ptr, size)
-      else
-        munmap_chunk(ptr)
+        # else # never reached block, redundant code in glibc.
+        # munmap_chunk(ptr)
       end
     end
 

data/lib/heapinfo/helper.rb

@@ -1,19 +1,9 @@
+require 'dentaku'
+
 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.zero? # process not exists yet
-      throw "pidof #{prog} fail" unless pid.between?(2, 65_535)
-      pid
-      # TODO: handle when multi processes exists
-    end
-
-    # Create read <tt>/proc/[pid]/*</tt> methods
+    # Create read +/proc/[pid]/*+ methods
     %w(exe maps).each do |method|
       define_singleton_method("#{method}_of".to_sym) do |pid|
         begin
@@ -24,106 +14,152 @@ module HeapInfo
       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(%r{^([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
+    # Define class methods here.
+    module ClassMethods
+      # Get the process id of a process
+      # @param [String] prog The request process name
+      # @return [Fixnum] process id
+      def pidof(prog)
+        # plz, don't cmd injection your self :p
+        pid = `pidof #{prog}`.strip.to_i
+        return nil if pid.zero? # process not exists yet
+        throw "pidof #{prog} fail" unless pid.between?(2, 65_535)
+        pid
+        # TODO: handle when multi processes exists
+      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
-    }.freeze
-    # Wrapper color codes 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 +#COLOR_CODE+.
-    # If this argument is not present, will detect according to the content of +s+.
-    # @return [String] wrapper with color codes.
-    def self.color(s, sev: nil)
-      s = s.to_s
-      cc = COLOR_CODE
-      color = if cc.key?(sev) then cc[sev]
-              elsif s =~ /^(0x)?[0-9a-f]+$/ then cc[:integer] # integers
-              else cc[:normal_s] # normal string
-              end
-      "#{color}#{s.sub(cc[:esc_m], color)}#{cc[:esc_m]}"
-    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 parse_maps(content)
+        lines = content.split("\n")
+        lines.map do |line|
+          s = line.scan(%r{^([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
 
-    # 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
+      # enable / disable the color function.
+      # @param [Boolean] on Enable or not.
+      # @return [void]
+      def toggle_color(on: false)
+        @disable_color = !on
+      end
 
-    # Convert number in hex format
-    #
-    # @param [Integer] num Non-negative integer.
-    # @return [String] number in hex format.
-    # @example
-    #   HeapInfo::Helper.hex(1000) # => '0x3e8'
-    def self.hex(num)
-      '0x' + num.to_s(16)
-    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
+      }.freeze
+      # Wrapper color codes 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 +#COLOR_CODE+.
+      # If this argument is not present, will detect according to the content of +s+.
+      # @return [String] wrapper with color codes.
+      def color(s, sev: nil)
+        s = s.to_s
+        return s if @disable_color
+        cc = COLOR_CODE
+        color = if cc.key?(sev) then cc[sev]
+                elsif integer?(s) then cc[:integer] # integers
+                else cc[:normal_s] # normal string
+                end
+        "#{color}#{s.sub(cc[:esc_m], color)}#{cc[:esc_m]}"
+      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
+      # 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 unpack(size_t, data)
+        data.unpack(size_t == 4 ? 'L*' : 'Q*')[0]
+      end
+
+      # Convert number in hex format
+      #
+      # @param [Integer] num Non-negative integer.
+      # @return [String] number in hex format.
+      # @example
+      #   HeapInfo::Helper.hex(1000) # => '0x3e8'
+      def hex(num)
+        return format('0x%x', num) if num >= 0
+        format('-0x%x', -num)
+      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 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 +str+ can be converted into integer.
+      # @example
+      #   Helper.integer? '1234'
+      #   # => true
+      #   Helper.integer? '0x1234'
+      #   # => true
+      #   Helper.integer? '0xheapoverflow'
+      #   # => false
+      def integer?(str)
+        true if Integer(str)
+      rescue ArgumentError, TypeError
+        false
+      end
+
+      # Safe-eval using dentaku.
+      # @param [String] formula Formula to be eval.
+      # @param [Hash{Symbol => Integer}] store Predefined values.
+      # @return [Integer] Evaluate result.
+      def evaluate(formula, store: {})
+        calc = Dentaku::Calculator.new
+        formula = formula.delete(':')
+        formula.gsub!(/0x[\da-z]+/) do |s|
+          s.to_i(16).to_s
+        end
+        calc.store(store).evaluate(formula)
+      end
 
-    # For checking a string is actually an integer.
-    # @param [String] str String to be checked.
-    # @return [Boolean] If +str+ can be converted into integer.
-    # @example
-    #   Helper.integer? '1234'
-    #   # => true
-    #   Helper.integer? '0x1234'
-    #   # => true
-    #   Helper.integer? '0xheapoverflow'
-    #   # => false
-    def self.integer?(str)
-      true if Integer(str)
-    rescue ArgumentError, TypeError
-      false
+      # Get temp filename.
+      # @param [String] name Filename.
+      # @return [String] Temp filename.
+      def tempfile(name)
+        Dir::Tmpname.create(name, HeapInfo::TMP_DIR) {}
+      end
     end
+
+    extend ClassMethods
   end
 end

data/lib/heapinfo/libc.rb

@@ -60,15 +60,17 @@ module HeapInfo
     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'
+      dir = HeapInfo::Helper.tempfile('')
+      FileUtils.mkdir(dir)
+      tmp_elf = File.join(dir, 'get_arena')
+      libc_file = File.join(dir, 'libc.so.6')
+      ld_file = File.join(dir, 'ld.so')
       flags = "-w #{size_t == 4 ? '-m32' : ''}"
       `cp #{name} #{libc_file} && \
          cp #{ld_name} #{ld_file} && \
          gcc #{flags} #{File.expand_path('../tools/get_arena.c', __FILE__)} -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)
+         #{ld_file} --library-path #{dir} #{tmp_elf} && \
+         rm -fr #{dir}`.to_i(16)
     end
   end
 end

data/lib/heapinfo/process.rb

@@ -2,18 +2,18 @@
 module HeapInfo
   # Main class of heapinfo.
   class Process
-    # The default options of libaries,
-    # use for matching glibc and ld segments in <tt>/proc/[pid]/maps</tt>
+    # The default options of libraries,
+    # use for matching glibc and ld segments in +/proc/[pid]/maps+.
     DEFAULT_LIB = {
       libc: /bc[^a-z]*\.so/,
       ld:   %r{/ld-.+\.so}
     }.freeze
-    # @return [Fixnum, NilClass] return the pid of process, <tt>nil</tt> if no such process found
+    # @return [Fixnum, NilClass] return the pid of process, +nil+ if no such process found
     attr_reader :pid
 
-    # Instantiate a <tt>HeapInfo::Process</tt> object
-    # @param [String, Fixnum] prog Process name or pid, see <tt>HeapInfo::heapinfo</tt> for more information
-    # @param [Hash] options libraries' filename, see <tt>HeapInfo::heapinfo</tt> for more information
+    # Instantiate a {HeapInfo::Process} object.
+    # @param [String, Fixnum] prog Process name or pid, see {HeapInfo::heapinfo} for more information
+    # @param [Hash{Symbol => RegExp, String}] options libraries' filename, see {HeapInfo::heapinfo} for more information
     def initialize(prog, options = {})
       @prog = prog
       @options = DEFAULT_LIB.merge options
@@ -23,7 +23,7 @@ module HeapInfo
 
     # Reload a new process with same program name
     #
-    # @return [HeapInfo::Process] return <tt>self</tt> for chainable.
+    # @return [HeapInfo::Process] return +self+ for chainable.
     # @example
     #   puts h.reload!
     def reload!
@@ -34,9 +34,9 @@ module HeapInfo
 
     # Use this method to wrapper all HeapInfo methods.
     #
-    # Since <tt>::HeapInfo</tt> is a tool(debugger) for local usage,
+    # Since {HeapInfo} is a tool(debugger) for local usage,
     # while exploiting remote service, all methods will not work properly.
-    # So I suggest to wrapper all methods inside <tt>#debug</tt>,
+    # So I suggest to wrapper all methods inside {#debug},
     # which will ignore the block while the victim process is not found.
     #
     # @example
@@ -57,49 +57,77 @@ module HeapInfo
     # Note: This method require you have permission of attaching another process.
     # If not, a warning message will present.
     #
-    # @param [Mixed] args Will be parsed into <tt>[base, offset, length]</tt>, see Examples for more information.
+    # @param [Mixed] args Will be parsed into +[base, length]+, see Examples for more information.
     # @return [String, HeapInfo::Nil]
     #   The content needed. When the request address is not readable or the process not exists,
-    #   <tt>HeapInfo::Nil.new</tt> is returned.
+    #   instance of {HeapInfo::Nil} is returned.
     #
     # @example
     #   h = heapinfo('victim')
-    #   h.dump(:heap) # &heap[0, 8]
-    #   h.dump(:heap, 64) # &heap[0, 64]
-    #   h.dump(:heap, 256, 64) # &heap[256, 64]
-    #   h.dump('heap+256, 64'  # &heap[256, 64]
-    #   h.dump('heap+0x100', 64) # &heap[256, 64]
+    #   h.dump(:heap) # heap[0, 8]
+    #   h.dump(:heap, 64) # heap[0, 64]
+    #   h.dump('heap+256', 64)  # heap[256, 64]
+    #   h.dump('heap+0x100', 64) # heap[256, 64]
+    #   h.dump('heap+0x100 * 2 + 0x300', 64) # heap[1024, 64]
     #   h.dump(<segment>, 8) # semgent can be [heap, stack, (program|elf), libc, ld]
     #   h.dump(addr, 64) # addr[0, 64]
     #
     #   # Invalid usage
     #   dump(:meow) # no such segment
-    #   dump('heap-1, 64') # not support '-'
     def dump(*args)
       return Nil.new unless load?
       dumper.dump(*args)
     end
 
     # Return the dump result as chunks.
-    # see <tt>HeapInfo::Dumper#dump_chunks</tt> for more information.
+    # see {HeapInfo::Dumper#dump_chunks} for more information.
     #
     # @return [HeapInfo::Chunks, HeapInfo::Nil] An array of chunk(s).
-    # @param [Mixed] args Same as arguments of <tt>#dump</tt>
+    # @param [Mixed] args Same as arguments of {#dump}.
     def dump_chunks(*args)
       return Nil.new unless load?
       dumper.dump_chunks(*args)
     end
 
+    # Show the offset in pretty way between the segment.
+    # Very useful in pwn when leak some address,
+    # see examples for more details.
+    # @param [Integer] addr The leaked address.
+    # @param [Symbol] sym
+    #   The segement symbol to be calculated offset.
+    #   If this parameter not given, will loop segments
+    #   and find the most close one. See examples for more details.
+    # @return [void] Offset will show to stdout.
+    # @example
+    #   h.offset(0x7f11f6ae1670, :libc)
+    #   #=> 0xf6670 after libc
+    #   h.offset(0x5559edc057a0, :heap)
+    #   #=> 0x9637a0 after heap
+    #   h.offset(0x7f11f6ae1670)
+    #   #=> 0xf6670 after :libc
+    #   h.offset(0x5559edc057a0)
+    #   #=> 0x9637a0 after :heap
+    def offset(addr, sym = nil)
+      return unless load?
+      segment = @info.send(sym) if HeapInfo::ProcessInfo::EXPORT.include?(sym)
+      segment = nil unless segment.is_a?(HeapInfo::Segment)
+      if segment.nil?
+        sym, segment = @info.segments.select { |_, seg| seg.base <= addr }.min_by { |_, seg| addr - seg.base }
+      end
+      return puts "Invalid address #{Helper.hex(addr)}" if segment.nil?
+      puts Helper.color(Helper.hex(addr - segment.base)) + ' after ' + Helper.color(sym, sev: :sym)
+    end
+
     # Gdb-like command
     #
-    # Show dump results like in gdb's command <tt>x</tt>,
-    # while will auto detect the current elf class to decide using <tt>gx</tt> or <tt>wx</tt>.
+    # Show dump results like gdb's command +x+.
+    # While will auto detect the current elf class to decide using +gx+ or +wx+.
     #
-    # The dump results wrapper with color codes and nice typesetting will output to <tt>stdout</tt> by default.
-    # @param [Integer] count The number of result need to dump, see examples for more information
-    # @param [Mixed] commands Same format as <tt>#dump(*args)</tt>, see <tt>#dump</tt> for more information
-    # @param [IO] io <tt>IO</tt> that use for printing, default is <tt>$stdout</tt>
-    # @return [NilClass] The return value of <tt>io.puts</tt>.
+    # The dump results wrapper with color codes and nice typesetting will output to +stdout+.
+    # @param [Integer] count The number of result need to dump, see examples for more information.
+    # @param [String, Symbol, Integer] address The base address to be dumped.
+    #   Same format as {#dump(*args)}, see {#dump} for more information.
+    # @return [void]
     # @example
     #   h.x 8, :heap
     #   # 0x1f0d000:      0x0000000000000000      0x0000000000002011
@@ -110,23 +138,23 @@ module HeapInfo
     #   h.x 3, 0x400000
     #   # 0x400000:       0x00010102464c457f      0x0000000000000000
     #   # 0x400010:       0x00000001003e0002
-    def x(count, *commands, io: $stdout)
-      return unless load? && io.respond_to?(:puts)
-      dumper.x(count, *commands, io: io)
+    def x(count, address)
+      return unless load?
+      dumper.x(count, address)
     end
 
     # Gdb-like command.
     #
     # Search a specific value/string/regexp in memory.
     # @param [Integer, String, Regexp] pattern
-    #   The desired search pattern, can be value(<tt>Integer</tt>), string, or regular expression.
+    #   The desired search pattern, can be value(+Integer+), string, or regular expression.
     # @param [Integer, String, Symbol] from
-    #   Start address for searching, can be segment(<tt>Symbol</tt>) or segments with offset.
+    #   Start address for searching, can be segment(+Symbol+) or segments with offset.
     #   See examples for more information.
     # @param [Integer] length
     #   The search length limit, default is unlimited,
     #   which will search until pattern found or reach unreadable memory.
-    # @return [Integer, NilClass] The first matched address, <tt>nil</tt> is returned when no such pattern found.
+    # @return [Integer, NilClass] The first matched address, +nil+ is returned when no such pattern found.
     # @example
     #   h.find(0xdeadbeef, 'heap+0x10', 0x1000)
     #   # => 6299664 # 0x602010
@@ -147,15 +175,14 @@ module HeapInfo
 
     # Pretty dump of bins layouts.
     #
-    # The request layouts will output to <tt>stdout</tt> by default.
+    # The request layouts will output to +stdout+.
     # @param [Array<Symbol>] args Bin type(s) you want to see.
-    # @param [IO] io <tt>IO</tt> that use for printing, default is <tt>$stdout</tt>
-    # @return [NilClass] The return value of <tt>io.puts</tt>.
+    # @return [void]
     # @example
-    #   h.layouts :fastbin, :unsorted_bin, :smallbin
-    def layouts(*args, io: $stdout)
-      return unless load? && io.respond_to?(:puts)
-      io.puts libc.main_arena.layouts(*args)
+    #   h.layouts(:fast, :unsorted, :small)
+    def layouts(*args)
+      return unless load?
+      puts libc.main_arena.layouts(*args)
     end
 
     # Show simple information of target process.

data/lib/heapinfo/process_info.rb

@@ -6,7 +6,7 @@ module HeapInfo
   class ProcessInfo
     # Methods to be transparent to +process+.
     # e.g. +process.libc+ alias to +process.info.libc+.
-    EXPORT = %i(libc ld heap elf program stack bits).freeze
+    EXPORT = %i(libc ld heap program elf stack bits).freeze
 
     # @return [Integer] 32 or 64.
     attr_reader :bits
@@ -20,14 +20,14 @@ module HeapInfo
     attr_reader :ld
     alias elf program
 
-    # Instantiate a {ProcessInfo} object
+    # Instantiate a {ProcessInfo} object.
     #
     # @param [HeapInfo::Process] process Load information from maps/memory for +process+.
     def initialize(process)
       @pid = process.pid
       options = process.instance_variable_get(:@options)
       maps!
-      @bits = bits_of Helper.exe_of @pid
+      @bits = bits_of(Helper.exe_of(@pid))
       @elf = @program = Segment.find(maps, File.readlink("/proc/#{@pid}/exe"))
       @stack = Segment.find(maps, '[stack]')
       # well.. stack is a strange case because it will grow in runtime..
@@ -39,18 +39,27 @@ module HeapInfo
     # Heap will not be mmapped if the process not use heap yet, so create a lazy loading method.
     # Will re-read maps when heap segment not found yet.
     #
-    # @return [HeapInfo::Segment] The {Segment} of heap
+    # @return [HeapInfo::Segment] The {Segment} of heap.
     def heap # special handle because heap might not be initialized in the beginning
       @heap ||= Segment.find(maps!, '[heap]')
     end
 
+    # Return segemnts load currently.
+    # @return [Hash{Symbol => Segment}] The segments in hash format.
+    def segments
+      EXPORT.map do |sym|
+        seg = send(sym)
+        [sym, seg] if seg.is_a?(Segment)
+      end.compact.to_h
+    end
+
     private
 
     attr_reader :maps
 
     # force reload maps
     def maps!
-      @maps = Helper.parse_maps Helper.maps_of @pid
+      @maps = Helper.parse_maps(Helper.maps_of(@pid))
     end
 
     def bits_of(elf)

data/lib/heapinfo/version.rb

@@ -1,3 +1,3 @@
 module HeapInfo
-  VERSION = '0.1.0'.freeze
+  VERSION = '1.0.0'.freeze
 end

metadata

@@ -1,15 +1,29 @@
 --- !ruby/object:Gem::Specification
 name: heapinfo
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 1.0.0
 platform: ruby
 authors:
 - david942j
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2017-02-14 00:00:00.000000000 Z
+date: 2017-03-06 00:00:00.000000000 Z
 dependencies:
+- !ruby/object:Gem::Dependency
+  name: dentaku
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2'
 - !ruby/object:Gem::Dependency
   name: rspec
   requirement: !ruby/object:Gem::Requirement