pwntools 1.0.1 → 1.1.0

data/lib/pwnlib/tubes/sock.rb

@@ -2,30 +2,30 @@
 
 require 'socket'
 
+require 'pwnlib/errors'
 require 'pwnlib/tubes/tube'
 
 module Pwnlib
   module Tubes
     # Socket!
     class Sock < Tube
+      attr_reader :sock # @return [TCPSocket] The socket object.
+
       # Instantiate a {Pwnlib::Tubes::Sock} object.
       #
       # @param [String] host
       #   The host to connect.
       # @param [Integer] port
       #   The port to connect.
-      def initialize(host, port)
-        super()
+      # @param [Float?] timeout
+      #   See {Pwnlib::Tubes::Tube#initialize}.
+      def initialize(host, port, timeout: nil)
+        super(timeout: timeout)
         @sock = TCPSocket.new(host, port)
         @sock.binmode
         @timeout = nil
-        @closed = { recv: false, send: false }
-      end
-
-      def io
-        @sock
+        @closed = { read: false, write: false }
       end
-      alias sock io
 
       # Close the TCPSocket if no arguments passed.
       # Or close the direction in +sock+.
@@ -35,59 +35,58 @@ module Pwnlib
       #   * Disallow further read in +sock+ if +:recv+ or +:read+ passed.
       #   * Disallow further write in +sock+ if +:send+ or +:write+ passed.
       #
+      # @return [void]
+      #
       # @diff In pwntools-python, method +shutdown(direction)+ is for closing socket one side,
       #   +close()+ is for closing both side. We merge these two methods into one here.
       def close(direction = :both)
-        case direction
-        when :both
+        if direction == :both
           return if @sock.closed?
-          @closed[:recv] = @closed[:send] = true
+          @closed[:read] = @closed[:write] = true
           @sock.close
-        when :recv, :read
-          shutdown(:recv)
-        when :send, :write
-          shutdown(:send)
         else
-          raise ArgumentError, 'Only allow :both, :recv, :read, :send and :write passed'
+          shutdown(*normalize_direction(direction))
         end
       end
 
       private
 
+      alias io_out sock
+
       def shutdown(direction)
         return if @closed[direction]
         @closed[direction] = true
 
-        if direction.equal?(:recv)
+        if direction.equal?(:read)
           @sock.close_read
-        elsif direction.equal?(:send)
+        elsif direction.equal?(:write)
           @sock.close_write
         end
       end
 
       def timeout_raw=(timeout)
-        @timeout = timeout == :forever ? nil : timeout
+        @timeout = timeout
       end
 
       def send_raw(data)
-        raise EOFError if @closed[:send]
+        raise ::Pwnlib::Errors::EndOfTubeError if @closed[:write]
         begin
           @sock.write(data)
         rescue Errno::EPIPE, Errno::ECONNRESET, Errno::ECONNREFUSED
-          shutdown(:send)
-          raise EOFError
+          shutdown(:write)
+          raise ::Pwnlib::Errors::EndOfTubeError
         end
       end
 
       def recv_raw(size)
-        raise EOFError if @closed[:recv]
+        raise ::Pwnlib::Errors::EndOfTubeError if @closed[:read]
         begin
           rs, = IO.select([@sock], [], [], @timeout)
           return if rs.nil?
           return @sock.readpartial(size)
         rescue Errno::ECONNREFUSED, Errno::ECONNRESET, Errno::ECONNABORTED, EOFError
-          shutdown(:recv)
-          raise EOFError
+          shutdown(:read)
+          raise ::Pwnlib::Errors::EndOfTubeError
         end
       end
     end

data/lib/pwnlib/tubes/tube.rb

@@ -1,19 +1,39 @@
 # encoding: ASCII-8BIT
 
 require 'pwnlib/context'
+require 'pwnlib/errors'
 require 'pwnlib/logger'
 require 'pwnlib/timer'
 require 'pwnlib/tubes/buffer'
 require 'pwnlib/util/hexdump'
 
 module Pwnlib
