pwntools 0.1.0 → 1.0.0

data/test/shellcraft/registers_test.rb

@@ -0,0 +1,32 @@
+# encoding: ASCII-8BIT
+
+require 'test_helper'
+
+require 'pwnlib/shellcraft/registers'
+
+class RegistersTest < MiniTest::Test
+  include ::Pwnlib::Shellcraft::Registers
+
+  def test_get_register
+    assert_instance_of(::Pwnlib::Shellcraft::Registers::Register, get_register('rdi'))
+    assert_nil(get_register('meow'))
+  end
+
+  def test_regtsiter?
+    assert_equal(true, register?(:ax))
+    assert_equal(true, register?('ax'))
+    assert_equal(true, register?('r8'))
+    assert_equal(true, register?('r15b'))
+    assert_equal(true, register?('r15w'))
+    assert_equal(false, register?('xdd'))
+    assert_equal(false, register?(''))
+  end
+
+  def test_register
+    assert_equal(8, get_register('rdi').bytes)
+    assert_equal(16, get_register('ax').bits)
+    assert_equal('r10d', get_register('r10d').name)
+    assert_equal('r10d', get_register('r10d').to_s)
+    assert_equal('Register(r10d)', get_register('r10d').inspect)
+  end
+end

data/test/shellcraft/ret_test.rb

@@ -0,0 +1,30 @@
+# encoding: ASCII-8BIT
+
+require 'test_helper'
+
+require 'pwnlib/context'
+require 'pwnlib/shellcraft/shellcraft'
+
+class RetTest < MiniTest::Test
+  include ::Pwnlib::Context
+
+  def setup
+    @shellcraft = ::Pwnlib::Shellcraft::Shellcraft.instance
+  end
+
+  def test_amd64
+    context.local(arch: 'amd64') do
+      assert_equal("  ret\n", @shellcraft.ret)
+      assert_equal("  mov rax, rdi\n  ret\n", @shellcraft.ret(:rdi))
+      assert_equal("  xor eax, eax /* 0 */\n  ret\n", @shellcraft.ret(0))
+      assert_equal(<<-'EOS', @shellcraft.ret(0x100000000))
+  mov rax, 0x101010201010101
+  push rax
+  mov rax, 0x101010301010101
+  xor [rsp], rax /* 0x100000000 == 0x101010201010101 ^ 0x101010301010101 */
+  pop rax
+  ret
+      EOS
+    end
+  end
+end

data/test/shellcraft/setregs_test.rb

@@ -0,0 +1,62 @@
+# encoding: ASCII-8BIT
+
+require 'test_helper'
+
+require 'pwnlib/context'
+require 'pwnlib/shellcraft/shellcraft'
+
+class SetregsTest < MiniTest::Test
+  include ::Pwnlib::Context
+
+  def setup
+    @shellcraft = ::Pwnlib::Shellcraft::Shellcraft.instance
+  end
+
+  def test_amd64
+    context.local(arch: 'amd64') do
+      assert_equal(<<-'EOS', @shellcraft.setregs(rax: 1, rbx: 'rax'))
+  mov rbx, rax
+  push 1
+  pop rax
+      EOS
+      assert_equal(<<-'EOS', @shellcraft.setregs(rax: 'SYS_write', rbx: 'rax'))
+  mov rbx, rax
+  push 1 /* (SYS_write) */
+  pop rax
+      EOS
+      assert_equal(<<-'EOS', @shellcraft.setregs(rax: 'rbx', rbx: 'rax', rcx: 'rbx'))
+  mov rcx, rbx
+  xchg rax, rbx
+      EOS
+      assert_equal(<<-'EOS', @shellcraft.setregs(rax: 1, rdx: 0))
+  push 1
+  pop rax
+  cdq /* rdx=0 */
+      EOS
+      # issue #50
+      assert_equal(<<-'EOS', @shellcraft.setregs(rdi: :rax, rsi: :rdi))
+  mov rsi, rdi
+  mov rdi, rax
+      EOS
+    end
+  end
+
+  def test_i386
+    context.local(arch: 'i386') do
+      assert_equal(<<-EOS, @shellcraft.setregs(eax: 1, ebx: 'eax'))
+  mov ebx, eax
+  push 1
+  pop eax
+      EOS
+      assert_equal(<<-EOS, @shellcraft.setregs(eax: 'ebx', ebx: 'eax', ecx: 'ebx'))
+  mov ecx, ebx
+  xchg eax, ebx
+      EOS
+      assert_equal(<<-'EOS', @shellcraft.setregs(eax: 1, edx: 0))
+  push 1
+  pop eax
+  cdq /* edx=0 */
+      EOS
+    end
+  end
+end

data/test/shellcraft/shellcraft_test.rb

