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

module Contrast
  module Agent
    module Assess
      module Policy
        module Propagator
          # Rack::Protection offers several protections against vulnerabilities. Of these, some apply to dataflow and
          # need to be accounted for in order to properly tag data. Others apply to configurations and may be used to
          # suppress configuration vulnerabilities in the future.
          class RackProtection < Contrast::Agent::Assess::Policy::Propagator::Base
            class << self
              # Our custom instrumentation for the Rack::Protection::EscapedParams#escape_string method
              # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode] the node that governs this
              #   propagation event.
              # @param preshift [Contrast::Agent::Assess::PreShift] The capture of the state of the code just prior to
              #   the invocation of the patched method.
              # @param ret [nil, String] the target to which to propagate.
              # @return [nil, String] ret
              def escaped_params propagation_node, preshift, ret, _block
                Contrast::Agent::Assess::Policy::Propagator::Splat.propagate(propagation_node, preshift, ret)
                apply_escaper_tags(preshift.object, ret)
                ret
              end

              private

              # Rack::Protection::EscapedParams can be configured such that it only applies certain escape. We need
              # to account for the configuration of the individual escapes when applying tags.
              #
              # @param escaper [Rack::Protection::EscapedParams] the instance of Rack::Protection::EscapedParams
              #   applying the escape_string
              # @param ret [String] the result of the escape
              def apply_escaper_tags escaper, ret
                # I don't know how this could not be an instance of Rack::Protection::EscapedParams, but I don't want
                # to chance it.
                return unless escaper.cs__is_a?(Rack::Protection::EscapedParams)
                return unless (properties = Contrast::Agent::Assess::Tracker.properties(ret))

                tags = []
                untags = []
                if escaper.instance_variable_get(:@html)
                  tags << 'HTML_ENCODED'
                  untags << 'HTML_DECODED'
                end

                if escaper.instance_variable_get(:@javascript)
                  tags << 'JAVASCRIPT_ENCODED'
                  untags << 'JAVASCRIPT_DECODED'
                end

                if escaper.instance_variable_get(:@url)
                  tags << 'URL_ENCODED'
                  untags << 'URL_DECODED'
                end

                length = Contrast::Utils::StringUtils.ret_length(ret)
                tags.each do |tag|
                  properties.add_tag(tag, 0...length)
                end

                untags.each do |tag|
                  properties.delete_tags(tag)
                end
              end
            end
          end
        end
      end
    end
  end
end