+  # Module that contains all kinds of tubes.
   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.
+    #
+    # @!macro [new] raise_eof
+    #   @raise [Pwnlib::Errors::EndOfTubeError]
+    #     If the request is not satisfied when all data is received.
+    #
+    # @!macro [new] raise_timeout
+    #   @raise [Pwnlib::Errors::TimeoutError]
+    #     If the request is not satisfied when timeout exceeded.
+
+    # Things common to all tubes (sockets, tty, ...)
     class Tube
+      # Receive 4096 bytes each time.
       BUFSIZE = 4096
 
       # Instantiate a {Pwnlib::Tubes::Tube} object.
@@ -34,6 +54,9 @@ module Pwnlib
       # @return [String]
       #   A string contains bytes received from the tube, or +''+ if a timeout occurred while
       #   waiting.
+      #
+      # @!macro raise_eof
+      # @!macro raise_timeout
       def recv(num_bytes = nil, timeout: nil)
         return '' if @buffer.empty? && !fillbuffer(timeout: timeout)
         @buffer.get(num_bytes)
@@ -44,8 +67,12 @@ module Pwnlib
       #
       # @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+.
@@ -67,6 +94,9 @@ module Pwnlib
       #
       # @raise [ArgumentError]
       #   If the block is not given.
+      #
+      # @!macro raise_eof
+      # @!macro raise_timeout
       def recvpred(timeout: nil)
         raise ArgumentError, 'Need a block for recvpred' unless block_given?
         @timer.countdown(timeout) do
@@ -75,12 +105,7 @@ module Pwnlib
             until yield(data)
               return '' unless @timer.active?
 
-              begin
-                # TODO(Darkpi): Some form of binary search to speed up?
-                c = recv(1)
-              rescue
-                return ''
-              end
+              c = recv(1)
 
               return '' if c.empty?
               data << c
@@ -103,25 +128,31 @@ module Pwnlib
       # @return [String]
       #   A string contains bytes received from the tube, or +''+ if a timeout occurred while
       #   waiting.
+      #
+      # @!macro raise_eof
+      # @!macro raise_timeout
       def recvn(num_bytes, timeout: nil)
         @timer.countdown(timeout) do
           fillbuffer while @timer.active? && @buffer.size < num_bytes
           @buffer.size >= num_bytes ? @buffer.get(num_bytes) : ''
         end
       end
+      alias readn recvn
 
-      # 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]
       #   A string contains bytes, which ends string in +delims+, received from the tube.
       #
+      # @!macro raise_eof
+      # @!macro raise_timeout
+      #
       # @diff We return the string that ends the earliest, rather then starts the earliest,
       #       since the latter can't be done greedly. Still, it would be bad to call this
       #       for case with ambiguity.
@@ -133,11 +164,7 @@ module Pwnlib
           matching = ''
           begin
             while @timer.active?
-              begin
-                s = recv(1)
-              rescue # TODO(Darkpi): QQ
-                return ''
-              end
+              s = recv(1)
 
               return '' if s.empty?
               matching << s
@@ -169,12 +196,11 @@ module Pwnlib
         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"+.
+      # which defaults to +"\\n"+.
       #
-      # @param [Boolean] drop
-      #   Whether drop the line ending.
+      # @!macro drop_definition
       # @!macro timeout_definition
       #
       # @return [String]
@@ -183,7 +209,48 @@ module Pwnlib
       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 [Pwnlib::Errors::EndOfTubeError]
+      #   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,73 +265,158 @@ module Pwnlib
         recvpred(timeout: timeout) { |data| data =~ regex }
       end
 
-      # Sends data
+      # Sends data.
       #
       # @param [String] data
