pwntools 0.1.0 → 1.0.0

data/lib/pwnlib/tubes/tube.rb

@@ -0,0 +1,270 @@
+# encoding: ASCII-8BIT
+
+require 'pwnlib/context'
+require 'pwnlib/logger'
+require 'pwnlib/timer'
+require 'pwnlib/tubes/buffer'
+require 'pwnlib/util/hexdump'
+
+module Pwnlib
+  module Tubes
+    # Things common to all tubes (sockets, tty, ...)
+    # @!macro [new] timeout_definition
+    #   @param [Float] timeout
+    #     Any positive floating number, indicates timeout in seconds.
+    #     Using +context.timeout+ if +timeout+ equals to +nil+.
+    class Tube
+      BUFSIZE = 4096
+
+      # Instantiate a {Pwnlib::Tubes::Tube} object.
+      #
+      # @!macro timeout_definition
+      def initialize(timeout: nil)
+        @timer = Timer.new(timeout)
+        @buffer = Buffer.new
+      end
+
+      # Receives up to +num_bytes+ bytes of data from the tube, and returns as soon as any quantity
+      # of data is available.
+      #
+      # @param [Integer] num_bytes
+      #   The maximum number of bytes to receive.
+      # @!macro timeout_definition
+      #
+      # @return [String]
+      #   A string contains bytes received from the tube, or +''+ if a timeout occurred while
+      #   waiting.
+      def recv(num_bytes = nil, timeout: nil)
+        return '' if @buffer.empty? && !fillbuffer(timeout: timeout)
+        @buffer.get(num_bytes)
+      end
+      alias read recv
+
+      # Puts the specified data back at the beginning of the receive buffer.
+      #
+      # @param [String] data
+      #   A string to put back.
+      def unrecv(data)
+        @buffer.unget(data)
+      end
+
+      # Receives one byte at a time from the tube, until the predicate evaluates to +true+.
+      #
+      # @!macro timeout_definition
+      #
+      # @return [String]
+      #   A string contains bytes received from the tube, or +''+ if a timeout occurred while
+      #   waiting.
+      #
+      # @yield
+      #   A predicate to evaluate whether the data satisfy the condition.
+      #
+      # @yieldparam [String] data
+      #   A string data to be validated by the predicate.
+      #
+      # @yieldreturn [Boolean]
+      #   Whether the data satisfy the condition.
+      #
+      # @raise [ArgumentError]
+      #   If the block is not given.
+      def recvpred(timeout: nil)
+        raise ArgumentError, 'Need a block for recvpred' unless block_given?
+        @timer.countdown(timeout) do
+          data = ''
+          begin
+            until yield(data)
+              return '' unless @timer.active?
+
+              begin
+                # TODO(Darkpi): Some form of binary search to speed up?
+                c = recv(1)
+              rescue
+                return ''
+              end
+
+              return '' if c.empty?
+              data << c
+            end
+            data.slice!(0..-1)
+          ensure
+            unrecv(data)
+          end
+        end
+      end
+
+      # Receives exactly +num_bytes+ bytes.
+      # If the request is not satisfied before +timeout+ seconds pass, all data is buffered and an
+      # empty string +''+ is returned.
+      #
+      # @param [Integer] num_bytes
+      #   The number of bytes to receive.
+      # @!macro timeout_definition
+      #
+      # @return [String]
+      #   A string contains bytes received from the tube, or +''+ if a timeout occurred while
+      #   waiting.
+      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
+
+      # Receive 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 timeout_definition
+      #
+      # @return [String]
+      #   A string contains bytes, which ends string in +delims+, received from the tube.
+      #
+      # @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.
+      def recvuntil(delims, drop: false, timeout: nil)
+        delims = Array(delims)
+        max_len = delims.map(&:size).max
+        @timer.countdown(timeout) do
+          data = Buffer.new
+          matching = ''
+          begin
+            while @timer.active?
+              begin
+                s = recv(1)
+              rescue # TODO(Darkpi): QQ
+                return ''
+              end
+
+              return '' if s.empty?
+              matching << s
+
+              sidx = matching.size
+              match_len = 0
+              delims.each do |d|
+                idx = matching.index(d)
+                next unless idx
+                if idx + d.size <= sidx + match_len
+                  sidx = idx
+                  match_len = d.size
+                end
+              end
+
+              if sidx < matching.size
+                r = data.get + matching.slice!(0, sidx + match_len)
+                r.slice!(-match_len..-1) if drop
+                return r
+              end
+
+              data << matching.slice!(0...-max_len) if matching.size > max_len
+            end
+            ''
+          ensure
+            unrecv(matching)
+            unrecv(data)
+          end
+        end
+      end
+
+      # Receive 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 timeout_definition
+      #
+      # @return [String]
+      #   All bytes received over the tube until the first newline is received.
+      #   Optionally retains the ending.
+      def recvline(drop: false, timeout: nil)
+        recvuntil(context.newline, drop: drop, timeout: timeout)
+      end
+      alias gets recvline
+
+      # Wrapper around +recvpred+, which will return when a regex matches the string in the buffer.
+      #
+      # @param [Regexp] regex
+      #   A regex to match.
+      # @!macro timeout_definition
+      #
+      # @return [String]
+      #   A string contains bytes received from the tube, or +''+ if a timeout occurred while
+      #   waiting.
+      def recvregex(regex, timeout: nil)
+        recvpred(timeout: timeout) { |data| data =~ regex }
+      end
+
+      # Sends data
+      #
+      # @param [String] data
+      #   The +data+ string to send.
+      def send(data)
+        data = data.to_s
+        log.debug(format('Sent %#x bytes:', data.size))
+        log.indented(hexdump(data), level: DEBUG)
+        send_raw(data)
+      end
+      alias write send
+
+      # Sends data 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)
+      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])
+          if rs.include?($stdin)
+            s = $stdin.readpartial(BUFSIZE)
+            write(s)
+          end
+          if rs.include?(io)
+            s = recv
+            $stdout.write(s)
+          end
+        end
+      end
+
+      private
+
+      def fillbuffer(timeout: nil)
+        data = @timer.countdown(timeout) do
+          self.timeout_raw = @timer.timeout
+          recv_raw(BUFSIZE)
+        end
+        if data
+          @buffer << data
+          log.debug(format('Received %#x bytes:', data.size))
+          log.indented(hexdump(data), level: DEBUG)
+        end
+        data
+      end
+
+      def send_raw(_data); raise NotImplementedError, 'Not implemented'
+      end
+
+      def recv_raw(_size); raise NotImplementedError, 'Not implemented'
+      end
+
+      def timeout_raw=(_timeout); raise NotImplementedError, 'Not implemented'
+      end
+
+      include ::Pwnlib::Context
+      include ::Pwnlib::Logger
+      include ::Pwnlib::Util::HexDump
+    end
+  end
+end

