pwntools 0.1.0

data/test/constants/constants_test.rb

@@ -0,0 +1,31 @@
+# encoding: ASCII-8BIT
+
+require 'test_helper'
+require 'pwnlib/constants/constants'
+require 'pwnlib/context'
+
+class ConstantsTest < MiniTest::Test
+  include ::Pwnlib::Context
+  Constants = ::Pwnlib::Constants
+
+  def test_amd64
+    context.local(arch: 'amd64') do
+      assert_equal('Constant("SYS_read", 0x0)', Constants.SYS_read.inspect)
+      assert_equal('__NR_arch_prctl', Constants.__NR_arch_prctl.to_s)
+      assert_equal('Constant("(O_CREAT)", 0x40)', Constants.eval('O_CREAT').inspect)
+      # TODO(david942j): implement 'real' Constants.eval
+      # assert_equal('Constant("(O_CREAT | O_WRONLY)", 0x41)', Constants.eval('O_CREAT | O_WRONLY').inspect)
+    end
+  end
+
+  def test_i386
+    context.local(arch: 'i386') do
+      assert_equal('Constant("SYS_read", 0x3)', Constants.SYS_read.inspect)
+      assert_equal('__NR_prctl', Constants.__NR_prctl.to_s)
+      assert_equal('Constant("(O_CREAT)", 0x40)', Constants.eval('O_CREAT').inspect)
+      assert_equal(0x40, Constants.method(:O_CREAT).call.to_i)
+      # 2 < 3
+      assert_operator(2, :<, Constants.SYS_read)
+    end
+  end
+end

data/test/context_test.rb

@@ -0,0 +1,131 @@
+# encoding: ASCII-8BIT
+
+require 'test_helper'
+require 'pwnlib/context'
+
+class ContextTest < MiniTest::Test
+  include ::Pwnlib::Context
+
+  def test_update
+    context.update(arch: 'arm', os: 'windows')
+    assert_equal('arm', context.arch)
+    assert_equal('windows', context.os)
+  end
+
+  def test_local
+    context.timeout = 1
+    assert_equal(1, context.timeout)
+
+    context.local(timeout: 2) do
+      assert_equal(2, context.timeout)
+      context.timeout = 3
+      assert_equal(3, context.timeout)
+    end
+
+    assert_equal(1, context.timeout)
+
+    assert_raises(RuntimeError) do
+      context.local(timeout: 3) { raise 'QQ failed in block' }
+    end
+
+    assert_equal(1, context.timeout)
+  end
+
+  def test_clear
+    default_arch = context.arch
+    context.arch = 'arm'
+    context.clear
+    assert_equal(default_arch, context.arch)
+  end
+
+  def test_arch
+    context.arch = 'mips'
+    assert_equal('mips', context.arch)
+
+    err = assert_raises(ArgumentError) { context.arch = 'shik' }
+    assert_match(/arch must be one of/, err.message)
+    assert_equal('mips', context.arch)
+
+    context.clear
+    assert_equal(32, context.bits)
+    context.arch = 'powerpc64'
+    assert_equal(64, context.bits)
+    assert_equal('big', context.endian)
+  end
+
+  def test_bits
+    context.bits = 64
+    assert_equal(64, context.bits)
+
+    err = assert_raises(ArgumentError) { context.bits = 0 }
+    assert_match(/bits must be > 0/, err.message)
+  end
+
+  def test_bytes
+    context.bytes = 8
+    assert_equal(64, context.bits)
+    assert_equal(8, context.bytes)
+
+    context.bits = 32
+    assert_equal(4, context.bytes)
+  end
+
+  def test_endian
+    context.endian = 'le'
+    assert_equal('little', context.endian)
+
+    context.endian = 'big'
+    assert_equal('big', context.endian)
+
+    err = assert_raises(ArgumentError) { context.endian = 'SUPERBIG' }
+    assert_match(/endian must be one of/, err.message)
+  end
+
+  def test_log_level
+    context.log_level = 'error'
+    assert_equal(Logger::ERROR, context.log_level)
+
+    context.log_level = 514
+    assert_equal(514, context.log_level)
+
+    err = assert_raises(ArgumentError) { context.log_level = 'BOOM' }
+    assert_match(/log_level must be an integer or one of/, err.message)
+  end
+
+  def test_os
+    context.os = 'windows'
+    assert_equal('windows', context.os)
+
+    err = assert_raises(ArgumentError) { context.os = 'deepblue' }
+    assert_match(/os must be one of/, err.message)
+  end
+
+  def test_signed
+    context.signed = true
+    assert_equal(true, context.signed)
+
+    context.signed = 'unsigned'
+    assert_equal(false, context.signed)
+
+    err = assert_raises(ArgumentError) { context.signed = 'partial' }
+    assert_match(/signed must be boolean or one of/, err.message)
+  end
+
+  def test_timeout
+    context.timeout = 123
+    assert_equal(123, context.timeout)
+  end
+
+  def test_newline
+    context.newline = "\r\n"
+    assert_equal("\r\n", context.newline)
+  end
+
+  def test_to_s
+    assert_match(/\APwnlib::Context::ContextType\(.+\)\Z/, context.to_s)
+  end
+
+  def teardown
+    context.clear
+  end
+end

