pwntools 0.1.0

data/lib/pwnlib/pwn.rb

@@ -0,0 +1,26 @@
+# encoding: ASCII-8BIT
+
+# require this file would also require all things in pwnlib, but would not
+# pollute anything.
+
+require 'pwnlib/constants/constant'
+require 'pwnlib/constants/constants'
+require 'pwnlib/context'
+require 'pwnlib/dynelf'
+require 'pwnlib/reg_sort'
+
+require 'pwnlib/util/cyclic'
+require 'pwnlib/util/fiddling'
+require 'pwnlib/util/hexdump'
+require 'pwnlib/util/packing'
+
+# include this module in a class to use all pwnlib functions in that class
+# instance.
+module Pwn
+  include ::Pwnlib::Context
+
+  include ::Pwnlib::Util::Cyclic::ClassMethods
+  include ::Pwnlib::Util::Fiddling::ClassMethods
+  include ::Pwnlib::Util::HexDump::ClassMethods
+  include ::Pwnlib::Util::Packing::ClassMethods
+end

data/lib/pwnlib/reg_sort.rb

@@ -0,0 +1,147 @@
+# encoding: ASCII-8BIT
+
+require 'pwnlib/context'
+
+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?
+
+        # 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 }
+
+        # Check input
+        if (in_out.keys - all_regs).any?
+          raise ArgumentError, format('Unknown register! Know: %p.  Got: %p', all_regs, in_out)
+        end
+
+        # 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
+
+        graph = in_out.dup
+        result = []
+
+        # 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 }
+
+        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
+
+        # 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
+
+        # Now assign those collapsed registers.
+        post_mov.sort.each do |dreg, sreg|
+          result << ['mov', dreg, sreg]
+        end
+
+        result
+      end
+
+      private
+
+      # 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']
+      #   check_cycle('a', {'a' => 'b', 'b' => 'c', 'c' => 'b', 'd' => 'a'}) #=> []
+      #   check_cycle('a', {'a' => 'b', 'b' => 'c', 'c' => 'd', 'd' => 'a'})
+      #   #=> ['a', 'b', 'c', 'd']
+      def check_cycle(reg, assignments)
+        check_cycle_(reg, assignments, [])
+      end
+
+      def check_cycle_(reg, assignments, path) # :nodoc:
+        target = assignments[reg]
+        path << reg
+        # No cycle, some other value (e.g. 1)
+        return [] unless assignments.key?(target)
+        # Found a cycle
+        return target == path.first ? path : [] if path.include?(target)
+        check_cycle_(target, assignments, path)
+      end
+    end
+    extend ClassMethods
+  end
+end

data/lib/pwnlib/util/cyclic.rb