-      #   The +data+ string to send.
+      #   The +data+ string to be sent.
+      #
+      # @!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)
+        log.indented(::Pwnlib::Util::HexDump.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 be sent.
+      #
+      # @!macro send_return_definition
+      def sendline(obj)
+        s = obj.to_s + context.newline
+        write(s)
+      end
+
+      # Sends the given object(s).
+      # The difference with +IO#puts+ is using +context.newline+ as default newline.
+      #
+      # @param [Array<Object>] objs
+      #   The objects to be sent.
+      #
+      # @!macro send_return_definition
+      #
+      # @example
+      #   s.puts
+      #   puts client.recv
+      #   #
+      #   #=> nil
+      #
+      # @example
+      #   s.puts('shik', "hao\n", 123)
+      #   puts client.recv
+      #   # shik
+      #   # hao
+      #   # 123
+      #   #=> nil
+      #
+      # @example
+      #   s.puts(["darkhh\n\n", 'wei shi', 360])
+      #   puts client.recv
+      #   # darkhh
+      #   #
+      #   # wei shi
+      #   # 360
+      #   #=> nil
+      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.
       def interact
         log.info('Switching to interactive mode')
         $stdout.write(@buffer.get)
-        until io.closed?
-          rs, = IO.select([$stdin, io])
+        until io_out.closed?
+          rs, = IO.select([$stdin, io_out])
           if rs.include?($stdin)
             s = $stdin.readpartial(BUFSIZE)
             write(s)
           end
-          if rs.include?(io)
+          if rs.include?(io_out)
             s = recv
             $stdout.write(s)
           end
         end
+      rescue ::Pwnlib::Errors::EndOfTubeError
+        log.info('Got EOF in interactive mode')
       end
 
       private
 
+      # Normalize direction.
+      #
+      # @param [Symbol] direction
+      #
+      # @return [Array<Symbol>]
+      #   If +direction+ equals to
+      #   * +:both+, returns +[:read, :write]+
+      #   * +:read+ or +:recv+, returns [:read]
+      #   * +:write+ or +:send+, returns [:write]
+      #   Otherwise, raise +ArgumentError+.
+      def normalize_direction(direction)
+        case direction
+        when :both then %i[read write]
+        when :read, :recv then [:read]
+        when :write, :send then [:write]
+        else
+          raise ArgumentError, 'Only allow :both, :recv, :read, :send and :write passed'
+        end
+      end
+
       def fillbuffer(timeout: nil)
         data = @timer.countdown(timeout) do
-          self.timeout_raw = @timer.timeout
+          self.timeout_raw = (@timer.timeout == :forever ? nil : @timer.timeout)
           recv_raw(BUFSIZE)
         end
         if data
           @buffer << data
           log.debug(format('Received %#x bytes:', data.size))
-          log.indented(hexdump(data), level: DEBUG)
+          log.indented(::Pwnlib::Util::HexDump.hexdump(data), level: DEBUG)
         end
         data
       end
 
-      def send_raw(_data); raise NotImplementedError, 'Not implemented'
+      # The IO object of output, will be used for IO.select([io_out]) in interactive mode.
+      #
+      # @return [IO]
+      def io_out
+        raise NotImplementedError, 'Not implemented'
       end
 
-      def recv_raw(_size); raise NotImplementedError, 'Not implemented'
+      # @param [String] _data
+      #
+      # @return [void]
+      def send_raw(_data)
+        raise NotImplementedError, 'Not implemented'
       end
 
-      def timeout_raw=(_timeout); raise NotImplementedError, 'Not implemented'
+      # @param [Integer] _size
+      #
+      # @return [String]
+      def recv_raw(_size)
+        raise NotImplementedError, 'Not implemented'
+      end
+
+      # @param [Float?] _timeout
+      #
+      # @return [void]
+      def timeout_raw=(_timeout)
+        raise NotImplementedError, 'Not implemented'
       end
 
       include ::Pwnlib::Context
       include ::Pwnlib::Logger
-      include ::Pwnlib::Util::HexDump
     end
   end
 end