pwntools 0.1.0 → 1.2.1

data/test/shellcraft/registers_test.rb

@@ -0,0 +1,33 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+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,31 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+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,63 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+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,30 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+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
+      err = assert_raises(::Pwnlib::Errors::UnsupportedArchError) { @shellcraft.respond_to?(:mov) }
+      assert_equal("Can't use shellcraft under architecture \"arm\".", err.message)
+    end
+  end
+end

data/test/test_helper.rb

@@ -1,8 +1,11 @@
-require 'codeclimate-test-reporter'
+# frozen_string_literal: true
 
+require 'rainbow'
 require 'simplecov'
+require 'tty/platform'
+
 SimpleCov.formatter = SimpleCov::Formatter::MultiFormatter.new(
-  [SimpleCov::Formatter::HTMLFormatter, CodeClimate::TestReporter::Formatter]
+  [SimpleCov::Formatter::HTMLFormatter]
 )
 SimpleCov.start do
   add_filter '/test/'
@@ -11,3 +14,59 @@ 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
+
+    def linux_only(msg = 'Only tested on Linux')
+      skip msg unless TTY::Platform.new.linux?
+    end
+
+    def skip_windows(msg = 'Skip on Windows')
+      skip msg if TTY::Platform.new.windows?
+    end
+
+    # Methods for hooking logger,
+    # require 'pwnlib/logger' before using these methods.
+
+    def log_null(&block)
+      File.open(File::NULL, 'w') { |f| log_hook(f, &block) }
+    end
+
+    def log_stdout(&block)
+      log_hook($stdout, &block)
+    end
+
+    def log_hook(obj)
+      old = ::Pwnlib::Logger.log.instance_variable_get(:@logdev)
+      ::Pwnlib::Logger.log.instance_variable_set(:@logdev, obj)
+      begin
+        yield
+      ensure
+        ::Pwnlib::Logger.log.instance_variable_set(:@logdev, old)
+      end
+    end
+
+    # Hooks +$stdin+.
+    #
+    # @param [IO] io
+    #   IO object to be used as +$stdin+.
+    #
+    # @return [Object]
+    #   Returns what the block returned.
+    def hook_stdin(io)
+      org_stdin = $stdin
+      $stdin = io
+      begin
+        yield
+      ensure
+        $stdin = org_stdin
+      end
+    end
+  end
+end

data/test/timer_test.rb

@@ -0,0 +1,42 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+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) do
+      t.countdown(0.1) { t.timeout = :forever }
+    end
+    assert_equal("Can't change timeout when countdown", exception.message)
+
+    t.countdown(0.1) { assert(t.started?) }
+    t.countdown(0.1) { assert(t.active?) }
+
+    assert_raises(::Pwnlib::Errors::TimeoutError) do
+      t.countdown(0.1) { sleep(0.2) }
+    end
+  end
+
+  # Check #115 fixed
+  def test_nested_countdown_false_positve_115
+    t = Timer.new
+    assert_raises(::Pwnlib::Errors::TimeoutError) do
+      t.countdown(0.01) { sleep(1) }
+    end
+    t.countdown(0.01) {}
+  end
+end

data/test/tubes/buffer_test.rb

@@ -0,0 +1,46 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+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/process_test.rb