@@ -0,0 +1,120 @@
+# encoding: ASCII-8BIT
+
+module Pwnlib
+  module Util
+    # Generate string with easy-to-find pattern.
+    # See {ClassMethods} for method details.
+    # @example Call by specifying full module path.
+    #   require 'pwnlib/util/cyclic'
+    #   Pwnlib::Util::Cyclic.cyclic_find(Pwnlib::Util::Cyclic.cyclic(200)[123, 4]) #=> 123
+    # @example require 'pwn' and have all methods.
+    #   require 'pwn'
+    #   cyclic_find(cyclic(200)[123, 4]) #=> 123
+    module Cyclic
+      # @note Do not create and call instance method here. Instead, call module method on {Cyclic}.
+      module ClassMethods
+        # TODO(Darkpi): Should we put this constant in some 'String' module?
+        ASCII_LOWERCASE = ('a'..'z').to_a.join
+        private_constant :ASCII_LOWERCASE
+
+        # Generator for a sequence of unique substrings of length +n+.
+        # This is implemented using a De Bruijn Sequence over the given +alphabet+.
+        # Returns an Enumerator if no block given.
+        #
+        # @overload de_bruijn(alphabet: ASCII_LOWERCASE, n: 4)
+        #   @param [String, Array] alphabet
+        #     Alphabet to be used.
+        #   @param [Integer] n
+        #     Length of substring that should be unique.
+        #   @return [void]
+        #   @yieldparam c
+        #     Item of the result sequence in order.
+        # @overload de_bruijn(alphabet: ASCII_LOWERCASE, n: 4)
+        #   @param [String, Array] alphabet
+        #     Alphabet to be used.
+        #   @param [Integer] n
+        #     Length of substring that should be unique.
+        #   @return [Enumerator]
+        #     The result sequence.
+        def de_bruijn(alphabet: ASCII_LOWERCASE, n: 4)
+          return to_enum(__method__, alphabet: alphabet, n: n) { alphabet.size**n } unless block_given?
+          k = alphabet.size
+          a = [0] * (k * n)
+
+          db = lambda do |t, p|
+            if t > n
+              (1..p).each { |j| yield alphabet[a[j]] } if (n % p).zero?
+            else
+              a[t] = a[t - p]
+              db.call(t + 1, p)
+              (a[t - p] + 1...k).each do |j|
+                a[t] = j
+                db.call(t + 1, t)
+              end
+            end
+          end
+
+          db[1, 1]
+        end
+
+        # Simple wrapper over {#de_bruijn}, returning at most +length+ items.
+        #
+        # @param [Integer, nil] length
+        #   Desired length of the sequence,
+        #   or +nil+ for the entire sequence.
+        # @param [String, Array] alphabet
+        #   Alphabet to be used.
+        # @param [Integer] n
+        #   Length of substring that should be unique.
+        # @return [String, Array]
+        #   The result sequence of at most +length+ items,
+        #   with same type as +alphabet+.
+        #
+        # @example
+        #   cyclic(alphabet: 'ABC', n: 3) #=> 'AAABAACABBABCACBACCBBBCBCCC'
+        #   cyclic(20) #=> 'aaaabaaacaaadaaaeaaa'
+        def cyclic(length = nil, alphabet: ASCII_LOWERCASE, n: 4)
+          enum = de_bruijn(alphabet: alphabet, n: n)
+          r = length.nil? ? enum.to_a : enum.take(length)
+          alphabet.is_a?(String) ? r.join : r
+        end
+
+        # Find the position of a substring in a De Bruijn sequence
+        #
+        # @todo Speed! See comment in Python pwntools.
+        # @param [String, Array] subseq
+        #   The substring to be found in the sequence.
+        # @param [String, Array] alphabet
+        #   Alphabet to be used.
+        # @param [Integer] n
+        #   Length of substring that should be unique.
+        #   Default to +subseq.size+.
+        # @return [Integer, nil]
+        #   The index +subseq+ first appear in the sequence,
+        #   or +nil+ if not found.
+        #
+        # @example
+        #   cyclic_find(cyclic(300)[217, 4]) #=> 217
+        def cyclic_find(subseq, alphabet: ASCII_LOWERCASE, n: nil)
+          n ||= subseq.size
+          subseq = subseq.chars if subseq.is_a?(String)
+          return nil unless subseq.all? { |c| alphabet.include?(c) }
+
+          pos = 0
+          saved = []
+          de_bruijn(alphabet: alphabet, n: n).each do |c|
+            saved << c
+            if saved.size > subseq.size
+              saved.shift
+              pos += 1
+            end
+            return pos if saved == subseq
+          end
+          nil
+        end
+      end
+
+      extend ClassMethods
+    end
+  end
+end

data/lib/pwnlib/util/fiddling.rb

