kronk 1.4.0 → 1.5.0

data/lib/kronk/response.rb

@@ -1,254 +1,435 @@
 class Kronk
 
   ##
-  # Wrapper to add a few niceties to the Net::HTTPResponse class.
+  # Mock File IO to allow rewinding on Windows platforms.
 
-  class Response < Net::HTTPResponse
+  class WinFileIO < StringIO
+    attr_accessor :path
+
+    def initialize path, str=""
+      @path = path
+      super str
+    end
+  end
+
+  ##
+  # Wraps an http response.
+
+  class Response
 
     class MissingParser < Exception; end
 
+    ENCODING_MATCHER = /(^|;\s?)charset=(.*?)\s*(;|$)/
+
     ##
-    # Create a new Response instance from an IO object.
+    # Read http response from a file and return a HTTPResponse instance.
 
-    def self.read_new io
-      io = Net::BufferedIO.new io unless Net::BufferedIO === io
-      io.debug_output = socket_io = StringIO.new
+    def self.read_file path
+      Kronk::Cmd.verbose "Reading file:  #{path}\n"
 
-      begin
-        resp = super io
-        resp.reading_body io, true do;end
-      rescue EOFError
+      file     = File.open(path, "rb")
+      resp     = new file
+      resp.uri = path
+      file.close
+
+      resp
+    end
+
+
+    attr_accessor :body, :bytes, :byterate, :code, :headers, :parser,
+                  :raw, :request, :uri
+
+    attr_reader :encoding, :time
+
+    alias to_hash headers
+    alias to_s raw
+
+    ##
+    # Create a new Response object from a String or IO.
+
+    def initialize io=nil, res=nil, request=nil
+      return unless io
+      io = StringIO.new io if String === io
+
+      if io && res
+        @_res, debug_io = res, io
+      else
+        @_res, debug_io = request_from_io(io)
       end
 
-      resp.extend Helpers
-      resp.set_helper_attribs socket_io, true, true
+      @headers  = @_res.to_hash
 
-      resp
+      @encoding = "utf-8" unless @_res["Content-Type"]
+      c_type = [*@headers["content-type"]].find{|ct| ct =~ ENCODING_MATCHER}
+      @encoding = $2 if c_type
+      @encoding ||= "ASCII-8BIT"
+      @encoding = Encoding.find(@encoding) if defined?(Encoding)
+
+      raw_req, raw_resp, bytes = read_raw_from debug_io
+      @raw      = try_force_encoding raw_resp
+
+      @request  = request ||
+                  raw_req = try_force_encoding(raw_req) &&
+                  Request.parse(raw_req)
+
+      @time   = 0
+
+      @body   = try_force_encoding(@_res.body) if @_res.body
+      @body ||= @raw.split("\r\n\r\n",2)[1]
+
+      @bytes = (@_res['Content-Length'] || @body.bytes.count).to_i
+
+      @code = @_res.code
+
+      @parser = Kronk.parser_for @_res['Content-Type']
+
+      @uri = @request.uri if @request && @request.uri
+
+      @byterate = 0
     end
 
 
+    ##
+    # Accessor for the HTTPResponse instance []
+
+    def [] key
+      @_res[key]
+    end
+
 
     ##
-    # Helper methods for Net::HTTPResponse objects.
+    # Accessor for the HTTPResponse instance []
 
-    module Helpers
+    def []= key, value
+      @_res[key] = value
+    end
 
-      ##
-      # Read the raw response from a debug_output instance and return an array
-      # containing the raw request, response, and number of bytes received.
 
-      def read_raw_from debug_io
-        req = nil
-        resp = ""
-        bytes = nil
+    ##
+    # Cookie header accessor.
 
-        debug_io.rewind
-        output = debug_io.read.split "\n"
+    def cookie
+      @_res['Cookie']
+    end
 
-        if output.first =~ %r{<-\s(.*)}
-          req = instance_eval $1
-          output.delete_at 0
-        end
 
-        if output.last =~ %r{read (\d+) bytes}
-          bytes = $1.to_i
-          output.delete_at(-1)
-        end
+    ##
+    # Force the encoding of the raw response and body
+
+    def force_encoding new_encoding
+      new_encoding = Encoding.find new_encoding unless Encoding === new_encoding
+      @encoding = new_encoding
+      try_force_encoding @body
+      try_force_encoding @raw
+      @encoding
+    end
 