@@ -0,0 +1,105 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+require 'socket'
+require 'tty-platform'
+
+require 'test_helper'
+
+require 'pwnlib/context'
+require 'pwnlib/errors'
+require 'pwnlib/tubes/process'
+
+class ProcessTest < MiniTest::Test
+  include ::Pwnlib::Context
+
+  def setup
+    skip_windows
+  end
+
+  def test_io
+    cat = ::Pwnlib::Tubes::Process.new('cat')
+    cat.puts('HAHA')
+    assert_equal("HAHA\n", cat.gets)
+    assert_raises(::Pwnlib::Errors::TimeoutError) { cat.gets(timeout: 0.1) }
+    cat.puts('HAHA2')
+    assert_equal("HAHA2\n", cat.gets)
+  end
+
+  def test_env
+    data = ::Pwnlib::Tubes::Process.new('env').read
+    assert_match('PATH=', data)
+    data = ::Pwnlib::Tubes::Process.new('env', env: { 'FOO' => 'BAR' }).read
+    assert_equal("FOO=BAR\n", data)
+  end
+
+  def test_aslr
+    linux_only
+
+    map1 = ::Pwnlib::Tubes::Process.new('cat /proc/self/maps', aslr: false).read
+    map2 = ::Pwnlib::Tubes::Process.new(['cat', '/proc/self/maps'], aslr: false).read
+    assert_match('/bin/cat', map1) # make sure it read something
+    assert_equal(map1, map2)
+  end
+
+  def test_eof
+    ls = ::Pwnlib::Tubes::Process.new(['ls', '-la'])
+    assert_match(/total/, ls.gets)
+    assert_raises(::Pwnlib::Errors::EndOfTubeError) { loop { ls.write('anything') } }
+    assert_raises(::Pwnlib::Errors::EndOfTubeError) { loop { ls.gets } }
+  end
+
+  def test_shutdown
+    cat = ::Pwnlib::Tubes::Process.new('cat')
+    assert_raises(::Pwnlib::Errors::TimeoutError) { cat.recvn(1, timeout: 0.1) }
+    cat.shutdown(:write) # This should cause `cat` dead
+    assert_raises(::Pwnlib::Errors::EndOfTubeError) { cat.recvn(1, timeout: 0.1) }
+    cat.shutdown
+
+    cat = ::Pwnlib::Tubes::Process.new('cat')
+    cat.shutdown(:read)
+    assert_raises(::Pwnlib::Errors::EndOfTubeError) { cat.recvn(1) }
+    cat.shutdown
+    assert_raises(ArgumentError) { cat.shutdown(:zz) }
+  end
+
+  def test_kill
+    cat = ::Pwnlib::Tubes::Process.new('cat')
+    cat.kill
+    assert_raises(::Pwnlib::Errors::EndOfTubeError) { cat.recvn(1) }
+  end
+
+  def test_tty
+    tty_test = proc do |*args, raw: true|
+      in_, out = args.map { |v| v ? :pty : :pipe }
+      process = ::Pwnlib::Tubes::Process.new('ruby -e "p [STDIN.tty?, STDOUT.tty?]"',
+                                             in: in_, out: out, raw: raw)
+      process.gets
+    end
+    assert_equal("[false, false]\n", tty_test.call(false, false))
+    assert_equal("[true, false]\n", tty_test.call(true, false))
+    assert_equal("[false, true]\n", tty_test.call(false, true))
+    assert_equal("[false, true]\r\n", tty_test.call(false, true, raw: false))
+
+    cat = ::Pwnlib::Tubes::Process.new('cat', in: :pty, out: :pty, raw: false)
+    cat.puts('Hi')
+    # In cooked mode, tty should echo the input, so we can gets twice.
+    assert_equal("Hi\r\n", cat.gets)
+    assert_equal("Hi\r\n", cat.gets)
+    class << cat
+      # hook shutdown to silence the +cat: -: Input/output error+ message.
+      def shutdown; end
+    end
+    cat.close
+  end
+
+  def test_interact
+    ls = ::Pwnlib::Tubes::Process.new('ls Gemfile*')
+    saved = $stdin
+    # prevents stdin being closed
+    $stdin = UDPSocket.new
+    assert_output("Gemfile\nGemfile.lock\n") { context.local(log_level: :fatal) { ls.interact } }
+    $stdin.close
+    $stdin = saved
+  end
+end

data/test/tubes/serialtube_test.rb