data/lib/pwnlib/util/cyclic.rb

@@ -3,7 +3,7 @@
 module Pwnlib
   module Util
     # Generate string with easy-to-find pattern.
-    # See {ClassMethods} for method details.
+    #
     # @example Call by specifying full module path.
     #   require 'pwnlib/util/cyclic'
     #   Pwnlib::Util::Cyclic.cyclic_find(Pwnlib::Util::Cyclic.cyclic(200)[123, 4]) #=> 123
@@ -11,110 +11,111 @@ module Pwnlib
     #   require 'pwn'
     #   cyclic_find(cyclic(200)[123, 4]) #=> 123
     module Cyclic
-      # @note Do not create and call instance method here. Instead, call module method on {Cyclic}.
-      module ClassMethods
-        # TODO(Darkpi): Should we put this constant in some 'String' module?
-        ASCII_LOWERCASE = ('a'..'z').to_a.join
-        private_constant :ASCII_LOWERCASE
+      # TODO(Darkpi): Should we put this constant in some 'String' module?
+      ASCII_LOWERCASE = ('a'..'z').to_a.join
+      private_constant :ASCII_LOWERCASE
 
-        # Generator for a sequence of unique substrings of length +n+.
-        # This is implemented using a De Bruijn Sequence over the given +alphabet+.
-        # Returns an Enumerator if no block given.
-        #
-        # @overload de_bruijn(alphabet: ASCII_LOWERCASE, n: 4)
-        #   @param [String, Array] alphabet
-        #     Alphabet to be used.
-        #   @param [Integer] n
-        #     Length of substring that should be unique.
-        #   @return [void]
-        #   @yieldparam c
-        #     Item of the result sequence in order.
-        # @overload de_bruijn(alphabet: ASCII_LOWERCASE, n: 4)
-        #   @param [String, Array] alphabet
-        #     Alphabet to be used.
-        #   @param [Integer] n
-        #     Length of substring that should be unique.
-        #   @return [Enumerator]
-        #     The result sequence.
-        def de_bruijn(alphabet: ASCII_LOWERCASE, n: 4)
-          return to_enum(__method__, alphabet: alphabet, n: n) { alphabet.size**n } unless block_given?
-          k = alphabet.size
-          a = [0] * (k * n)
+      module_function
 