@@ -0,0 +1,262 @@
+# encoding: ASCII-8BIT
+
+require 'pwnlib/context'
+
+module Pwnlib
+  module Util
+    # Some fiddling methods.
+    # See {ClassMethods} for method details.
+    # @example Call by specifying full module path.
+    #   require 'pwnlib/util/fiddling'
+    #   Pwnlib::Util::Fiddling.enhex('217') #=> '323137'
+    # @example require 'pwn' and have all methods.
+    #   require 'pwn'
+    #   enhex('217') #=> '323137'
+    module Fiddling
+      # @note Do not create and call instance method here. Instead, call module method on {Fiddling}.
+      module ClassMethods
+        # Hex-encodes a string.
+        #
+        # @param [String] s
+        #   String to be encoded.
+        # @return [String]
+        #   Hex-encoded string.
+        #
+        # @example
+        #   enhex('217') #=> '323137'
+        def enhex(s)
+          s.unpack('H*')[0]
+        end
+
+        # Hex-decodes a string.
+        #
+        # @param [String] s
+        #   String to be decoded.
+        # @return [String]
+        #   Hex-decoded string.
+        #
+        # @example
+        #   unhex('353134') #=> '514'
+        def unhex(s)
+          [s].pack('H*')
+        end
+
+        # Present number in hex format, same as python hex() do.
+        #
+        # @param [Integer] n
+        #   The number.
+        #
+        # @return [String]
+        #   The hex format string.
+        #
+        # @example
+        #   hex(0) #=> '0x0'
+        #   hex(-10) #=> '-0xa'
+        #   hex(0xfaceb00cdeadbeef) #=> '0xfaceb00cdeadbeef'
+        def hex(n)
+          (n < 0 ? '-' : '') + format('0x%x', n.abs)
+        end
+
+        # URL-encodes a string.
+        #
+        # @param [String] s
+        #   String to be encoded.
+        # @return [String]
+        #   URL-encoded string.
+        #
+        # @example
+        #   urlencode('shikway') #=> '%73%68%69%6b%77%61%79'
+        def urlencode(s)
+          s.bytes.map { |b| format('%%%02x', b) }.join
+        end
+
+        # URL-decodes a string.
+        #
+        # @param [String] s
+        #   String to be decoded.
+        # @param [Boolean] ignore_invalid
+        #   Whether invalid encoding should be ignore.
+        #   If set to +true+,
+        #   invalid encoding in input are left intact to output.
+        # @return [String]
+        #   URL-decoded string.
+        # @raise [ArgumentError]
+        #   If +ignore_invalid+ is +false+,
+        #   and there are invalid encoding in input.
+        #
+        # @example
+        #   urldecode('test%20url') #=> 'test url'
+        #   urldecode('%qw%er%ty') #=> raise ArgumentError
+        #   urldecode('%qw%er%ty', true) #=> '%qw%er%ty'
+        def urldecode(s, ignore_invalid = false)
+          res = ''
+          n = 0
+          while n < s.size
+            if s[n] != '%'
+              res << s[n]
+              n += 1
+            else
+              cur = s[n + 1, 2]
+              if cur =~ /[0-9a-fA-F]{2}/
+                res << cur.to_i(16).chr
+                n += 3
+              elsif ignore_invalid
+                res << '%'
+                n += 1
+              else
+                raise ArgumentError, 'Invalid input to urldecode'
+              end
+            end
+          end
+          res
+        end
+
+        # Converts the argument to an array of bits.
+        #
+        # @param [String, Integer] s
+        #   Input to be converted into bits.
+        #   If input is integer,
+        #   output would be padded to byte aligned.
+        # @param [String] endian
+        #   Endian for conversion.
+        #   Can be any value accepted by context (See {Context::ContextType}).
+        # @param zero
+        #   Object representing a 0-bit.
+        # @param one
+        #   Object representing a 1-bit.
+        # @return [Array]
+        #   An array consisting of +zero+ and +one+.
+        #
+        # @example
+        #   bits(314) #=> [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 1, 0, 1, 0]
+        #   bits('orz', zero: '-', one: '+').join #=> '-++-++++-+++--+--++++-+-'
+        #   bits(128, endian: 'little') #=> [0, 0, 0, 0, 0, 0, 0, 1]
+        def bits(s, endian: 'big', zero: 0, one: 1)
+          context.local(endian: endian) do
+            is_little = context.endian == 'little'
+            case s
+            when String
+              v = 'B*'
+              v.downcase! if is_little
+              s.unpack(v)[0].chars.map { |ch| ch == '1' ? one : zero }
+            when Integer
+              # TODO(Darkpi): What should we do to negative number?
+              raise ArgumentError, 's must be non-negative' unless s >= 0
+              r = s.to_s(2).chars.map { |ch| ch == '1' ? one : zero }
+              r.unshift(zero) until (r.size % 8).zero?
+              is_little ? r.reverse : r
+            else
+              raise ArgumentError, 's must be either String or Integer'
+            end
+          end
+        end
+
+        # Simple wrapper around {#bits}, which converts output to string.
+        #
+        # @param (see #bits)
+        # @return [String]
+        #
+        # @example
+        #   bits_str('GG') #=> '0100011101000111'
+        def bits_str(s, endian: 'big', zero: 0, one: 1)
+          bits(s, endian: endian, zero: zero, one: one).join
+        end
+
+        # Reverse of {#bits} and {#bits_str}, convert an array of bits back to string.
+        #
+        # @param [String, Array<String, Integer, Boolean>] s
+        #   String or array of bits to be convert back to string.
+        #   <tt>[0, '0', false]</tt> represents 0-bit,
+        #   and <tt>[1, '1', true]</tt> represents 1-bit.
+        # @param [String] endian
+        #   Endian for conversion.
+        #   Can be any value accepted by context (See {Context::ContextType}).
+        # @raise [ArgumentError]
+        #   If input contains value not in <tt>[0, 1, '0', '1', true, false]</tt>.
+        #
+        # @example
+        #   unbits('0100011101000111') #=> 'GG'
+        #   unbits([0, 1, 0, 1, 0, 1, 0, 0]) #=> 'T'
+        #   unbits('0100011101000111', endian: 'little') #=> "\xE2\xE2"
+        def unbits(s, endian: 'big')
+          s = s.chars if s.is_a?(String)
+          context.local(endian: endian) do
+            is_little = context.endian == 'little'
+            bytes = s.map do |c|
+              case c
+              when '1', 1, true then '1'
+              when '0', 0, false then '0'
+              else raise ArgumentError, "cannot decode value #{c.inspect} into a bit"
+              end
+            end
+            [bytes.join].pack(is_little ? 'b*' : 'B*')
+          end
+        end
+
+        # Reverse the bits of each byte in input string.
+        #
+        # @param [String] s
+        #   Input string.
+        # @return [String]
+        #   The string with bits of each byte reversed.
+        #
+        # @example
+        #   bitswap('rb') #=> 'NF'
+        def bitswap(s)
+          unbits(bits(s, endian: 'big'), endian: 'little')
+        end
+
+        # Reverse the bits of a number, and returns the result as number.
+        #
+        # @param [Integer] n
+        # @param [Integer] bits
+        #   The bit length of +n+,
+        #   only the lower +bits+ bits of +n+ would be used.
+        #   Default to context.bits
+        # @return [Integer]
+        #   The number with bits reversed.
+        #
+        # @example
+        #   bitswap_int(217, bits: 8) #=> 155
+        def bitswap_int(n, bits: nil)
+          context.local(bits: bits) do
+            bits = context.bits
+            n &= (1 << bits) - 1
+            bits_str(n, endian: 'little').ljust(bits, '0').to_i(2)
+          end
+        end
+
+        # Base64-encodes a string.
+        # Do NOT contains those stupid newline (with RFC 4648)
+        #
+        # @param [String] s
+        #   String to be encoded.
+        # @return [String]
+        #   Base64-encoded string.
+        #
+        # @example
+        #   b64e('desu') #=> 'ZGVzdQ=='
+        def b64e(s)
+          [s].pack('m0')
+        end
+
+        # Base64-decodes a string.
+        #
+        # @param [String] s
+        #   String to be decoded.
+        # @return [String]
+        #   Base64-decoded string.
+        #
+        # @example
+        #   b64d('ZGVzdQ==') #=> 'desu'
+        def b64d(s)
+          s.unpack('m0')[0]
+        end
+
+        include ::Pwnlib::Context
+      end
+
+      extend ClassMethods
+    end
+  end
+end