pwntools 0.1.0 → 1.0.0

data/lib/pwnlib/ext/helper.rb

@@ -2,7 +2,7 @@
 
 module Pwnlib
   module Ext
-    # Helper methods for defining extension
+    # Helper methods for defining extension.
     module Helper
       def def_proxy_method(mod, *ms, **m2)
         ms.flatten
@@ -10,9 +10,9 @@ module Pwnlib
           .concat(m2.to_a)
           .each do |method, proxy_to|
             class_eval <<-EOS
-            def #{method}(*args, &block)
-            #{mod}.#{proxy_to}(self, *args, &block)
-            end
+              def #{method}(*args, &block)
+              #{mod}.#{proxy_to}(self, *args, &block)
+              end
             EOS
           end
       end

data/lib/pwnlib/logger.rb

@@ -0,0 +1,87 @@
+# encoding: ASCII-8BIT
+
+require 'logger'
+require 'rainbow'
+
+require 'pwnlib/context'
+
+module Pwnlib
+  # Logging module for printing status during an exploitation, and internally within {Pwnlib}.
+  #
+  # An exploit developer can use +log+ to print out status messages during an exploitation.
+  module Logger
+    # The type for logger which inherits Ruby builtin Logger.
+    # Main difference is using +context.log_level+ instead of +level+ in logging methods.
+    class LoggerType < ::Logger
+      SEV_COLOR = {
+        'DEBUG' => '#ff5f5f',
+        'INFO' => '#87ff00',
+        'WARN' => '#ffff00',
+        'ERROR' => '#ff5f00',
+        'FATAL' => '#ff0000'
+      }.freeze
+
+      # Instantiate a {Pwnlib::Logger::LoggerType} object.
+      def initialize
+        super(STDOUT)
+        @formatter = proc do |severity, _datetime, _progname, msg|
+          format("[%s] %s\n", Rainbow(severity).color(SEV_COLOR[severity]), msg)
+        end
+      end
+
+      # Log the message with indent.
+      #
+      # @param [String] message
+      #   The message to log.
+      # @param [DEBUG, INFO, WARN, ERROR, FATAL, UNKNOWN] level
+      #   The severity of the message.
+      def indented(message, level: DEBUG)
+        return if @logdev.nil? || level < context.log_level
+        @logdev.write(
+          message.lines.map { |s| '    ' + s }.join + "\n"
+        )
+        true
+      end
+
+      private
+
+      def add(severity, message = nil, progname = nil)
+        severity ||= UNKNOWN
+        return true if severity < context.log_level
+        super(severity, message, progname)
+      end
+
+      include ::Pwnlib::Context
+    end
+
+    @log = LoggerType.new
+
+    # @!attribute [r] logger
+    #   @return [LoggerType] the singleton logger for all classes.
+    class << self
+      attr_reader :log
+    end
+
+    # Include this module to use logger.
+    # Including {Pwnlib::Logger} from module M would add +log+ as a private instance method and a private class
+    # method for module M.
+    # @!visibility private
+    module IncludeLogger
+      private
+
+      def log
+        ::Pwnlib::Logger.log
+      end
+    end
+
+    # @!visibility private
+    def self.included(base)
+      base.include(IncludeLogger)
+      class << base
+        include IncludeLogger
+      end
+    end
+
+    include ::Logger::Severity
+  end
+end

data/lib/pwnlib/memleak.rb

@@ -3,59 +3,88 @@
 require 'pwnlib/util/packing'
 
 module Pwnlib
-  # MemLeak is a caching and heuristic tool for exploiting memory leaks.
+  # A class caching and heuristic tool for exploiting memory leaks.
   class MemLeak
-    PAGE_SIZE = 0x1000
-    PAGE_MASK = ~(PAGE_SIZE - 1)
-
+    # Instantiate a {Pwnlib::MemLeak} object.
+    #
+    # @yieldparam [Integer] leak_addr
+    #   The start address that the leaker should leak from.
+    #
+    # @yieldreturn [String]
+    #   A leaked non-empty byte string, starting from +leak_addr+.
     def initialize(&block)
       @leak = block