-        output.map do |line|
-          next unless line[0..2] == "-> "
-          resp << instance_eval(line[2..-1])
-        end
 
-        [req, resp, bytes]
-      end
+    ##
+    # If there was an error parsing the input as a standard http response,
+    # the input is assumed to be a body and HeadlessResponse is used.
+
+    def headless?
+      HeadlessResponse === @_res
+    end
+
 
+    ##
+    # Ruby inspect.
+
+    def inspect
+      "#<#{self.class}:#{self.code} #{self['Content-Type']}>"
+    end
 
-      ##
-      # Instantiates helper attributes from the debug socket io.
 
-      def set_helper_attribs socket_io, socket=nil, body_read=nil
-        @raw      = udpate_encoding read_raw_from(socket_io)[1]
+    ##
+    # Returns the body data parsed according to the content type.
+    # If no parser is given will look for the default parser based on
+    # the Content-Type, or will return the cached parsed body if available.
 
-        @read     = body_read unless body_read.nil?
-        @socket ||= socket
-        @body   ||= @raw.split("\r\n\r\n",2)[1]
+    def parsed_body parser=nil
+      @parsed_body ||= nil
 
-        udpate_encoding @body
+      return @parsed_body if @parsed_body && !parser
 
-        puts "#{@raw.length} - #{@body.length}" if
-          @raw.length == 0 && @body.length != 0
-        self
+      if String === parser
+        parser = Kronk.parser_for(parser) || Kronk.find_const(parser)
       end
 
+      parser ||= @parser
 
-      ##
-      # Returns the encoding provided in the Content-Type header or
-      # "binary" if charset is unavailable.
-      # Returns "utf-8" if no content type header is missing.
+      raise MissingParser,
+        "No parser for Content-Type: #{@_res['Content-Type']}" unless parser
 
-      def encoding
-        content_types = self.to_hash["content-type"]
+      @parsed_body = parser.parse self.body
+    end
+
+
+    ##
+    # Returns the parsed header hash.
 
-        return "utf-8" if !content_types
+    def parsed_header include_headers=true
+      headers = @_res.to_hash.dup
 
-        content_types.each do |c_type|
-          return $2 if c_type =~ /(^|;\s?)charset=(.*?)\s*(;|$)/
+      case include_headers
+      when nil, false
+        nil
+
+      when Array, String
+        include_headers = [*include_headers].map{|h| h.to_s.downcase}
+
+        headers.each do |key, value|
+          headers.delete key unless
+            include_headers.include? key.to_s.downcase
         end
 
-        "binary"
+        headers
+
+      when true
+        headers
       end
+    end
 
 
-      ##
-      # Assigns self.encoding to the passed string if
-      # it responds to 'force_encoding'.
-      # Returns the string given with the new encoding.
+    ##
+    # Returns the header portion of the raw http response.
 
-      def udpate_encoding str
-        str.force_encoding self.encoding if str.respond_to? :force_encoding
-        str
-      end
+    def raw_header include_headers=true
+      headers = "#{@raw.split("\r\n\r\n", 2)[0]}\r\n"
 
+      case include_headers
+      when nil, false
+        nil
 