-          db = lambda do |t, p|
-            if t > n
-              (1..p).each { |j| yield alphabet[a[j]] } if (n % p).zero?
-            else
-              a[t] = a[t - p]
-              db.call(t + 1, p)
-              (a[t - p] + 1...k).each do |j|
-                a[t] = j
-                db.call(t + 1, t)
-              end
+      # Generator for a sequence of unique substrings of length +n+.
+      # This is implemented using a De Bruijn Sequence over the given +alphabet+.
+      # Returns an Enumerator if no block given.
+      #
+      # @overload de_bruijn(alphabet: ASCII_LOWERCASE, n: 4)
+      #   @param [String, Array] alphabet
+      #     Alphabet to be used.
+      #   @param [Integer] n
+      #     Length of substring that should be unique.
+      #
+      #   @return [void]
+      #
+      #   @yieldparam c
+      #     Item of the result sequence in order.
+      #
+      # @overload de_bruijn(alphabet: ASCII_LOWERCASE, n: 4)
+      #   @param [String, Array] alphabet
+      #     Alphabet to be used.
+      #   @param [Integer] n
+      #     Length of substring that should be unique.
+      #
+      #   @return [Enumerator]
+      #     The result sequence.
+      def de_bruijn(alphabet: ASCII_LOWERCASE, n: 4)
+        return to_enum(__method__, alphabet: alphabet, n: n) { alphabet.size**n } unless block_given?
+        k = alphabet.size
+        a = [0] * (k * n)
+
+        db = lambda do |t, p|
+          if t > n
+            (1..p).each { |j| yield alphabet[a[j]] } if (n % p).zero?
+          else
+            a[t] = a[t - p]
+            db.call(t + 1, p)
+            (a[t - p] + 1...k).each do |j|
+              a[t] = j
+              db.call(t + 1, t)
             end
           end
-
-          db[1, 1]
         end
 
-        # Simple wrapper over {#de_bruijn}, returning at most +length+ items.
-        #
-        # @param [Integer, nil] length
-        #   Desired length of the sequence,
-        #   or +nil+ for the entire sequence.
-        # @param [String, Array] alphabet
-        #   Alphabet to be used.
-        # @param [Integer] n
-        #   Length of substring that should be unique.
-        # @return [String, Array]
-        #   The result sequence of at most +length+ items,
-        #   with same type as +alphabet+.
-        #
-        # @example
-        #   cyclic(alphabet: 'ABC', n: 3) #=> 'AAABAACABBABCACBACCBBBCBCCC'
-        #   cyclic(20) #=> 'aaaabaaacaaadaaaeaaa'
-        def cyclic(length = nil, alphabet: ASCII_LOWERCASE, n: 4)
-          enum = de_bruijn(alphabet: alphabet, n: n)
-          r = length.nil? ? enum.to_a : enum.take(length)
-          alphabet.is_a?(String) ? r.join : r
-        end
+        db[1, 1]
+      end
 