-      @base = nil
       @cache = {}
     end
 
-    def find_elf_base(ptr)
-      ptr &= PAGE_MASK
-      loop do
-        return @base = ptr if n(ptr, 4) == "\x7fELF"
-        ptr -= PAGE_SIZE
-      end
-    end
-
-    # Call the leaker function on address `addr`.
-    # Store the result to @cache
-    def do_leak(addr)
-      unless @cache.key?(addr)
-        data = @leak.call(addr)
-        data.bytes.each.with_index(addr) { |b, i| @cache[i] = b }
-      end
-      @cache[addr]
-    end
-
-    # Leak `numb` bytes at `addr`.
+    # Leak +numb+ bytes at +addr+.
     # Returns a string with the leaked bytes.
+    #
+    # @param [Integer] addr
+    #   The starting address of the leak.
+    # @param [Integer] numb
+    #   Number of bytes to be leaked.
+    #
+    # @return [String]
+    #   The leaked byte string.
     def n(addr, numb)
       (0...numb).map { |i| do_leak(addr + i) }.pack('C*')
     end
 
-    # Leak byte at ``((uint8_t*) addr)[ndx]``
+    # Leak a byte at +*((uint8_t*) addr)+.
+    #
+    # @param [Integer] addr
+    #   The address of the leak.
+    #
+    # @return [Integer]
+    #   The leaked byte.
     def b(addr)
-      n(addr, 1)
+      Util::Packing.u8(n(addr, 1))
     end
 
-    # Leak word at ``((uint16_t*) addr)[ndx]``
+    # Leak a word at +*((uint16_t*) addr)+.
+    #
+    # @param [Integer] addr
+    #   The address of the leak.
+    #
+    # @return [Integer]
+    #   The leaked word.
     def w(addr)
       Util::Packing.u16(n(addr, 2))
     end
 
-    # Leak dword at ``((uint32_t*) addr)[ndx]``
+    # Leak a dword at +*((uint32_t*) addr)+.
+    #
+    # @param [Integer] addr
+    #   The address of the leak.
+    #
+    # @return [Integer]
+    #   The leaked dword.
     def d(addr)
       Util::Packing.u32(n(addr, 4))
     end
 
-    # Leak qword at ``((uint64_t*) addr)[ndx]``
+    # Leak a qword at +*((uint64_t*) addr)+.
+    #
+    # @param [Integer] addr
+    #   The address of the leak.
+    #
+    # @return [Integer]
+    #   The leaked qword.
     def q(addr)
       Util::Packing.u64(n(addr, 8))
     end
+
+    private
+
+    # Call the leaker function on address +addr+.
+    # The result would be cached.
+    def do_leak(addr)
+      unless @cache.key?(addr)
+        data = @leak.call(addr)
+        data.bytes.each.with_index(addr) { |b, i| @cache[i] = b }
+      end
+      @cache[addr]
+    end
   end
 end

data/lib/pwnlib/pwn.rb

@@ -1,26 +1,37 @@
 # encoding: ASCII-8BIT
 
-# require this file would also require all things in pwnlib, but would not
-# pollute anything.
+# require this file would also require all things in pwnlib, but would not pollute anything.
 
+require 'pwnlib/asm'
 require 'pwnlib/constants/constant'
 require 'pwnlib/constants/constants'
 require 'pwnlib/context'
 require 'pwnlib/dynelf'
+require 'pwnlib/elf/elf'
+require 'pwnlib/logger'
 require 'pwnlib/reg_sort'
+require 'pwnlib/shellcraft/shellcraft'
+require 'pwnlib/tubes/sock'
 
 require 'pwnlib/util/cyclic'
 require 'pwnlib/util/fiddling'