@@ -0,0 +1,28 @@
+# encoding: ASCII-8BIT
+
+require 'test_helper'
+
+require 'pwnlib/context'
+require 'pwnlib/shellcraft/shellcraft'
+
+class ShellcraftTest < MiniTest::Test
+  include ::Pwnlib::Context
+
+  def setup
+    @shellcraft = ::Pwnlib::Shellcraft::Shellcraft.instance
+  end
+
+  def test_respond
+    context.local(arch: 'amd64') do
+      # Check respond_to_missing? is well defined
+      assert @shellcraft.respond_to?(:mov)
+      assert @shellcraft.method(:sh)
+    end
+    refute @shellcraft.respond_to?(:linux)
+    assert_raises(NoMethodError) { @shellcraft.meow }
+
+    context.local(arch: 'arm') do
+      refute @shellcraft.respond_to?(:mov)
+    end
+  end
+end

data/test/test_helper.rb

@@ -1,6 +1,7 @@
 require 'codeclimate-test-reporter'
-
+require 'rainbow'
 require 'simplecov'
+
 SimpleCov.formatter = SimpleCov::Formatter::MultiFormatter.new(
   [SimpleCov::Formatter::HTMLFormatter, CodeClimate::TestReporter::Formatter]
 )
@@ -11,3 +12,13 @@ end
 require 'minitest/autorun'
 require 'minitest/unit'
 require 'minitest/hell'
+
+module MiniTest
+  class Test
+    def before_setup
+      super
+      # Default to disable coloring for easier testing.
+      Rainbow.enabled = false
+    end
+  end
+end

data/test/timer_test.rb

@@ -0,0 +1,23 @@
+# encoding: ASCII-8BIT
+
+require 'test_helper'
+
+require 'pwnlib/timer'
+
+class TimerTest < MiniTest::Test
+  include ::Pwnlib
+
+  def test_countdown
+    t = Timer.new
+    refute(t.started?)
+    refute(t.active?)
+    assert_equal('DARKHH QQ', t.countdown(0.1) { 'DARKHH QQ' })
+    exception = assert_raises(RuntimeError) { t.countdown(0.1) { t.countdown(0.1) {} } }
+    assert_equal('Nested countdown not permitted', exception.message)
+    t.timeout = 0.514
+    exception = assert_raises(RuntimeError) { t.countdown(0.1) { t.timeout = :forever } }
+    assert_equal("Can't change timeout when countdown", exception.message)
+    t.countdown(0.1) { assert(t.started?) }
+    t.countdown(0.1) { assert(t.active?) }
+  end
+end

data/test/tubes/buffer_test.rb

@@ -0,0 +1,45 @@
+# encoding: ASCII-8BIT
+
+require 'test_helper'
+
+require 'pwnlib/tubes/buffer'
+
+class BufferTest < MiniTest::Test
+  def test_add
+    b = ::Pwnlib::Tubes::Buffer.new
+    b.add('A' * 10)
+    b.add('B' * 10)
+    assert_equal(20, b.size)
+    assert_equal('A', b.get(1))
+    assert_equal(19, b.size)
+    assert_equal(false, b.empty?)
+    assert_equal('AAAAAAAAABBBBBBBBBB', b.get(9999))
+    assert_equal(0, b.size)
+    assert_equal(true, b.empty?)
+    assert_equal('', b.get(1))
+  end
+
+  def test_unget
+    b = ::Pwnlib::Tubes::Buffer.new
+    b.add('hello')
+    b.add('world')
+    assert_equal('hello', b.get(5))
+    b.unget('goodbye')
+    assert_equal('goodbyeworld', b.get)
+  end
+
+  def test_buffer
+    b = ::Pwnlib::Tubes::Buffer.new
+    b2 = ::Pwnlib::Tubes::Buffer.new
+    b.add('hello')
+    b2.add('world')
+    b.add(b2)
+    assert_equal('hello', b.get(5))
+    assert_equal(false, b2.empty?)
+    assert_equal('world', b2.get)
+    b2.add('goodbye')
+    b.unget(b2)
+    assert_equal('goodbyeworld', b.get)
+    assert_equal('goodbye', b2.get)
+  end
+end

data/test/tubes/sock_test.rb

