pwntools 0.1.0

data/lib/pwnlib/util/hexdump.rb

@@ -0,0 +1,145 @@
+# encoding: ASCII-8BIT
+require 'rainbow'
+
+module Pwnlib
+  module Util
+    # Method for output a pretty hexdump.
+    # Since this may be used in log module, to avoid cyclic dependency, it is put in a separate module as {Fiddling}
+    # See {ClassMethods} for method details.
+    # @todo Control coloring by context?
+    # @example Call by specifying full module path.
+    #   require 'pwnlib/util/hexdump'
+    #   Pwnlib::Util::HexDump.hexdump('217')
+    #   #=> "00000000  32 31 37                  │217│\n00000003"
+    # @example require 'pwn' and have all methods.
+    #   require 'pwn'
+    #   hexdump('217')
+    #   #=> "00000000  32 31 37                  │217│\n00000003"
+    module HexDump
+      # @note Do not create and call instance method here. Instead, call module method on {HexDump}.
+      module ClassMethods
+        MARKER = "\u2502".freeze
+        HIGHLIGHT_STYLE = ->(s) { Rainbow(s).bg(:red) }
+        DEFAULT_STYLE = {
+          0x00 => ->(s) { Rainbow(s).red },
+          0x0a => ->(s) { Rainbow(s).red },
+          0xff => ->(s) { Rainbow(s).green },
+          marker: ->(s) { Rainbow(s).dimgray },
+          printable: ->(s) { s },
+          unprintable: ->(s) { Rainbow(s).dimgray }
+        }.freeze
+
+        # @!macro [new] hexdump_options
+        #   Color is provided using +rainbow+ gem and only when output is a tty.
+        #   To force enable/disable coloring, call <tt>Rainbow.enabled = true / false</tt>.
+        #   @param [Integer] width
+        #     The max number of characters per line.
+        #   @param [Boolean] skip
+        #     Whether repeated lines should be replaced by a +"*"+.
+        #   @param [Integer] offset
+        #     Offset of the first byte to print in the left column.
+        #   @param [Hash{Integer, Symbol => Proc}] style
+        #     Color scheme to use.
+        #
+        #     Possible keys are:
+        #     * <tt>0x00..0xFF</tt>, for specified byte.
+        #     * +:marker+, for the separator in right column.
+        #     * +:printable+, for printable bytes that don't have style specified.
+        #     * +:unprintable+, for unprintable bytes that don't have style specified.
+        #     The proc is called with a single argument, the string to be formatted.
+        #   @param [String] highlight
+        #     Convenient argument to highlight (red background) some bytes in style.
+
+        # Yields lines of a hexdump-dump of a string. Unless you have massive
+        # amounts of data you probably want to use {#hexdump}.
+        # Returns an Enumerator if no block given.
+        #
+        # @param [#read] io
+        #   The object to be dumped.
+        # @!macro hexdump_options
+        # @return [Enumerator<String>]
+        #   The resulting hexdump, line by line.
+        def hexdump_iter(io, width: 16, skip: true, offset: 0, style: {}, highlight: '')
+          Enumerator.new do |y|
+            style = DEFAULT_STYLE.merge(style)
+            highlight.bytes.each { |b| style[b] = HIGHLIGHT_STYLE }
+            (0..255).each do |b|
+              next if style.include?(b)
+              style[b] = (b.chr =~ /[[:print:]]/ ? style[:printable] : style[:unprintable])
+            end
+
+            styled_bytes = (0..255).map do |b|
+              left_hex = format('%02x', b)
+              c = b.chr
+              right_char = (c =~ /[[:print:]]/ ? c : "\u00b7")
+              [style[b].call(left_hex), style[b].call(right_char)]
+            end
+
+            marker = style[:marker].call(MARKER)
+            spacer = ' '
+
+            byte_index = offset
+            skipping = false
+            last_chunk = ''
+
+            loop do
+              # We assume that chunk is in ASCII-8BIT encoding.
+              chunk = io.read(width)
+              break unless chunk
+              chunk_bytes = chunk.bytes
+              start_byte_index = byte_index
+              byte_index += chunk_bytes.size
+
+              # Yield * once for repeated lines.
+              if skip && last_chunk == chunk
+                y << '*' unless skipping
+                skipping = true
+                next
+              end
+              skipping = false
+              last_chunk = chunk
+
+              hex_bytes = ''
+              printable = ''
+              chunk_bytes.each_with_index do |b, i|
+                left_hex, right_char = styled_bytes[b]
+                hex_bytes << left_hex
+                printable << right_char
+                if i % 4 == 3 && i != chunk_bytes.size - 1
+                  hex_bytes << spacer
+                  printable << marker
+                end
+                hex_bytes << ' '
+              end
+
+              if chunk_bytes.size < width
+                padded_hex_length = 3 * width + (width - 1) / 4
+                hex_length = 3 * chunk_bytes.size + (chunk_bytes.size - 1) / 4
+                hex_bytes << ' ' * (padded_hex_length - hex_length)
+              end
+
+              y << format("%08x  %s #{MARKER}%s#{MARKER}", start_byte_index, hex_bytes, printable)
+            end
+
+            y << format('%08x', byte_index)
+          end
+        end
+
+        # Returns a hexdump-dump of a string.
+        #
+        # @param [String] str
+        #   The string to be hexdumped.
+        # @!macro hexdump_options
+        # @return [String]
+        #   The resulting hexdump.
+        def hexdump(str, width: 16, skip: true, offset: 0, style: {}, highlight: '')
+          hexdump_iter(StringIO.new(str),
+                       width: width, skip: skip, offset: offset, style: style,
+                       highlight: highlight).to_a.join("\n")
+        end
+      end
+
+      extend ClassMethods
+    end
+  end
+end