+require 'pwnlib/util/getdents'
 require 'pwnlib/util/hexdump'
+require 'pwnlib/util/lists'
 require 'pwnlib/util/packing'
 
-# include this module in a class to use all pwnlib functions in that class
-# instance.
+# include this module in a class to use all pwnlib functions in that class instance.
 module Pwn
+  include ::Pwnlib::Asm
   include ::Pwnlib::Context
+  include ::Pwnlib::Logger
+  include ::Pwnlib::Util::Cyclic
+  include ::Pwnlib::Util::Fiddling
+  include ::Pwnlib::Util::HexDump
+  include ::Pwnlib::Util::Lists
+  include ::Pwnlib::Util::Packing
 
-  include ::Pwnlib::Util::Cyclic::ClassMethods
-  include ::Pwnlib::Util::Fiddling::ClassMethods
-  include ::Pwnlib::Util::HexDump::ClassMethods
-  include ::Pwnlib::Util::Packing::ClassMethods
+  def shellcraft
+    ::Pwnlib::Shellcraft::Shellcraft.instance
+  end
 end

data/lib/pwnlib/reg_sort.rb

@@ -1,127 +1,122 @@
 # encoding: ASCII-8BIT
 
 require 'pwnlib/context'
+require 'pwnlib/util/ruby'
 
 module Pwnlib
   # Do topological sort on register assignments.
   module RegSort
-    # @note Do not create and call instance method here. Instead, call module method on {RegSort}.
-    module ClassMethods
-      # Sorts register dependencies.
-      #
-      # Given a dictionary of registers to desired register contents,
-      # return the optimal order in which to set the registers to
-      # those contents.
-      #
-      # The implementation assumes that it is possible to move from
-      # any register to any other register.
-      #
-      # @param [Hash<Symbol, String => Object>] in_out
-      #   Dictionary of desired register states.
-      #   Keys are registers, values are either registers or any other value.
-      # @param [Array<String>] all_regs
-      #   List of all possible registers.
-      #   Used to determine which values in +in_out+ are registers, versus
-      #   regular values.
-      # @option [Boolean] randomize
-      #   Randomize as much as possible about the order or registers.
-      #
-      # @return [Array]
-      #   Array of instructions, see examples for more details.
-      #
-      # @example
-      #   regs = %w(a b c d x y z)
-      #   regsort({a: 1, b: 2}, regs)
-      #   => [['mov', 'a', 1], ['mov', 'b', 2]]
-      #   regsort({a: 'b', b: 'a'}, regs)
-      #   => [['xchg', 'a', 'b']]
-      #   regsort({a: 1, b: 'a'}, regs)
-      #   => [['mov', 'b', 'a'], ['mov', 'a', 1]]
-      #   regsort({a: 'b', b: 'a', c: 3}, regs)
-      #   => [['mov', 'c', 3], ['xchg', 'a', 'b']]
-      #   regsort({a: 'b', b: 'a', c: 'b'}, regs)
-      #   => [['mov', 'c', 'b'], ['xchg', 'a', 'b']]
-      #   regsort({a: 'b', b: 'c', c: 'a', x: '1', y: 'z', z: 'c'}, regs)
-      #   => [['mov', 'x', '1'],
-      #       ['mov', 'y', 'z'],
-      #       ['mov', 'z', 'c'],
-      #       ['xchg', 'a', 'b'],
-      #       ['xchg', 'b', 'c']]
-      #
-      # @note
-      #   Different from python-pwntools, we don't support +tmp+/+xchg+ options
-      #   because there's no such usage at all.
-      def regsort(in_out, all_regs, randomize: nil)
-        # randomize = context.randomize if randomize.nil?
+    module_function
 
