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

require 'resolv'
require 'timeout'

require 'contrast/utils/object_share'
require 'contrast/utils/string_utils'
require 'contrast/utils/hash_digest'
require 'contrast/components/logger'

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::Logger::InstanceMethods

      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|
          append_pair(context_response.normalized_response_headers, key, value)
        end
        context_response.response_body_binary = Contrast::Utils::StringUtils.force_utf8(body)

        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

      def content_type
        return unless rack_response

        if @is_array
          headers[Rack::CONTENT_TYPE]
        else
          rack_response.content_type
        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.
      def body
        return unless rack_response

        body_content = @is_array ? rack_response[2] : rack_response.body
        extract_body(body_content)
      end

      private

      # From the dtm for normalized_response_headers:
      #   Key is UPPERCASE_UNDERSCORE
      #
      #   Example: Content-Type: text/html; charset=utf-8
      #   "CONTENT_TYPE" => Content-Type,["text/html; charset=utf8"]
      def append_pair map, key, value
        return unless key && value
        return if value.is_a?(Hash)

        safe_key = Contrast::Utils::StringUtils.force_utf8(key)
        hash_key = Contrast::Utils::StringUtils.normalized_key(safe_key)
        map[hash_key] ||= Contrast::Api::Dtm::Pair.new
        map[hash_key].key = safe_key
        map[hash_key].values << Contrast::Utils::StringUtils.force_utf8(value)
      end

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

      # Given some holder of the content of the response's body, extract that
      # content and return it as a String
      #
      # @param body [String, Rack::File, Rack::BodyProxy,
      #   ActionDispatch::Response::RackBody, Rack::Response] Something that
      #   holds, wraps, or is the body of the Response
      # @return [nil, String] the content of the body
      def extract_body body
        return 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 body.is_a?(Rack::BodyProxy)
          handle_rack_body_proxy(body)
        elsif (defined?(ActionDispatch::Response::RackBody) && body.is_a?(ActionDispatch::Response::RackBody)) ||
              body.is_a?(Rack::Response)

          extract_body(body.body)
        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)
        elsif ActionView::OutputBuffer
          # https://stackoverflow.com/questions/15654676/how-to-convert-activesupportsafebuffer-to-string
          body.to_str
        else
          read_or_string(body)
        end
      end

      def handle_rack_body_proxy body
        next_body = body.instance_variable_get(:@body)
        case next_body
        when Array
          extract_body(next_body[0])
        else
          extract_body(next_body)
        end
      end

      def read_or_string obj
        return 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