data/lib/pwnlib/util/packing.rb

@@ -0,0 +1,284 @@
+# encoding: ASCII-8BIT
+
+require 'pwnlib/context'
+
+module Pwnlib
+  module Util
+    # Methods for integer pack/unpack.
+    # See {ClassMethods} for method details.
+    # @example Call by specifying full module path.
+    #   require 'pwnlib/util/packing'
+    #   Pwnlib::Util::Packing.p8(217) #=> "\xD9"
+    # @example require 'pwn' and have all methods.
+    #   require 'pwn'
+    #   p8(217) #=> "\xD9"
+    module Packing
+      # @note Do not create and call instance method here. Instead, call module method on {Packing}.
+      module ClassMethods
+        # Pack arbitrary-sized integer.
+        #
+        # +bits+ indicates number of bits that packed output should use.
+        # The output would be padded to be byte-aligned.
+        #
+        # +bits+ can also be the string 'all',
+        # indicating that the result should be long enough to hold all bits of the number.
+        #
+        # @param [Integer] number
+        #   Number to be packed.
+        # @param [Integer, 'all'] bits
+        #   Number of bits the output should have,
+        #   or +'all'+ for all bits.
+        #   Default to +context.bits+.
+        # @param [String] endian
+        #   Endian to use when packing.
+        #   Can be any value accepted by context (See {Context::ContextType}).
+        #   Default to +context.endian+.
+        # @param [Boolean, String] signed
+        #   Whether the input number should be considered signed when +bits+ is +'all'+.
+        #   Can be any value accepted by context (See {Context::ContextType}).
+        #   Default to +context.signed+.
+        # @return [String]
+        #   The packed string.
+        # @raise [ArgumentError]
+        #   When input integer can't be packed into the size specified by +bits+ and +signed+.
+        #
+        # @example
+        #   pack(0x34, bits: 8) #=> '4'
+        #   pack(0x1234, bits: 16, endian: 'little') #=> "4\x12"
+        #   pack(0xFF, bits: 'all', signed: false) #=> "\xFF"
+        #   pack(0xFF, bits: 'all', endian: 'big', signed: true) #=> "\x00\xFF"
+        def pack(number, bits: nil, endian: nil, signed: nil)
+          if bits == 'all'
+            bits = nil
+            is_all = true
+          else
+            is_all = false
+          end
+
+          context.local(bits: bits, endian: endian, signed: signed) do
+            bits = context.bits
+            endian = context.endian
+            signed = context.signed
+
+            # Verify that bits make sense
+            if signed
+              bits = (number.bit_length | 7) + 1 if is_all
+
+              limit = 1 << (bits - 1)
+              unless -limit <= number && number < limit
+                raise ArgumentError, "signed number=#{number} does not fit within bits=#{bits}"
+              end
+            else
+              if is_all
+                if number < 0
+                  raise ArgumentError, "Can't pack negative number with bits='all' and signed=false"
+                end
+                bits = number.zero? ? 8 : ((number.bit_length - 1) | 7) + 1
+              end
+
+              limit = 1 << bits
+              unless 0 <= number && number < limit
+                raise ArgumentError, "unsigned number=#{number} does not fit within bits=#{bits}"
+              end
+            end
+
+            number &= (1 << bits) - 1
+            bytes = (bits + 7) / 8
+
+            out = []
+            bytes.times do
+              out << (number & 0xFF)
+              number >>= 8
+            end
+            out = out.pack('C*')
+
+            endian == 'little' ? out : out.reverse
+          end
+        end
+
+        # Unpack packed string back to integer.
+        #
+        # +bits+ indicates number of bits that should be used from input data.
+        #
+        # +bits+ can also be the string +'all'+,
+        # indicating that all bytes from input should be used.
+        #
+        # @param [String] data
+        #   String to be unpacked.
+        # @param [Integer, 'all'] bits
+        #   Number of bits to be used from +data+,
+        #   or +'all'+ for all bits.
+        #   Default to +context.bits+
+        # @param [String] endian
+        #   Endian to use when unpacking.
+        #   Can be any value accepted by context (See {Context::ContextType}).
+        #   Default to +context.endian+.
+        # @param [Boolean, String] signed
+        #   Whether the output number should be signed.
+        #   Can be any value accepted by context (See {Context::ContextType}).
+        #   Default to +context.signed+.
+        # @return [Integer]
+        #   The unpacked number.
+        # @raise [ArgumentError]
+        #   When +data.size+ doesn't match +bits+.
+        #
+        # @example
+        #   unpack('4', bits: 8) #=> 52
+        #   unpack("\x3F", bits: 6, signed: false) #=> 63
+        #   unpack("\x3F", bits: 6, signed: true) #=> -1
+        def unpack(data, bits: nil, endian: nil, signed: nil)
+          bits = data.size * 8 if bits == 'all'
+
+          context.local(bits: bits, endian: endian, signed: signed) do
+            bits = context.bits
+            endian = context.endian
+            signed = context.signed
+            bytes = (bits + 7) / 8
+
+            unless data.size == bytes
+              raise ArgumentError, "data.size=#{data.size} does not match with bits=#{bits}"
+            end
+
+            data = data.reverse if endian == 'little'
+            data = data.unpack('C*')
+            number = 0
+            data.each { |c| number = (number << 8) + c }
+            number &= (1 << bits) - 1
+            if signed
+              signbit = number & (1 << (bits - 1))
+              number -= 2 * signbit
+            end
+            number
+          end
+        end
+
+        # Split the data into chunks, and unpack each element.
+        #
+        # +bits+ indicates how many bits each chunk should be.
+        # This should be a multiple of 8,
+        # and size of +data+ should be divisible by +bits / 8+.
+        #
+        # +bits+ can also be the string +'all'+,
+        # indicating that all bytes from input would be used,
+        # and result would be an array with one element.
+        #
+        # @param [String] data
+        #   String to be unpacked.
+        # @param [Integer, 'all'] bits
+        #   Number of bits to be used for each chunk of +data+,
+        #   or +'all'+ for all bits.
+        #   Default to +context.bits+
+        # @param [String] endian
+        #   Endian to use when unpacking.
+        #   Can be any value accepted by context (See {Context::ContextType}).
+        #   Default to +context.endian+.
+        # @param [Boolean, String] signed
+        #   Whether the output number should be signed.
+        #   Can be any value accepted by context (See {Context::ContextType}).
+        #   Default to +context.signed+.
+        # @return [Array<Integer>]
+        #   The unpacked numbers.
+        # @raise [ArgumentError]
+        #   When +bits+ isn't divisible by 8 or +data.size+ isn't divisible by +bits / 8+.
+        #
+        # @todo
+        #   Support +bits+ not divisible by 8, if ever found this useful.
+        #
+        # @example
+        #   unpack_many('haha', bits: 8) #=> [104, 97, 104, 97]
+        #   unpack_many("\xFF\x01\x02\xFE", bits: 16, endian: 'little', signed: true) #=> [511, -510]
+        #   unpack_many("\xFF\x01\x02\xFE", bits: 16, endian: 'big', signed: false) #=> [65281, 766]
+        def unpack_many(data, bits: nil, endian: nil, signed: nil)
+          return [unpack(data, bits: bits, endian: endian, signed: signed)] if bits == 'all'
+
+          context.local(bits: bits, endian: endian, signed: signed) do
+            bits = context.bits
+
+            # TODO(Darkpi): Support this if found useful.
+            raise ArgumentError, 'bits must be a multiple of 8' if bits % 8 != 0
+
+            bytes = bits / 8
+
+            if data.size % bytes != 0
+              raise ArgumentError, "data.size=#{data.size} must be a multiple of bytes=#{bytes}"
+            end
+            ret = []
+            (data.size / bytes).times do |idx|
+              x1 = idx * bytes
+              x2 = x1 + bytes
+              # We already set local context, no need to pass things down.
+              ret << unpack(data[x1...x2], bits: bits)
+            end
+            ret
+          end
+        end
+
+        { 8 => 'c', 16 => 's', 32 => 'l', 64 => 'q' }.each do |sz, ch|
+          define_method("p#{sz}") do |num, **kwargs|
+            context.local(**kwargs) do
+              c = context.signed ? ch : ch.upcase
+              arrow = context.endian == 'little' ? '<' : '>'
+              arrow = '' if sz == 8
+              [num].pack("#{c}#{arrow}")
+            end
+          end
+
+          define_method("u#{sz}") do |data, **kwargs|
+            context.local(**kwargs) do
+              c = context.signed ? ch : ch.upcase
+              arrow = context.endian == 'little' ? '<' : '>'
+              arrow = '' if sz == 8
+              data.unpack("#{c}#{arrow}")[0]
+            end
+          end
+        end
+
+        # TODO(Darkpi): pwntools-python have this for performance reason,
+        #               but current implementation doesn't offer that much performance
+        #               relative to what pwntools-python do. Maybe we should initialize
+        #               those functions (p8lu, ...) like in pwntools-python?
+        [%w(pack p), %w(unpack u)].each do |v1, v2|
+          define_method("make_#{v1}er") do |bits: nil, endian: nil, signed: nil|
+            context.local(bits: bits, endian: endian, signed: signed) do
+              bits = context.bits
+              endian = context.endian
+              signed = context.signed
+
+              if [8, 16, 32, 64].include?(bits)
+                ->(num) { public_send("#{v2}#{bits}", num, endian: endian, signed: signed) }
+              else
+                ->(num) { public_send(v1, num, bits: bits, endian: endian, signed: signed) }
+              end
+            end
+          end
+        end
+
+        def flat(*args, **kwargs, &preprocessor)
+          ret = []
+          p = make_packer(**kwargs)
+          args.each do |it|
+            if preprocessor && !it.is_a?(Array)
+              r = preprocessor[it]
+              it = r unless r.nil?
+            end
+            v = case it
+                when Array then flat(*it, **kwargs, &preprocessor)
+                when Integer then p[it]
+                when String then it.force_encoding('ASCII-8BIT')
+                else
+                  raise ArgumentError, "flat does not support values of type #{it.class}"
+                end
+            ret << v
+          end
+          ret.join
+        end
+
+        # TODO(Darkpi): fit! Which seems super useful.
+
+        include ::Pwnlib::Context
+      end
+
+      extend ClassMethods
+    end
+  end
+end

data/lib/pwnlib/version.rb

@@ -0,0 +1,5 @@
+# encoding: ASCII-8BIT
+
+module Pwnlib
+  VERSION = '0.1.0'.freeze
+end

data/test/constants/constant_test.rb

@@ -0,0 +1,24 @@
+# encoding: ASCII-8BIT
+
+require 'test_helper'
+require 'pwnlib/constants/constant'
+
+class ConstantTest < MiniTest::Test
+  Constant = ::Pwnlib::Constants::Constant
+
+  def test_methods
+    a1 = Constant.new('a', 1)
+    assert_equal(1, a1.to_i)
+    assert_equal(3, a1 | 2)
+    assert_operator(a1, :==, 1)
+    # test coerce
+    assert_operator(1, :==, a1)
+    assert_operator(a1, :==, Constant.new('b', 1))
+    refute_operator(a1, :==, Constant.new('a', 3))
+    assert_equal(3, a1 | Constant.new('a', 2))
+    assert_equal(3, 2 | a1)
+    assert_equal(1, 2 - a1)
+
+    assert_equal(a1.method(:chr).call, "\x01")
+  end
+end