-        # TODO(david942j): stringify_keys
-        in_out = in_out.map { |k, v| [k.to_s, v] }.to_h
-        # Drop all registers which will be set to themselves.
-        # Ex. {eax: 'eax'}
-        in_out.reject! { |k, v| k == v }
+    # Sorts register dependencies.
+    #
+    # Given a dictionary of registers to desired register contents, return the optimal order in which to set the
+    # registers to those contents.
+    #
+    # The implementation assumes that it is possible to move from any register to any other register.
+    #
+    # @param [Hash<Symbol, String => Object>] in_out
+    #   Dictionary of desired register states.
+    #   Keys are registers, values are either registers or any other value.
+    # @param [Array<String>] all_regs
+    #   List of all possible registers.
+    #   Used to determine which values in +in_out+ are registers, versus regular values.
+    # @param [Boolean] randomize
+    #   Randomize as much as possible about the order or registers.
+    #
+    # @return [Array]
+    #   Array of instructions, see examples for more details.
+    #
+    # @diff We don't support +tmp+/+xchg+ options because there's no such usage at all.
+    #
+    # @example
+    #   regs = %w(a b c d x y z)
+    #   regsort({a: 1, b: 2}, regs)
+    #   #=> [['mov', 'a', 1], ['mov', 'b', 2]]
+    #   regsort({a: 'b', b: 'a'}, regs)
+    #   #=> [['xchg', 'a', 'b']]
+    #   regsort({a: 1, b: 'a'}, regs)
+    #   #=> [['mov', 'b', 'a'], ['mov', 'a', 1]]
+    #   regsort({a: 'b', b: 'a', c: 3}, regs)
+    #   #=> [['mov', 'c', 3], ['xchg', 'a', 'b']]
+    #   regsort({a: 'b', b: 'a', c: 'b'}, regs)
+    #   #=> [['mov', 'c', 'b'], ['xchg', 'a', 'b']]
+    #   regsort({a: 'b', b: 'c', c: 'a', x: '1', y: 'z', z: 'c'}, regs)
+    #   #=> [['mov', 'x', '1'],
+    #        ['mov', 'y', 'z'],
+    #        ['mov', 'z', 'c'],
+    #        ['xchg', 'a', 'b'],
+    #        ['xchg', 'b', 'c']]
+    def regsort(in_out, all_regs, randomize: nil)
+      # randomize = context.randomize if randomize.nil?
 
-        # Check input
-        if (in_out.keys - all_regs).any?
-          raise ArgumentError, format('Unknown register! Know: %p.  Got: %p', all_regs, in_out)
-        end
+      # TODO(david942j): stringify_keys
+      in_out = in_out.map { |k, v| [k.to_s, v] }.to_h
+      # Drop all registers which will be set to themselves.
+      # Ex. {eax: 'eax'}
+      in_out.reject! { |k, v| k == v }
 
-        # Collapse constant values
-        #
-        # Ex. {eax: 1, ebx: 1} can be collapsed to {eax: 1, ebx: 'eax'}.
-        # +post_mov+ are collapsed registers, set their values in the end.
-        post_mov = in_out.group_by { |_, v| v }.each_value.with_object({}) do |list, hash|
-          list.sort!
-          first_reg, val = list.shift
-          # Special case for val.zero? because zeroify registers cost cheaper than mov.
-          next if list.empty? || all_regs.include?(val) || val.zero?
-          list.each do |reg, _|
-            hash[reg] = first_reg
-            in_out.delete(reg)
-          end
-        end
+      # Check input
+      if (in_out.keys - all_regs).any?
+        raise ArgumentError, format('Unknown register! Know: %p.  Got: %p', all_regs, in_out)
+      end
 
-        graph = in_out.dup
-        result = []
+      # Collapse constant values.
+      #
+      # Ex. {eax: 1, ebx: 1} can be collapsed to {eax: 1, ebx: 'eax'}.
+      # +post_mov+ are collapsed registers, set their values in the end.
+      post_mov = in_out.group_by { |_, v| v }.each_value.with_object({}) do |list, hash|
+        list.sort!
+        first_reg, val = list.shift
+        # Special case for val.zero? because zeroify registers is cheaper than mov.
+        next if list.empty? || all_regs.include?(val) || val.zero?
+        list.each do |reg, _|
+          hash[reg] = first_reg
+          in_out.delete(reg)
+        end
+      end
 