@@ -0,0 +1,68 @@
+# encoding: ASCII-8BIT
+
+require 'open3'
+
+require 'test_helper'
+
+require 'pwnlib/tubes/sock'
+
+class SockTest < MiniTest::Test
+  include ::Pwnlib::Tubes
+  ECHO_FILE = File.expand_path('../data/echo.rb', __dir__)
+  BIND_PORT = 31_337
+
+  def popen_echo(data)
+    Open3.popen2("bundle exec ruby #{ECHO_FILE} #{BIND_PORT}") do |_i, o, _t|
+      o.gets
+      s = Sock.new('localhost', BIND_PORT)
+      yield s, data, o
+    end
+  end
+
+  def test_io
+    popen_echo('DARKHH') do |s, data, _o|
+      s.io.puts(data)
+      rs, = IO.select([s.io])
+      refute_nil(rs)
+      assert_equal(data, s.io.readpartial(data.size))
+    end
+  end
+
+  def test_sock
+    popen_echo('DARKHH') do |s, data, o|
+      s.puts(data)
+      assert_equal(data + "\n", s.gets)
+      o.gets
+      s.puts(514)
+      sleep(0.1) # Wait for connection reset
+      assert_raises(EOFError) { s.puts(514) }
+      assert_raises(EOFError) { s.puts(514) }
+      assert_raises(EOFError) { s.recv }
+      assert_raises(EOFError) { s.recv }
+    end
+  end
+
+  def test_close
+    popen_echo('DARKHH') do |s, _data, _o|
+      s.close
+      assert_raises(EOFError) { s.puts(514) }
+      assert_raises(EOFError) { s.puts(514) }
+      assert_raises(EOFError) { s.recv }
+      assert_raises(EOFError) { s.recv }
+      assert_raises(ArgumentError) { s.close(:hh) }
+    end
+
+    popen_echo('DARKHH') do |s, _data, _o|
+      3.times { s.close(:read) }
+      3.times { s.close(:recv) }
+      3.times { s.close(:send) }
+      3.times { s.close(:write) }
+      assert_raises(EOFError) { s.puts(514) }
+      assert_raises(EOFError) { s.puts(514) }
+      assert_raises(EOFError) { s.recv }
+      assert_raises(EOFError) { s.recv }
+      3.times { s.close }
+      assert_raises(ArgumentError) { s.close(:shik) }
+    end
+  end
+end

data/test/tubes/tube_test.rb

