Enhance Tubes

lib/pwnlib/tubes/tube.rb

@@ -9,10 +9,18 @@
 module Pwnlib
   module Tubes
     # Things common to all tubes (sockets, tty, ...)
+    # @!macro [new] drop_definition
+    #   @param [Boolean] drop
+    #     Whether drop the ending.
+    #
     # @!macro [new] timeout_definition
     #   @param [Float] timeout
     #     Any positive floating number, indicates timeout in seconds.
     #     Using +context.timeout+ if +timeout+ equals to +nil+.
+    #
+    # @!macro [new] send_return_definition
+    #   @return [Integer]
+    #     Returns the number of bytes had been sent.
     class Tube
       BUFSIZE = 4096
 
@@ -44,8 +52,12 @@ def recv(num_bytes = nil, timeout: nil)
       #
       # @param [String] data
       #   A string to put back.
+      #
+      # @return [Integer]
+      #   The length of the put back data.
       def unrecv(data)
         @buffer.unget(data)
+        data.size
       end
 
       # Receives one byte at a time from the tube, until the predicate evaluates to +true+.
@@ -110,13 +122,12 @@ def recvn(num_bytes, timeout: nil)
         end
       end
 
-      # Receive data until one of +delims+ is encountered. If the request is not satisfied before
+      # Receives data until one of +delims+ is encountered. If the request is not satisfied before
       # +timeout+ seconds pass, all data is buffered and an empty string is returned.
       #
       # @param [Array<String>] delims
       #   String of delimiters characters, or list of delimiter strings.
-      # @param [Boalean] drop
-      #   Whether drop the ending.
+      # @!macro drop_definition
       # @!macro timeout_definition
       #
       # @return [String]
@@ -169,12 +180,11 @@ def recvuntil(delims, drop: false, timeout: nil)
         end
       end
 
-      # Receive a single line from the tube.
+      # Receives a single line from the tube.
       # A "line" is any sequence of bytes terminated by the byte sequence set in +context.newline+,
       # which defaults to +"\n"+.
       #
-      # @param [Boolean] drop
-      #   Whether drop the line ending.
+      # @!macro drop_definition
       # @!macro timeout_definition
       #
       # @return [String]
@@ -183,7 +193,48 @@ def recvuntil(delims, drop: false, timeout: nil)
       def recvline(drop: false, timeout: nil)
         recvuntil(context.newline, drop: drop, timeout: timeout)
       end
-      alias gets recvline
+
+      # Receives the next "line" from the tube; lines are separated by +sep+.
+      # The difference with +IO#gets+ is using +context.newline+ as default newline.
+      #
+      # @param [String, Integer] sep
+      #   If +String+ is given, use +sep+ as the separator.
+      #   If +Integer+ is given, receive exactly +sep+ bytes.
+      # @!macro drop_definition
+      # @!macro timeout_definition
+      #
+      # @return [String]
+      #   The next "line".
+      #
+      # @raise [EOFError]
+      #   When the remaining data does not contain +sep+.
+      #   When the size of the remaining data is less than +sep+.
+      #
+      # @example
+      #   Sock.new('127.0.0.1', 1337).gets
+      #   #=> "This is line one\n"
+      #
+      #   Sock.new('127.0.0.1', 1337).gets(drop: true)
+      #   #=> "This is line one"
+      #
+      #   Sock.new('127.0.0.1', 1337).gets 'line'
+      #   #=> "This is line"
+      #
+      #   Sock.new('127.0.0.1', 1337).gets ''
+      #   #=> "This is line"
+      #
+      #   Sock.new('127.0.0.1', 1337).gets(4)
+      #   #=> "This"
+      def gets(sep = context.newline, drop: false, timeout: nil)
+        case sep
+        when Integer
+          recvn(sep, timeout: timeout)
+        when String
+          recvuntil(sep, drop: drop, timeout: timeout)
+        else
+          raise ArgumentError, 'only Integer and String are supported'
+        end
+      end
 
       # Wrapper around +recvpred+, which will return when a regex matches the string in the buffer.
       #