@@ -0,0 +1,162 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+require 'open3'
+
+require 'test_helper'
+
+require 'pwnlib/tubes/serialtube'
+
+module Pwnlib
+  module Tubes
+    class SerialTube
+      def break_encapsulation
+        @conn.close
+      end
+    end
+  end
+end
+
+class SerialTest < MiniTest::Test
+  include ::Pwnlib::Tubes
+
+  def open_pair
+    Open3.popen3('socat -d -d pty,raw,echo=0 pty,raw,echo=0') do |_i, _o, stderr, thread|
+      devs = []
+      2.times do
+        devs << stderr.readline.chomp.split.last
+        # First pattern matches Linux, second is macOS
+        raise IOError, 'Could not create serial crosslink' if devs.last !~ %r{^(/dev/pts/[0-9]+|/dev/ttys[0-9]+)$}
+      end
+      # To ensure socat have finished setup
+      stderr.gets('starting data transfer loop')
+
+      serial = SerialTube.new devs[1], convert_newlines: false
+
+      begin
+        File.open devs[0], 'r+' do |file|
+          file.set_encoding 'default'.encoding
+          yield file, serial, thread
+        end
+      ensure
+        ::Process.kill('SIGTERM', thread.pid) if thread.alive?
+      end
+    end
+  end
+
+  def random_string(length)
+    Random.rand(36**length).to_s(36).rjust(length, '0')
+  end
+
+  def test_raise
+    skip_windows
+    open_pair do |_file, serial, thread|
+      ::Process.kill('SIGTERM', thread.pid)
+      # ensure the process has been killed
+      thread.value
+      assert_raises(Pwnlib::Errors::EndOfTubeError) { serial.puts('a') }
+    end
+    open_pair do |_file, serial|
+      serial.break_encapsulation
+      assert_raises(Pwnlib::Errors::EndOfTubeError) { serial.recv(1, timeout: 2) }
+    end
+  end
+
+  def test_recv
+    skip_windows
+    open_pair do |file, serial|
+      # recv, recvline
+      rs = random_string 24
+      file.puts rs
+      result = serial.recv 8, timeout: 1
+
+      assert_equal(rs[0...8], result)
+      result = serial.recv 8
+      assert_equal(rs[8...16], result)
+      result = serial.recvline.chomp
+      assert_equal(rs[16..-1], result)
+
+      assert_raises(Pwnlib::Errors::TimeoutError) { serial.recv(1, timeout: 0.2) }
+
+      # recvpred
+      rs = random_string 12
+      file.print rs
+      result = serial.recvpred do |data|
+        data[-6..-1] == rs[-6..-1]
+      end
+      assert_equal rs, result
+
+      assert_raises(Pwnlib::Errors::TimeoutError) { serial.recv(1, timeout: 0.2) }
+
+      # recvn
+      rs = random_string 6
+      file.print rs
+      result = ''
+      assert_raises(Pwnlib::Errors::TimeoutError) do
+        result = serial.recvn 120, timeout: 1
+      end
+      assert_empty result
+      file.print rs
+      result = serial.recvn 12
+      assert_equal rs * 2, result
+
+      assert_raises(Pwnlib::Errors::TimeoutError) { serial.recv(1, timeout: 0.2) }
+
+      # recvuntil
+      rs = random_string 12
+      file.print "#{rs}|"
+      result = serial.recvuntil('|').chomp('|')
+      assert_equal rs, result
+
+      assert_raises(Pwnlib::Errors::TimeoutError) { serial.recv(1, timeout: 0.2) }
+
+      # gets
+      rs = random_string 24
+      file.puts rs
+      result = serial.gets 12
+      assert_equal rs[0...12], result
+      result = serial.gets.chomp
+      assert_equal rs[12..-1], result
+
+      assert_raises(Pwnlib::Errors::TimeoutError) { serial.recv(1, timeout: 0.2) }
+    end
+  end
+
+  def test_send
+    skip_windows
+    open_pair do |file, serial|
+      # send, sendline
+      rs = random_string 24
+      # rubocop:disable Style/Send
+      # Justification: This isn't Object#send, false positive.
+      serial.send rs[0...12]
+      # rubocop:enable Style/Send
+      serial.sendline rs[12...24]
+      result = file.readline.chomp
+      assert_equal rs, result
+
+      # puts
+      r1 = random_string 4
+      r2 = random_string 4
+      r3 = random_string 4
+      serial.puts r1, r2, r3
+      result = ''
+      3.times do
+        result += file.readline.chomp
+      end
+      assert_equal r1 + r2 + r3, result
+    end
+  end
+
+  def test_close
+    skip_windows
+    open_pair do |_file, serial|
+      serial.close
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { serial.puts(514) }
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { serial.puts(514) }
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { serial.recv }
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { serial.recv }
+      assert_raises(ArgumentError) { serial.close(:hh) }
+    end
+  end
+end