# Copyright (c) 2020 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
# frozen_string_literal: true

cs__scoped_require 'resolv'
cs__scoped_require 'timeout'

cs__scoped_require 'contrast/utils/object_share'
cs__scoped_require 'contrast/utils/string_utils'
cs__scoped_require 'contrast/utils/comment_range'
cs__scoped_require 'contrast/utils/hash_digest'
cs__scoped_require 'contrast/components/interface'

module Contrast
  module Agent
    # This class is the Contrast representation of the Rack::Response object. It
    # provides access to the original Rack::Response object as well as extracts
    # data in a format that the Agent expects, caching those transformations in
    # order to avoid repeatedly creating Strings & thrashing GC.
    class Response
      include Contrast::Components::Interface
      access_component :logging

      extend Forwardable

      attr_reader :rack_response

      def initialize rack_response
        @rack_response = rack_response
        @is_array = !rack_response.is_a?(Rack::Response)
      end

      def document_type
        case content_type
        when /xml/i
          :XML
        when /json/i
          :JSON
        when /html/i
          :NORMAL
        end
      end

      # B/c the response can change, we can't memoize this :(
      def dtm
        context_response = Contrast::Api::Dtm::HttpResponse.new
        context_response.response_code = response_code.to_i
        headers.each_pair do |key, value|
          k = Contrast::Utils::StringUtils.force_utf8(key)
          v = Contrast::Utils::StringUtils.force_utf8(value)
          context_response.response_headers[k] = v
        end
        context_response.parsed_response_headers = true

        context_response.response_body = Contrast::Utils::StringUtils.force_utf8(body)
        context_response.parsed_response_body = true

        doc_type = document_type
        context_response.document_type = doc_type if doc_type

        context_response
      end

      def response_code
        return unless @rack_response

        @is_array ? @rack_response[0].to_i : @rack_response.status
      end

      def headers
        return unless @rack_response

        if @is_array
          @rack_response[1]
        else
          @rack_response.headers
        end
      end

      # Header keys upcased and any underscores replaced with dashes
      # We cannot memoize this, b/c response headers can change during the
      # request lifecycle.
      def normalized_headers
        hash = {}
        headers.each_pair do |header_name, header_value|
          hash[Contrast::Utils::StringUtils.normalized_key(header_name)] = header_value
        end
        hash
      end

      CONTENT_TYPE_HEADER = 'CONTENT-TYPE'.cs__freeze
      def content_type
        return unless @rack_response

        if @is_array
          normalized_headers[CONTENT_TYPE_HEADER]
        else
          @rack_response.content_type
        end
      end

      def header key
        return unless @rack_response

        if @is_array
          normalized_headers[Contrast::Utils::StringUtils.normalized_key(key)]
        else
          @rack_response.get_header(key)
        end
      end

      def set_header key, value
        return unless @rack_response

        if @is_array
          Rack::Utils.set_cookie_header!(@rack_response[1], key, value)
        elsif @rack_response.is_a?(Rack::Response)
          @rack_response.set_header(key, value)
        end
      end

      # The response body can change during the request lifecycle
      # We should not extract it out as a variable here, or we'll miss those
      # changes. (headdesk)
      def body
        return unless @rack_response

        if @is_array
          extract_body(@rack_response[2])
        elsif Contrast::Utils::DuckUtils.quacks_to?(@rack_response, :body)
          extract_body(@rack_response.body)
        end
      end

      def update_body body_string
        return unless @rack_response

        successfully_updated_body = true
        if @is_array
          if @rack_response[2].is_a?(Rack::BodyProxy)
            successfully_updated_body = update_rack_body_proxy(body_string, true)
          else
            @rack_response[2] = [body_string]
          end
        elsif @rack_response.body.is_a?(Rack::BodyProxy)
          successfully_updated_body = update_rack_body_proxy(body_string)
        else
          @rack_response.body = body_string
        end
        update_content_length(body_string.bytesize) if successfully_updated_body
      end

      # Set the length header for this response. This value should be set ot the
      # bytesize, NOT THE LENGTH, of the response body. Otherwise, we may get
      # got by the Lint thing.
      CONTENT_LENGTH_HEADER = 'Content-Length'.cs__freeze
      def update_content_length length
        headers[CONTENT_LENGTH_HEADER] = length.to_s
      end

      private

      HTTP_PREFIX = /^[Hh][Tt][Tt][Pp][_-]/i.cs__freeze

      def update_rack_body_proxy body_string, is_array = false
        top_body_proxy = is_array ? @rack_response[2] : @rack_response
        parent_body_proxy = top_body_proxy
        until (next_body = parent_body_proxy.instance_variable_get(:@body)).cs__class != Rack::BodyProxy
          parent_body_proxy = next_body
        end

        if next_body.cs__class.to_s == 'ActionDispatch::Response::RackBody'
          modified_response = next_body.instance_variable_defined?(:@response) ? next_body.instance_variable_get(:@response) : nil
          if modified_response
            modified_response.body = body_string
            next_body.instance_variable_set(:@response, modified_response)
          end
        elsif next_body.is_a?(Array) && next_body[0].cs__class.to_s == 'ActionView::OutputBuffer'
          new_body = ActionView::OutputBuffer.new(body_string)
          next_body[0] = new_body
        else
          logger.warn("Detected unsupported Rack::BodyProxy internal response class #{ next_body.cs__class }")
          return false
        end
        true
      end

      def extract_body body
        return nil unless body

        if defined?(Rack::File) && body.is_a?(Rack::File)
          # not sure what to do in this situation, so don't do anything.
          nil
        elsif Contrast::Utils::DuckUtils.quacks_to?(body, :each)
          acc = []
          body.each { |tmp| acc << read_or_string(tmp) }
          acc.compact.join(Contrast::Utils::ObjectShare::NEW_LINE)
        else
          read_or_string(body)
        end
      end

      def read_or_string obj
        return nil unless obj

        if Contrast::Utils::DuckUtils.quacks_to?(obj, :read)
          tmp = obj.read
          obj.rewind
          tmp
        else
          obj.to_s
        end
      end
    end
  end
end