pwntools 0.1.0 → 1.2.1

data/test/tubes/sock_test.rb

@@ -0,0 +1,68 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+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__)
+
+  def popen_echo(data)
+    Open3.popen2("ruby #{ECHO_FILE}") do |_i, o, _t|
+      port = o.gets.split.last.to_i
+      s = Sock.new('127.0.0.1', port)
+      yield s, data, o
+    end
+  end
+
+  def test_sock
+    popen_echo('DARKHH') do |s, data, _o|
+      s.sock.puts(data)
+      rs, = IO.select([s.sock])
+      refute_nil(rs)
+      assert_equal(data, s.sock.readpartial(data.size))
+    end
+  end
+
+  def test_eof
+    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(::Pwnlib::Errors::EndOfTubeError) { s.puts(514) }
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { s.puts(514) }
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { s.recv }
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { s.recv }
+    end
+  end
+
+  def test_close
+    popen_echo('DARKHH') do |s, _data, _o|
+      s.close
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { s.puts(514) }
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { s.puts(514) }
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { s.recv }
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { 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(::Pwnlib::Errors::EndOfTubeError) { s.puts(514) }
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { s.puts(514) }
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { s.recv }
+      assert_raises(::Pwnlib::Errors::EndOfTubeError) { s.recv }
+      3.times { s.close }
+      assert_raises(ArgumentError) { s.close(:shik) }
+    end
+  end
+end

data/test/tubes/tube_test.rb

@@ -0,0 +1,320 @@
+# encoding: UTF-8
+# frozen_string_literal: true
+
+# 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 ::Pwnlib::Errors::EndOfTubeError 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
+        @io ||= Tempfile.new('pwntools_ruby_test')
+      end
+
+      def io_out
+        io
+      end
+    end
+
+    t
+  end
+
+  def test_not_implement
+    t = Tube.new
+    # io_out
+    assert_raises(NotImplementedError) { context.local(log_level: :fatal) { t.interact } }
+    # send_raw
+    assert_raises(NotImplementedError) { t.puts }
+    # timeout_raw=
+    assert_raises(NotImplementedError) { t.recv(timeout: 1) }
+    class << t
+      def timeout_raw=(_) end
+    end
+    # recv_raw
+    assert_raises(NotImplementedError) { t.gets }
+  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))
+
+    # test timeout
+    assert_raises(::Pwnlib::Errors::TimeoutError) { t.recvuntil('DARKHH', drop: true, timeout: 0.1) }
+
+    t = basic_tube
+    t.unrecv('meow')
+    # test EOF
+    assert_raises(::Pwnlib::Errors::EndOfTubeError) { 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')
+    # test EOF
+    assert_raises(::Pwnlib::Errors::EndOfTubeError) { assert_equal('', t.recvline) }
+    assert_equal('Hello, world', t.recv)
+  end
+
+  def test_gets
+    t = Tube.new
+    t.unrecv("Foo\nBar\r\nBaz\n")
+    assert_equal("Foo\n", t.gets)
+    assert_equal("Bar\r\n", t.gets)
+    assert_equal('Baz', t.gets(drop: true))
+
+    t = hello_tube
+    assert_equal('Hello,', t.gets(','))
+    assert_equal(' world', t.gets('H', drop: true))
+    assert_equal('ello', t.gets(4))
+    assert_raises(ArgumentError) { t.gets(t) }
+  end
+
+  def test_recvpred
+    t = hello_tube
+    r = /H.*w/
+    10.times { assert_match(r, t.recvpred { |data| data =~ r }) }
+    r = /H.*W/
+
+    # test timeout
+    assert_raises(::Pwnlib::Errors::TimeoutError) { t.recvpred(timeout: 0.01) { |data| data =~ r } }
+    t = basic_tube
+    t.unrecv('darkhh')
+    # test EOF
+    assert_raises(::Pwnlib::Errors::EndOfTubeError) { 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_recvall
+    t = basic_tube
+    t.unrecv('meow')
+    assert_equal('meow', t.recvall)
+    assert_equal('', t.recvall)
+  end
+
+  def test_send
+    t = hello_tube
+    assert_equal(6, t.write('DARKHH'))
+    assert_equal('DARKHH', t.buf)
+    assert_equal(3, t.write(' QQ'))
+    assert_equal('DARKHH QQ', t.buf)
+    assert_equal(3, 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)
+    assert_equal(4, t.sendline(' QQ'))
+    assert_equal("DARKHH QQ\n", t.buf)
+    assert_equal(1, t.sendline(''))
+    assert_equal("DARKHH QQ\n\n", t.buf)
+  end
+
+  def test_puts
+    t = hello_tube
+    assert_equal(1, t.puts)
+    assert_equal("\n", t.buf)
+    assert_equal(17, t.puts("darkhh i4 so sad\n"))
+    assert_equal("\ndarkhh i4 so sad\n", t.buf)
+
+    t = hello_tube
+    assert_equal(14, t.puts('shik', 'hao', '', 123))
+    assert_equal("shik\nhao\n\n123\n", t.buf)
+
+    t = hello_tube
+    assert_equal(15, t.puts(['shik', '', "\n", 'hao', 123]))
+    assert_equal("shik\n\n\nhao\n123\n", t.buf)
+    assert_equal(0, t.puts([]))
+    assert_equal("shik\n\n\nhao\n123\n", t.buf)
+
+    t = hello_tube
+    assert_equal(20, t.puts(["darkhh\n\n", 'wei shi', 360]))
+    assert_equal("darkhh\n\nwei shi\n360\n", t.buf)
+
+    context.local(newline: '!!!') do
+      t = hello_tube
+      assert_equal(5, t.puts('hi'))
+      assert_equal('hi!!!', t.buf)
+    end
+  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 ::Pwnlib::Errors::EndOfTubeError
+      t.io.rewind
+      assert_equal(IO.binread(FLAG_FILE), t.io.read)
+    end
+    assert_equal("[INFO] Switching to interactive mode\n[INFO] Got EOF in 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(:@io, File.new(FLAG_FILE, File::RDONLY))
+      t.interact
+    rescue ::Pwnlib::Errors::EndOfTubeError
+      $stdout.rewind
+      assert_equal(IO.binread(FLAG_FILE), $stdout.read)
+    end
+    assert_equal("[INFO] Switching to interactive mode\n[INFO] Got EOF in 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

data/test/ui_test.rb

@@ -0,0 +1,18 @@
+# encoding: ASCII-8BIT
+# frozen_string_literal: true
+
+require 'stringio'
+
+require 'test_helper'
+
+require 'pwnlib/ui'
+
+class UITest < MiniTest::Test
+  def test_pause
+    hook_stdin(StringIO.new("\n")) do
+      assert_output(<<-EOS) { log_stdout { ::Pwnlib::UI.pause } }
+[INFO] Paused (press enter to continue)
+      EOS
+    end
+  end
+end

data/test/util/cyclic_test.rb

@@ -1,10 +1,12 @@
 # encoding: ASCII-8BIT
+# frozen_string_literal: true
 
 require 'test_helper'
+
 require 'pwnlib/util/cyclic'
 
 class CyclicTest < MiniTest::Test
-  include ::Pwnlib::Util::Cyclic::ClassMethods
+  include ::Pwnlib::Util::Cyclic
 
   def test_cyclic
     assert_equal('AAABAACABBABCACBACCBBBCBCCC', cyclic(alphabet: 'ABC', n: 3))

data/test/util/fiddling_test.rb

@@ -1,10 +1,12 @@
 # encoding: ASCII-8BIT
+# frozen_string_literal: true
 
 require 'test_helper'
+
 require 'pwnlib/util/fiddling'
 
 class FiddlingTest < MiniTest::Test
-  include ::Pwnlib::Util::Fiddling::ClassMethods
+  include ::Pwnlib::Util::Fiddling
 
   def test_enhex
     assert_equal('4141313233', enhex('AA123'))
@@ -38,11 +40,11 @@ class FiddlingTest < MiniTest::Test
     assert_equal('test A', urldecode('te%73t%20%41'))
     assert_equal("\x00\xff\x01\xfe", urldecode('%00%ff%01%fe'))
 
-    assert_equal('%qq', urldecode('%qq', true))
+    assert_equal('%qq', urldecode('%qq', ignore_invalid: true))
     err = assert_raises(ArgumentError) { urldecode('%qq') }
     assert_match(/Invalid input to urldecode/, err.message)
 
-    assert_equal('%%1z2%orz%%%%%#$!#)@%', urldecode('%%1z2%orz%%%%%#$!#)@%', true))
+    assert_equal('%%1z2%orz%%%%%#$!#)@%', urldecode('%%1z2%orz%%%%%#$!#)@%', ignore_invalid: true))
     err = assert_raises(ArgumentError) { urldecode('%ff%') }
     assert_match(/Invalid input to urldecode/, err.message)
   end
@@ -103,4 +105,11 @@ class FiddlingTest < MiniTest::Test
     assert_equal('test', b64d('dGVzdA=='))
     assert_equal("\xb2\x18\xa4" * 100, b64d('shik' * 100))
   end
+
+  def test_xor
+    assert_equal('happy', xor("\xE8\xE1\xF0\xF0\xF9", "\x80"))
+    assert_equal('happy', xor("\x80", "\xE8\xE1\xF0\xF0\xF9"))
+    assert_equal("\x04\x04\x04\x02\v\r\x11\x10\x11", xor('plaintext', 'thekey'))
+    assert_equal('2172172172', xor('217', "\x00" * 10))
+  end
 end