data/test/data/victim.c

@@ -0,0 +1,8 @@
+#include <stdio.h>
+
+int main() {
+  setvbuf(stdout, NULL, _IONBF, 0);
+  printf("%p\n", __builtin_return_address(0));
+  scanf("%c");
+  return 0;
+}

data/test/dynelf_test.rb

@@ -0,0 +1,48 @@
+# encoding: ASCII-8BIT
+
+require 'open3'
+
+require 'tty-platform'
+
+require 'test_helper'
+require 'pwnlib/dynelf'
+
+class DynELFTest < MiniTest::Test
+  def test_lookup
+    skip 'Only tested on linux' unless TTY::Platform.new.linux?
+    [32, 64].each do |b|
+      # TODO(hh): Use process instead of popen2
+      Open3.popen2(File.expand_path("../data/victim#{b}", __FILE__)) do |i, o, t|
+        main_ra = Integer(o.readline)
+        libc_path = nil
+        IO.readlines("/proc/#{t.pid}/maps").map(&:split).each do |s|
+          st, ed = s[0].split('-').map { |x| x.to_i(16) }
+          next unless main_ra.between?(st, ed)
+          libc_path = s[-1]
+          break
+        end
+        refute_nil(libc_path)
+
+        # TODO(hh): Use ELF instead of objdump
+        # Methods in libc might have multi-versions, so we record and check if
+        # we can find one of them.
+        h = Hash.new { |hsh, key| hsh[key] = [] }
+        symbols = `objdump -T #{libc_path}`.lines.map(&:split).select { |a| a[2] == 'DF' }
+        symbols.map { |a| h[a[-1]] << a[0].to_i(16) }
+
+        mem = open("/proc/#{t.pid}/mem", 'rb')
+        d = ::Pwnlib::DynELF.new(main_ra) do |addr|
+          mem.seek(addr)
+          mem.getc
+        end
+
+        assert_nil(d.lookup('pipi_hao_wei!'))
+        h.each do |sym, off|
+          assert_includes(off, d.lookup(sym) - d.libbase)
+        end
+
+        i.write('bye')
+      end
+    end
+  end
+end

data/test/ext_test.rb

@@ -0,0 +1,26 @@
+# encoding: ASCII-8BIT
+
+require 'test_helper'
+require 'pwnlib/ext/array'
+require 'pwnlib/ext/integer'
+require 'pwnlib/ext/string'
+
+class ExtTest < MiniTest::Test
+  # Thought that test one method in each module for each type is enough, since it's quite
+  # stupid (and meaningless) to copy the list of proxied functions to here...
+  def test_ext_string
+    assert_equal(0x4142, 'AB'.u16(endian: 'be'))
+    assert_equal([1, 1, 0, 0, 0, 1, 0, 0], "\xC4".bits)
+  end
+
+  def test_ext_integer
+    assert_equal('AB', 0x4241.p16)
+    assert_equal([0, 0, 1, 1, 0, 1, 0, 0], 0x34.bits)
+    assert_equal(2**31, 1.bitswap)
+  end
+
+  def test_ext_array
+    assert_equal("\xfe", [1, 1, 1, 1, 1, 1, 1, 0].unbits)
+    assert_equal("XX\xef\xbe\xad\xdeXX", ['XX', 0xdeadbeef, 'XX'].flat)
+  end
+end

