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 [Boalean] 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,20 @@ 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] sep
+      #   The separator.
+      # @!macro drop_definition
+      # @!macro timeout_definition
+      #
+      # @return [String]
+      #   The next "line".
+      def gets(sep = context.newline, drop: false, timeout: nil)
+        recvuntil(sep, drop: drop, timeout: timeout)
+      end
 
       # Wrapper around +recvpred+, which will return when a regex matches the string in the buffer.
       #
@@ -198,27 +221,49 @@ 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
+      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 if elem.empty? || !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 +281,8 @@ def interact
             $stdout.write(s)
           end
         end
+      rescue
+        log.info('Got EOF in interactive mode')
       end
 
       private

test/tubes/tube_test.rb

@@ -136,6 +136,18 @@ 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))
+  end
+
   def test_recvpred
     t = hello_tube
     r = /H.*w/
@@ -157,11 +169,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 +204,34 @@ 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', '', 'wei'))
+    assert_equal("shik\nhao\n\nwei\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)
+
+    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 +246,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 +266,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