-      ##
-      # Returns the raw http response.
+      when Array, String
+        includes = [*include_headers].join("|")
+        headers.scan(%r{^((?:#{includes}): [^\n]*\n)}im).flatten.join
 
-      def raw
-        @raw
+      when true
+        headers
       end
+    end
 
 
-      ##
-      # Returns the body data parsed according to the content type.
-      # If no parser is given will look for the default parser based on
-      # the Content-Type, or will return the cached parsed body if available.
+    ##
+    # Check if this is a redirect response.
 
-      def parsed_body parser=nil
-        @parsed_body ||= nil
+    def redirect?
+      @code.to_s =~ /^30\d$/
+    end
 
-        return @parsed_body if @parsed_body && !parser
 
-        if String === parser
-          parser = Kronk.parser_for(parser) || Kronk.find_const(parser)
-        end
+    ##
+    # Follow the redirect and return a new Response instance.
+    # Returns nil if not redirect-able.
 
-        parser ||= Kronk.parser_for self['Content-Type']
+    def follow_redirect opts={}
+      return if !redirect?
+      Request.new(@_res['Location'], opts).retrieve
+    end
 
-        raise MissingParser,
-          "No parser for Content-Type: #{self['Content-Type']}" unless parser
 
-        @parsed_body = parser.parse self.body
-      end
+    ##
+    # Returns the raw response with selective headers and/or the body of
+    # the response. Supports the following options:
+    # :no_body:: Bool - Don't return the body; default nil
+    # :with_headers:: Bool/String/Array - Return headers; default nil
 
+    def selective_string options={}
+      str = @body unless options[:no_body]
 
-      ##
-      # Returns the parsed header hash.
+      if options[:with_headers]
+        header = raw_header(options[:with_headers])
+        str = [header, str].compact.join "\r\n"
+      end
+
+      str
+    end
 
-      def parsed_header include_headers=true
-        headers = self.to_hash.dup
 
-        case include_headers
-        when nil, false
-          nil
+    ##
+    # Returns the parsed response with selective headers and/or the body of
+    # the response. Supports the following options:
+    # :no_body:: Bool - Don't return the body; default nil
+    # :with_headers:: Bool/String/Array - Return headers; default nil
+    # :parser:: Object - The parser to use for the body; default nil
+    # :ignore_data:: String/Array - Removes the data from given data paths
+    # :only_data:: String/Array - Extracts the data from given data paths
+
+    def selective_data options={}
+      data = nil
+
+      unless options[:no_body]
+        data = parsed_body options[:parser]
+      end
 
-        when Array, String
-          include_headers = [*include_headers].map{|h| h.to_s.downcase}
+      if options[:with_headers]
+        data = [parsed_header(options[:with_headers]), data].compact
+      end
 
-          headers.each do |key, value|
-            headers.delete key unless
-              include_headers.include? key.to_s.downcase
-          end
+      Path::Transaction.run data, options do |t|
+        t.select(*options[:only_data])
+        t.delete(*options[:ignore_data])
+      end
+    end
 
-          headers
 
-        when true
-          headers
-        end
+    ##
+    # Returns a String representation of the response, the response body,
+    # or the response headers, parsed or in raw format.
+    # Options supported are:
+    # :parser:: Object/String - the parser for the body; default nil (raw)
+    # :struct:: Boolean - Return data types instead of values
+    # :only_data:: String/Array - extracts the data from given data paths
+    # :ignore_data:: String/Array - defines which data points to exclude
+    # :raw:: Boolean - Force using the unparsed raw response
+    # :keep_indicies:: Boolean - indicies of modified arrays display as hashes.
+    # :with_headers:: Boolean/String/Array - defines which headers to include
+
+    def stringify options={}
+      if !options[:raw] && (options[:parser] || @parser)
+        data = selective_data options
+        Diff.ordered_data_string data, options[:struct]
+      else
+        selective_string options
       end
 
+    rescue MissingParser
+      Cmd.verbose "Warning: No parser for #{@_res['Content-Type']} [#{@uri}]"
+        selective_string options
+    end
 
-      ##
-      # Returns the header portion of the raw http response.
 
-      def raw_header include_headers=true
-        headers = "#{raw.split("\r\n\r\n", 2)[0]}\r\n"
+    ##
+    # Check if this is a 2XX response.
 
-        case include_headers
-        when nil, false
-          nil
+    def success?
+      @code.to_s =~ /^2\d\d$/
+    end
 
-        when Array, String
-          includes = [*include_headers].join("|")
-          headers.scan(%r{^((?:#{includes}): [^\n]*\n)}im).flatten.join
 
-        when true
-          headers
-        end
+    ##
+    # Assign how long the request took in seconds.
+
+    def time= new_time
+      @time = new_time
+      @byterate = self.total_bytes / @time.to_f if @raw && @time > 0
+      @time
+    end
+
+
+    ##
+    # Number of bytes of the response including the header.
+
+    def total_bytes
+      self.raw.bytes.count
+    end
+
+
+    private
+
+
+    ##
+    # Creates a Net::HTTPRequest instance from an IO instance.
+
+    def request_from_io resp_io
+      # On windows, read the full file and insert contents into
+      # a StringIO to avoid failures with IO#read_nonblock
+      if Kronk::Cmd.windows? && File === resp_io
+        resp_io = WinFileIO.new resp_io.path, io.read
       end
 
+      io = Net::BufferedIO === resp_io ? resp_io : Net::BufferedIO.new(resp_io)
+      io.debug_output = debug_io = StringIO.new
+
+      begin
+        resp = Net::HTTPResponse.read_new io
+        resp.reading_body io, true do;end
 
-      ##
-      # Returns the raw response with selective headers and/or the body of
-      # the response. Supports the following options:
-      # :no_body:: Bool - Don't return the body; default nil
-      # :with_headers:: Bool/String/Array - Return headers; default nil
+      rescue Net::HTTPBadResponse
+        ext = "text/html"
+        ext = File.extname(resp_io.path) if WinFileIO === resp_io
 
-      def selective_string options={}
-        str = self.body unless options[:no_body]
+        resp_io.rewind
+        resp = HeadlessResponse.new resp_io.read, ext
 
-        if options[:with_headers]
-          header = raw_header(options[:with_headers])
-          str = [header, str].compact.join "\r\n"
+      rescue EOFError
+        # If no response was read because it's too short
+        unless resp
+          resp_io.rewind
+          resp = HeadlessResponse.new resp_io.read, "html"
         end
+      end
 
-        str
+      resp.instance_eval do
+        @socket ||= true
+        @read   ||= true
       end
 
+      [resp, debug_io]
+    end
+
 
-      ##
-      # Returns the parsed response with selective headers and/or the body of
-      # the response. Supports the following options:
-      # :no_body:: Bool - Don't return the body; default nil
-      # :with_headers:: Bool/String/Array - Return headers; default nil
-      # :parser:: Object - The parser to use for the body; default nil
-      # :ignore_data:: String/Array - Removes the data from given data paths
-      # :only_data:: String/Array - Extracts the data from given data paths
+    ##
+    # Read the raw response from a debug_output instance and return an array
+    # containing the raw request, response, and number of bytes received.
 
-      def selective_data options={}
-        data = nil
+    def read_raw_from debug_io
+      req = nil
+      resp = ""
+      bytes = nil
 
-        unless options[:no_body]
-          data = parsed_body options[:parser]
-        end
+      debug_io.rewind
+      output = debug_io.read.split "\n"
 
-        if options[:with_headers]
-          data = [parsed_header(options[:with_headers]), data].compact
-        end
+      if output.first =~ %r{<-\s(.*)}
+        req = instance_eval $1
+        output.delete_at 0
+      end
+
+      if output.last =~ %r{read (\d+) bytes}
+        bytes = $1.to_i
+        output.delete_at(-1)
+      end
 
-        DataSet.new(data).fetch options
+      output.map do |line|
+        next unless line[0..2] == "-> "
+        resp << instance_eval(line[2..-1])
       end
+
+      [req, resp, bytes]
+    end
+
+
+    ##
+    # Assigns self.encoding to the passed string if
+    # it responds to 'force_encoding'.
+    # Returns the string given with the new encoding.
+
+    def try_force_encoding str
+      str.force_encoding @encoding if str.respond_to? :force_encoding
+      str
     end
   end
 
 
+
   ##
   # Mock response object without a header for body-only http responses.
 
   class HeadlessResponse
 
-    include Response::Helpers
-
     attr_accessor :body, :code
 
     def initialize body, file_ext=nil
       @body = body
       @raw  = body
+      @code = "200"
 
-      encoding = body.encoding rescue "UTF-8"
+      encoding = body.respond_to?(:encoding) ? body.encoding : "UTF-8"
 
       @header = {
-        'Content-Type' => ["#{file_ext}; charset=#{encoding}"]
+        'Content-Type' => ["text/#{file_ext}; charset=#{encoding}"]
       }
     end
 
@@ -269,7 +450,7 @@ class Kronk
     # Interface method only. Returns empty hash.
 
     def to_hash
-      Hash.new
+      @header
     end
   end
 end