data/test/files/use_pwn.rb

@@ -0,0 +1,34 @@
+# encoding: ASCII-8BIT
+
+# Make sure we're using local copy for local testing.
+$LOAD_PATH.unshift File.expand_path(File.join(__FILE__, '..', '..', '..', 'lib'))
+
+require 'pwn'
+
+context[arch: 'amd64']
+
+raise 'pack fail' unless pack(1) == "\x01\0\0\0\0\0\0\0"
+unless ::Pwnlib::Util::Fiddling.__send__(:context).object_id == context.object_id
+  raise 'not unique context'
+end
+unless ::Pwnlib::Context.context.object_id == context.object_id
+  raise 'not unique context'
+end
+
+# Make sure things aren't polluting Object
+begin
+  1.__send__(:context)
+  raise 'context polluting Object.'
+rescue NoMethodError
+  puts 'good'
+end
+
+begin
+  '1'.__send__(:context)
+  raise 'context polluting Object.'
+rescue NoMethodError
+  puts 'good'
+end
+
+# Make sure we can use Util::xxx::yyy directly
+raise 'pack fail' unless Util::Packing.pack(1) == "\x01\0\0\0\0\0\0\0"

data/test/files/use_pwnlib.rb

@@ -0,0 +1,19 @@
+# encoding: ASCII-8BIT
+
+# Make sure we're using local copy for local testing.
+$LOAD_PATH.unshift File.expand_path(File.join(__FILE__, '..', '..', '..', 'lib'))
+
+# TODO(Darkpi): Should we make sure ALL module works? (maybe we should).
+require 'pwnlib/util/packing'
+
+raise 'call from module fail' unless ::Pwnlib::Util::Packing.p8(0x61) == 'a'
+
+include ::Pwnlib::Util::Packing::ClassMethods
+raise 'include module and call fail' unless p8(0x61) == 'a'
+
+begin
+  ::Pwnlib::Util::Packing.context
+  raise 'context public in Pwnlib module'
+rescue NoMethodError
+  puts 'good'
+end

data/test/full_file_test.rb

@@ -0,0 +1,16 @@
+# encoding: ASCII-8BIT
+
+require 'open3'
+
+require 'test_helper'
+
+class FullFileTest < MiniTest::Test
+  parallelize_me!
+  Dir['test/files/*.rb'].each do |f|
+    fn = File.basename(f, '.rb')
+    define_method("test_#{fn}") do
+      _, stderr, status = Open3.capture3('ruby', f, binmode: true)
+      assert(status.success?, stderr)
+    end
+  end
+end

data/test/memleak_test.rb

@@ -0,0 +1,72 @@
+# encoding: ASCII-8BIT
+
+require 'open3'
+
+require 'tty-platform'
+
+require 'test_helper'
+require 'pwnlib/memleak'
+
+class MemLeakTest < MiniTest::Test
+  def setup
+    @victim = IO.binread(File.expand_path('../data/victim32', __FILE__))
+    @leak = ::Pwnlib::MemLeak.new { |addr| @victim[addr] }
+  end
+
+  def test_find_elf_base_basic
+    assert_equal(0, @leak.find_elf_base(@victim.length * 2 / 3))
+  end
+
+  def test_find_elf_base_running
+    skip 'Only tested on linux' unless TTY::Platform.new.linux?
+    [32, 64].each do |b|
+      # TODO(hh): Use process instead of popen2
+      Open3.popen2(File.expand_path("../data/victim#{b}", __FILE__)) do |i, o, t|
+        main_ra = o.readline[2...-1].to_i(16)
+        realbase = nil
+        IO.readlines("/proc/#{t.pid}/maps").map(&:split).each do |s|
+          st, ed = s[0].split('-').map { |x| x.to_i(16) }
+          next unless main_ra.between?(st, ed)
+          realbase = st
+          break
+        end
+        refute_nil(realbase)
+        mem = open("/proc/#{t.pid}/mem", 'rb')
+        l2 = ::Pwnlib::MemLeak.new do |addr|
+          mem.seek(addr)
+          mem.getc
+        end
+        assert_equal(realbase, l2.find_elf_base(main_ra))
+        mem.close
+        i.write('bye')
+      end
+    end
+  end
+
+  def test_n
+    assert_equal("\x7fELF", @leak.n(0, 4))
+    assert_equal(@victim[0xf0, 0x20], @leak.n(0xf0, 0x20))
+    assert_equal(@victim[514, 0x20], @leak.n(514, 0x20))
+  end
+
+  def test_b
+    assert_equal(@victim[0x100], @leak.b(0x100))
+    assert_equal(@victim[514], @leak.b(514))
+  end
+
+  def test_w
+    assert_equal(::Pwnlib::Util::Packing.u16(@victim[0x100, 2]), @leak.w(0x100))
+    assert_equal(::Pwnlib::Util::Packing.u16(@victim[514, 2]), @leak.w(514))
+  end
+
+  def test_d
+    assert_equal(::Pwnlib::Util::Packing.u32(@victim[0, 4]), @leak.d(0))
+    assert_equal(::Pwnlib::Util::Packing.u32(@victim[0x100, 4]), @leak.d(0x100))
+    assert_equal(::Pwnlib::Util::Packing.u32(@victim[514, 4]), @leak.d(514))
+  end
+
+  def test_q
+    assert_equal(::Pwnlib::Util::Packing.u64(@victim[0x100, 8]), @leak.q(0x100))
+    assert_equal(::Pwnlib::Util::Packing.u64(@victim[514, 8]), @leak.q(514))
+  end
+end