@@ -0,0 +1,241 @@
+# encoding: UTF-8
+
+# This test use UTF-8 encoding for strings since the output for hexdump contains lots of UTF-8 characters.
+
+require 'test_helper'
+
+require 'pwnlib/context'
+require 'pwnlib/logger'
+require 'pwnlib/tubes/tube'
+
+class TubeTest < MiniTest::Test
+  include ::Pwnlib::Context
+  include ::Pwnlib::Tubes
+
+  def setup
+    @old_log = ::Pwnlib::Logger.log.dup
+    @log = ::Pwnlib::Logger::LoggerType.new
+
+    class << @log
+      def clear
+        @logdev = StringIO.new
+      end
+
+      def string
+        @logdev.string
+      end
+    end
+
+    ::Pwnlib::Logger.instance_variable_set(:@log, @log)
+  end
+
+  def hello_tube
+    t = Tube.new
+
+    class << t
+      def buf
+        @buf ||= ''
+      end
+
+      private
+
+      def recv_raw(_n)
+        'Hello, world'
+      end
+
+      def timeout_raw=(timeout)
+        @timeout = timeout == :forever ? nil : timeout
+      end
+
+      def send_raw(data)
+        buf << data
+      end
+    end
+
+    t
+  end
+
+  def basic_tube
+    t = Tube.new
+
+    class << t
+      def recv_raw(_n)
+        raise EOFError if io.eof?
+        io.read
+      end
+
+      def send_raw(data)
+        io.write(data)
+      end
+
+      def timeout_raw=(timeout)
+        @timeout = timeout == :forever ? nil : timeout
+      end
+
+      def io
+        @fakeio ||= Tempfile.new('pwntools_ruby_test')
+      end
+    end
+
+    t
+  end
+
+  def test_recv
+    t = hello_tube
+    assert_equal('Hello, world', t.recv)
+    assert_equal('Hello, world', t.recv)
+    t.unrecv('Woohoo')
+    assert_equal('Woohoo', t.recv)
+    assert_equal('Hello, world', t.recv)
+    assert_equal('H', t.recvn(1))
+    assert_equal('ello, w', t.recvn(7))
+    assert_equal('orldH', t.recvn(5))
+    assert_equal('ello, world', t.recv)
+
+    context.local(log_level: 'debug') do
+      @log.clear
+      t = hello_tube
+      assert_equal('Hello, world', t.recv)
+      assert_equal(<<-'EOS', @log.string.encode('UTF-8'))
+[DEBUG] Received 0xc bytes:
+    00000000  48 65 6c 6c  6f 2c 20 77  6f 72 6c 64               │Hell│o, w│orld│
+    0000000c
+      EOS
+    end
+  end
+
+  def test_recvuntil
+    t = hello_tube
+    assert_equal('Hello, ', t.recvuntil(' '))
+    assert_equal('worldHello, ', t.recvuntil(' '))
+    t.unrecv('Hello, world!')
+    assert_equal('Hello,', t.recvuntil(' wor', drop: true))
+    assert_equal('', t.recvuntil('DARKHH', drop: true, timeout: 0.1))
+
+    t = basic_tube
+    t.unrecv('meow')
+    assert_equal('', t.recvuntil('DARKHH'))
+    assert_equal('meow', t.recv)
+  end
+
+  def test_recvline
+    t = Tube.new
+    t.unrecv("Foo\nBar\r\nBaz\n")
+    assert_equal("Foo\n", t.recvline)
+    assert_equal("Bar\r\n", t.recvline)
+    assert_equal('Baz', t.recvline(drop: true))
+    context.local(newline: "\r\n") do
+      t = Tube.new
+      t.unrecv("Foo\nBar\r\nBaz\n")
+      assert_equal("Foo\nBar", t.recvline(drop: true))
+    end
+
+    t = basic_tube
+    t.unrecv('Hello, world')
+    assert_equal('', t.recvline)
+    assert_equal('Hello, world', t.recv)
+  end
+
+  def test_recvpred
+    t = hello_tube
+    r = /H.*w/
+    10.times { assert_match(r, t.recvpred { |data| data =~ r }) }
+    r = /H.*W/
+    assert_match('', t.recvpred(timeout: 0.01) { |data| data =~ r })
+    t = basic_tube
+    t.unrecv('darkhh')
+    assert_match('', t.recvpred { |data| data =~ r })
+  end
+
+  def test_recvregex
+    t = hello_tube
+    r = /[aeiou]/
+    5.times { assert_match(r, t.recvregex(r)) }
+    r = /[wl][aeiou]/
+    5.times { assert_match(r, t.recvregex(r)) }
+  end
+
+  def test_send
+    t = hello_tube
+    t.write('DARKHH')
+    assert_equal('DARKHH', t.buf)
+    t.write(' QQ')
+    assert_equal('DARKHH QQ', t.buf)
+    t.write(333)
+    assert_equal('DARKHH QQ333', t.buf)
+
+    context.local(log_level: 'debug') do
+      @log.clear
+      data = (0..40).map(&:chr).join
+      t = hello_tube
+      t.write(data)
+      assert_equal(data, t.buf)
+      assert_equal(<<-'EOS', @log.string.encode('UTF-8'))
+[DEBUG] Sent 0x29 bytes:
+    00000000  00 01 02 03  04 05 06 07  08 09 0a 0b  0c 0d 0e 0f  │····│····│····│····│
+    00000010  10 11 12 13  14 15 16 17  18 19 1a 1b  1c 1d 1e 1f  │····│····│····│····│
+    00000020  20 21 22 23  24 25 26 27  28                        │ !"#│$%&'│(│
+    00000029
+      EOS
+
+      @log.clear
+      t.puts('meow')
+      assert_equal(<<-'EOS', @log.string.encode('UTF-8'))
+[DEBUG] Sent 0x5 bytes:
+    00000000  6d 65 6f 77  0a                                     │meow│·│
+    00000005
+      EOS
+    end
+  end
+
+  def test_sendline
+    t = hello_tube
+    t.write('DARKHH')
+    assert_equal('DARKHH', t.buf)
+    t.puts(' QQ')
+    assert_equal("DARKHH QQ\n", t.buf)
+  end
+
+  FLAG_FILE = File.expand_path('../data/flag', __dir__)
+  def test_interact_send
+    save_stdin = $stdin.dup
+    $stdin = File.new(FLAG_FILE, File::RDONLY)
+    @log.clear
+    begin
+      t = basic_tube
+      t.interact
+    rescue EOFError
+      t.io.rewind
+      assert_equal(IO.binread(FLAG_FILE), t.io.read)
+    end
+    assert_equal("[INFO] Switching to interactive mode\n", @log.string)
+    $stdin.close
+    t.io.close
+    $stdin = save_stdin
+  end
+
+  def test_interact_recv
+    save_stdin = $stdin.dup
+    save_stdout = $stdout.dup
+    $stdin = UDPSocket.new
+    $stdout = Tempfile.new('pwntools_ruby_test')
+    @log.clear
+    begin
+      t = basic_tube
+      t.instance_variable_set(:@fakeio, File.new(FLAG_FILE, File::RDONLY))
+      t.interact
+    rescue EOFError
+      $stdout.rewind
+      assert_equal(IO.binread(FLAG_FILE), $stdout.read)
+    end
+    assert_equal("[INFO] Switching to interactive mode\n", @log.string)
+    $stdout.close
+    t.io.close
+    $stdin = save_stdin
+    $stdout = save_stdout
+  end
+
+  def teardown
+    ::Pwnlib::Logger.instance_variable_set(:@log, @old_log)
+  end
+end