@@ -198,27 +249,69 @@ def recvregex(regex, timeout: nil)
         recvpred(timeout: timeout) { |data| data =~ regex }
       end
 
-      # Sends data
+      # Sends data.
       #
       # @param [String] data
       #   The +data+ string to send.
+      #
+      # @!macro send_return_definition
       def send(data)
         data = data.to_s
         log.debug(format('Sent %#x bytes:', data.size))
         log.indented(hexdump(data), level: DEBUG)
         send_raw(data)
+        data.size
       end
       alias write send
 
-      # Sends data with +context.newline+.
+      # Sends the given object with +context.newline+.
       #
-      # @param [String] data
-      #   The +data+ string to send.
-      def sendline(data)
-        # Logged by +write+, not +send_raw+
-        write(data.to_s + context.newline)
+      # @param [Object] obj
+      #   The object to send.
+      #
+      # @!macro send_return_definition
+      def sendline(obj)
+        s = obj.to_s + context.newline
+        write(s)
+      end
+
+      # Sends the given object(s) to the tube.
+      # The difference with +IO#puts+ is using +context.newline+ as default newline.
+      #
+      # @param [Array<Object>] objs
+      #   The objects to send.
+      #
+      # @!macro send_return_definition
+      #
+      # @example
+      #   s.puts
+      #   puts client.recv
+      #   #
+      #
+      # @example
+      #   s.puts('shik', "hao\n", 123)
+      #   puts client.recv
+      #   # shik
+      #   # hao
+      #   # 123
+      #
+      # @example
+      #   s.puts(["darkhh\n\n", 'wei shi', 360])
+      #   puts client.recv
+      #   # darkhh
+      #   #
+      #   # wei shi
+      #   # 360
+      def puts(*objs)
+        return write(context.newline) if objs.empty?
+        objs = *objs.flatten
+        s = ''
+        objs.map(&:to_s).each do |elem|
+          s << elem
+          s << context.newline unless elem.end_with?(context.newline)
+        end
+        write(s)
       end
-      alias puts sendline
 
       # Does simultaneous reading and writing to the tube. In principle this just connects the tube
       # to standard in and standard out.
@@ -236,6 +329,9 @@ def interact
             $stdout.write(s)
           end
         end
+      # TODO(darkhh): Use our own Exception class.
+      rescue EOFError
+        log.info('Got EOF in interactive mode')
       end
 
       private

test/tubes/tube_test.rb

@@ -136,6 +136,20 @@ def test_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/
@@ -157,11 +171,11 @@ def test_recvregex
 
   def test_send
     t = hello_tube
-    t.write('DARKHH')
+    assert_equal(6, t.write('DARKHH'))
     assert_equal('DARKHH', t.buf)
-    t.write(' QQ')
+    assert_equal(3, t.write(' QQ'))
     assert_equal('DARKHH QQ', t.buf)
-    t.write(333)
+    assert_equal(3, t.write(333))
     assert_equal('DARKHH QQ333', t.buf)
 
     context.local(log_level: 'debug') do
@@ -192,8 +206,38 @@ def test_sendline
     t = hello_tube
     t.write('DARKHH')
     assert_equal('DARKHH', t.buf)
-    t.puts(' QQ')
+    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__)
@@ -208,7 +252,7 @@ def test_interact_send
       t.io.rewind
       assert_equal(IO.binread(FLAG_FILE), t.io.read)
     end
-    assert_equal("[INFO] Switching to interactive mode\n", @log.string)
+    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
@@ -228,7 +272,7 @@ def test_interact_recv
       $stdout.rewind
       assert_equal(IO.binread(FLAG_FILE), $stdout.read)
     end
-    assert_equal("[INFO] Switching to interactive mode\n", @log.string)
+    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