data/test/reg_sort_test.rb

@@ -0,0 +1,41 @@
+# encoding: ASCII-8BIT
+require 'test_helper'
+require 'pwnlib/reg_sort'
+
+class RegSortTest < MiniTest::Test
+  include ::Pwnlib::RegSort::ClassMethods
+
+  def setup
+    @regs = %w(a b c d x y z)
+  end
+
+  def test_normal
+    assert_equal([['mov', 'a', 1], ['mov', 'b', 2]], regsort({ a: 1, b: 2 }, @regs))
+  end
+
+  def test_post_mov
+    assert_equal([['mov', 'a', 1], %w(mov b a)], regsort({ a: 1, b: 1 }, @regs))
+    assert_equal([%w(mov c a), ['mov', 'a', 1], %w(mov b a)], regsort({ a: 1, b: 1, c: 'a' }, @regs))
+  end
+
+  def test_pseudoforest
+    # only one connected component
+    assert_equal([%w(mov b a), ['mov', 'a', 1]], regsort({ a: 1, b: 'a' }, @regs))
+    assert_equal([['mov', 'c', 3], %w(xchg a b)], regsort({ a: 'b', b: 'a', c: 3 }, @regs))
+    assert_equal([%w(mov c b), %w(xchg a b)], regsort({ a: 'b', b: 'a', c: 'b' }, @regs))
+    assert_equal([%w(mov x 1), %w(mov y z), %w(mov z c), %w(xchg a b), %w(xchg b c)],
+                 regsort({ a: 'b', b: 'c', c: 'a', x: '1', y: 'z', z: 'c' }, @regs))
+
+    # more than one connected components
+    assert_equal([%w(xchg a b), %w(xchg c d)], regsort({ a: 'b', b: 'a', c: 'd', d: 'c' }, @regs))
+    assert_equal([%w(mov c b), %w(mov d b), %w(mov z x), %w(xchg a b), %w(xchg x y)],
+                 regsort({ a: 'b', b: 'a', c: 'b', d: 'b', x: 'y', y: 'x', z: 'x' }, @regs))
+  end
+
+  def test_raise
+    err = assert_raises(ArgumentError) do
+      regsort({ a: 1 }, ['b'])
+    end
+    assert_match(/Unknown register!/, err.message)
+  end
+end

data/test/test_helper.rb

@@ -0,0 +1,13 @@
+require 'codeclimate-test-reporter'
+
+require 'simplecov'
+SimpleCov.formatter = SimpleCov::Formatter::MultiFormatter.new(
+  [SimpleCov::Formatter::HTMLFormatter, CodeClimate::TestReporter::Formatter]
+)
+SimpleCov.start do
+  add_filter '/test/'
+end
+
+require 'minitest/autorun'
+require 'minitest/unit'
+require 'minitest/hell'