-        # Find the position of a substring in a De Bruijn sequence
-        #
-        # @todo Speed! See comment in Python pwntools.
-        # @param [String, Array] subseq
-        #   The substring to be found in the sequence.
-        # @param [String, Array] alphabet
-        #   Alphabet to be used.
-        # @param [Integer] n
-        #   Length of substring that should be unique.
-        #   Default to +subseq.size+.
-        # @return [Integer, nil]
-        #   The index +subseq+ first appear in the sequence,
-        #   or +nil+ if not found.
-        #
-        # @example
-        #   cyclic_find(cyclic(300)[217, 4]) #=> 217
-        def cyclic_find(subseq, alphabet: ASCII_LOWERCASE, n: nil)
-          n ||= subseq.size
-          subseq = subseq.chars if subseq.is_a?(String)
-          return nil unless subseq.all? { |c| alphabet.include?(c) }
+      # Simple wrapper over {.de_bruijn}, returning at most +length+ items.
+      #
+      # @param [Integer, nil] length
+      #   Desired length of the sequence, or +nil+ for the entire sequence.
+      # @param [String, Array] alphabet
+      #   Alphabet to be used.
+      # @param [Integer] n
+      #   Length of substring that should be unique.
+      #
+      # @return [String, Array]
+      #   The result sequence of at most +length+ items, with same type as +alphabet+.
+      #
+      # @example
+      #   cyclic(alphabet: 'ABC', n: 3) #=> 'AAABAACABBABCACBACCBBBCBCCC'
+      #   cyclic(20) #=> 'aaaabaaacaaadaaaeaaa'
+      def cyclic(length = nil, alphabet: ASCII_LOWERCASE, n: 4)
+        enum = de_bruijn(alphabet: alphabet, n: n)
+        r = length.nil? ? enum.to_a : enum.take(length)
+        alphabet.is_a?(String) ? r.join : r
+      end
 
-          pos = 0
-          saved = []
-          de_bruijn(alphabet: alphabet, n: n).each do |c|
-            saved << c
-            if saved.size > subseq.size
-              saved.shift
-              pos += 1
-            end
-            return pos if saved == subseq
+      # Find the position of a substring in a De Bruijn sequence.
+      #
+      # @param [String, Array] subseq
+      #   The substring to be found in the sequence.
+      # @param [String, Array] alphabet
+      #   Alphabet to be used.
+      # @param [Integer] n
+      #   Length of substring that should be unique.
+      #   Default to +subseq.size+.
+      #
+      # @return [Integer, nil]
+      #   The index +subseq+ first appear in the sequence, or +nil+ if not found.
+      #
+      # @todo Speed! See comment in Python pwntools.
+      #
+      # @example
+      #   cyclic_find(cyclic(300)[217, 4]) #=> 217
+      def cyclic_find(subseq, alphabet: ASCII_LOWERCASE, n: nil)
+        n ||= subseq.size
+        subseq = subseq.chars if subseq.is_a?(String)
+        return nil unless subseq.all? { |c| alphabet.include?(c) }
+
+        pos = 0
+        saved = []
+        de_bruijn(alphabet: alphabet, n: n).each do |c|
+          saved << c
+          if saved.size > subseq.size
+            saved.shift
+            pos += 1
           end
-          nil
+          return pos if saved == subseq
         end
+        nil
       end
-
-      extend ClassMethods
     end
   end
 end