-        # Let's do the topological sort.
-        # so sad ruby 2.1 doesn't have +itself+...
-        deg = graph.values.group_by { |i| i }.map { |k, v| [k, v.size] }.to_h
-        graph.each_key { |k| deg[k] ||= 0 }
+      graph = in_out.dup
+      result = []
 
-        until deg.empty?
-          min_deg = deg.min_by { |_, v| v }[1]
-          break unless min_deg.zero? # remain are all cycles
-          min_pivs = deg.select { |_, v| v == min_deg }
-          piv = randomize ? min_pivs.sample : min_pivs.first
-          dst = piv.first
-          deg.delete(dst)
-          next unless graph.key?(dst) # Reach an end node.
-          deg[graph[dst]] -= 1
-          result << ['mov', dst, graph[dst]]
-          graph.delete(dst)
-        end
+      # Let's do the topological sort.
+      # so sad ruby 2.1 doesn't have +itself+...
+      deg = graph.values.group_by { |i| i }.map { |k, v| [k, v.size] }.to_h
+      graph.each_key { |k| deg[k] ||= 0 }
 
-        # Remain must be cycles.
-        graph.each_key do |reg|
-          cycle = check_cycle(reg, graph)
-          cycle.each_cons(2) do |d, s|
-            result << ['xchg', d, s]
-          end
-          cycle.each { |r| graph.delete(r) }
-        end
+      until deg.empty?
+        min_deg = deg.min_by { |_, v| v }[1]
+        break unless min_deg.zero? # remain are all cycles
+        min_pivs = deg.select { |_, v| v == min_deg }
+        piv = randomize ? min_pivs.sample : min_pivs.first
+        dst = piv.first
+        deg.delete(dst)
+        next unless graph.key?(dst) # Reach an end node.
+        deg[graph[dst]] -= 1
+        result << ['mov', dst, graph[dst]]
+        graph.delete(dst)
+      end
 
-        # Now assign those collapsed registers.
-        post_mov.sort.each do |dreg, sreg|
-          result << ['mov', dreg, sreg]
+      # Remain must be cycles.
+      graph.each_key do |reg|
+        cycle = check_cycle(reg, graph)
+        cycle.each_cons(2) do |d, s|
+          result << ['xchg', d, s]
         end
+        cycle.each { |r| graph.delete(r) }
+      end
 
-        result
+      # Now assign those collapsed registers.
+      post_mov.sort.each do |dreg, sreg|
+        result << ['mov', dreg, sreg]
       end
 
-      private
+      result
+    end
 
-      # Walk down the assignment list of a register,
-      # return the path walked if it is encountered again.
+    Pwnlib::Util::Ruby.private_class_method_block do
+      # Walk down the assignment list of a register, return the path walked if it is encountered again.
+      #
       # @example
       #   check_cycle('a', {'a' => 1}) #=> []
       #   check_cycle('a', {'a' => 'a'}) #=> ['a']
@@ -132,16 +127,15 @@ module Pwnlib
         check_cycle_(reg, assignments, [])
       end
 
-      def check_cycle_(reg, assignments, path) # :nodoc:
+      def check_cycle_(reg, assignments, path)
         target = assignments[reg]
         path << reg
-        # No cycle, some other value (e.g. 1)
+        # No cycle, some other value (e.g. 1).
         return [] unless assignments.key?(target)
-        # Found a cycle
+        # Found a cycle.
         return target == path.first ? path : [] if path.include?(target)
         check_cycle_(target, assignments, path)
       end
     end
-    extend ClassMethods
   end
 end