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

require 'set'
require 'contrast/agent/assess/policy/source_validation/source_validation'
require 'contrast/agent/excluder/excluder'
require 'contrast/components/logger'
require 'contrast/utils/object_share'
require 'contrast/utils/sha256_builder'
require 'contrast/utils/assess/source_method_utils'
require 'contrast/utils/assess/event_limit_utils'
require 'contrast/agent/assess/events/event_data'

module Contrast
  module Agent
    module Assess
      module Policy
        # This class controls the actions we take on Sources, as determined by our Assess policy. It indicates what
        # actions we should take in order to mark data as User Input and treat it as untrusted, starting the dataflows
        # used in Assess vulnerability detection.
        module SourceMethod
          extend Contrast::Components::Logger::InstanceMethods
          extend Contrast::Utils::Assess::SourceMethodUtils
          extend Contrast::Utils::Assess::EventLimitUtils

          PARAMETER_TYPE     = 'PARAMETER'
          PARAMETER_KEY_TYPE = 'PARAMETER_KEY'
          HEADER_TYPE        = 'HEADER'
          HEADER_KEY_TYPE    = 'HEADER_KEY'
          COOKIE_TYPE        = 'COOKIE'
          COOKIE_KEY_TYPE    = 'COOKIE_KEY'

          class << self
            # This is called from within our woven proc. It will be called as if it were inline in the Rack application.
            #
            # @param method_policy [Contrast::Agent::Patching::Policy::MethodPolicy] the policy that applies to the
            #   method being called
            # @param object [Object] the Object on which the method was invoked
            # @param ret [Object] the Return of the invoked method
            # @param args [Array<Object>] the Arguments with which the method was invoked
            def apply_source method_policy, object, ret, args
              logger.trace_with_time('Elapsed time for Contrast::Agent::Assess::Policy::SourceMethod#apply_source') do
                return unless analyze?(method_policy, object, ret, args)
                return if event_limit?(method_policy)
                return unless (source_node = method_policy.source_node)
                return if excluded_by_url?

                # used to hold the object and ret
                source_data = Contrast::Agent::Assess::Events::EventData.new(nil, nil, object, ret, nil)

                return unless (target = determine_target(source_node, source_data, args))
                return if target.cs__frozen? && !Contrast::Agent::Assess::Tracker.trackable?(target)

                process_source(source_node, target, source_data, source_node.type, nil, *args)
              end
            end

            private

            # This is our method that actually taints the object our source_node targets.
            #
            # @param source_node [Contrast::Agent::Assess::Policy::SourceNode] the node to direct applying this source
            #   event
            # @param target [Object] the target of the Source Event
            # @param source_data [Contrast::Agent::Asssess::Events::EventData] used to hold object and ret
            # @param source_type [String] the type of this source, from the source_node, or a KEY_TYPE if invoked for a
            #   map
            # @param source_name [String, nil] the name of this source, i.e. the key used to accessed if from a map or
            #   nil if a type like BODY
            # @param args [Array<Object>] the Arguments with which the method was invoked
            def process_source source_node, target, source_data, source_type, source_name = nil, *args # rubocop:disable Metrics/PerceivedComplexity
              context = Contrast::Agent::REQUEST_TRACKER.current
              return unless context && source_node && target

              increment_event_count(source_node)

              source_name ||= determine_source_name(source_node, source_data.object, source_data.ret, *args)

              return if excluded_by_input?(source_type, source_name)

              # We know we only work on certain things.
              # Skip if this isn't one of them
              if Contrast::Agent::Assess::Tracker.trackable?(target)
                apply_tags(source_node, target, source_data, source_type, source_name, *args)
              elsif Contrast::Utils::DuckUtils.iterable_hash?(target)
                apply_hash_tags(source_node, target, source_data, source_type, *args)
                # While we don't taint arrays themselves, we may taint the things they hold. Let's pass their keys and
                # values back to ourselves and try again
              elsif Contrast::Utils::DuckUtils.iterable_enumerable?(target)
                target.each do |value|
                  process_source(source_node, value, source_data, source_type, source_name, *args)
                end
              end
              nil
            rescue StandardError => e
              logger.warn('Unable to apply source', e, node_id: source_node.id)
            end

            # While we don't taint hashes themselves, we may taint the things they hold. Let's pass their keys and
            # values back to ourselves and try again
            #
            # @param source_node [Contrast::Agent::Assess::Policy::SourceNode] the node to direct applying this source
            #   event
            # @param target [Object] the target of the Source Event
            # @param source_data [Contrast::Agent::Asssess::Events::EventData] used to hold object and ret
            # @param source_type [String] the type of this source, from the source_node, or a KEY_TYPE if invoked for a
            #   map
            # @param args [Array<Object>] the Arguments with which the method was invoked
            def apply_hash_tags source_node, target, source_data, source_type, *args
              target.each_pair do |key, value|
                process_source(source_node, key, source_data, key_type(source_type), key, *args)
                process_source(source_node, value, source_data, source_type, key, *args)
              end
            end

            # @param source_node [Contrast::Agent::Assess::Policy::SourceNode] the node to direct applying this source
            #   event
            # @param target [Object] the target of the Source Event
            # @param source_data [Contrast::Agent::Asssess::Events::EventData] used to hold object and ret
            # @param source_type [String] the type of this source, from the source_node, or a KEY_TYPE if invoked for a
            #   map
            # @param source_name [String, nil] the name of this source, i.e. the key used to accessed if from a map or
            #   nil if a type like BODY
            # @param args [Array<Object>] the Arguments with which the method was invoked
            def apply_tags source_node, target, source_data, source_type, source_name, *args
              # don't apply tags if we can't track the thing
              return unless Contrast::Agent::Assess::Tracker.trackable?(target)
              # don't apply second source -- probably needs tuning later if we use more than 'UNTRUSTED' in our sources
              return if Contrast::Agent::Assess::Tracker.tracked?(target)
              return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))

              # otherwise for each tag this source_node applies, create a tag range on the target object. I realize
              # this looping is counter-intuitive from the above message, that's why we're revisiting.
              source_node.tags.each do |tag|
                next unless Contrast::Agent::Assess::Policy::SourceValidation.valid?(tag, source_type, source_name)

                length = Contrast::Utils::StringUtils.ret_length(target)
                properties.add_tag(tag, 0...length)
                properties.add_properties(source_node.properties)
                logger.trace('Source detected', node_id: source_node.id, target_id: target.__id__, tag: tag)
              end
              # make a representation of this method that TeamServer can render
              event_data = Contrast::Agent::Assess::Events::EventData.new(source_node,
                                                                          target,
                                                                          source_data.object,
                                                                          source_data.ret,
                                                                          args)
              properties.build_event(event_data, source_type, source_name)
            end

            # Find the literal target of the propagation
            #
            # @param source_node [Contrast::Agent::Assess::Policy::SourceNode] the node to direct applying this source
            #   event
            # @param source_data [Contrast::Agent::Asssess::Events::EventData] used to hold object and ret
            # @param args [Array<Object>] the Arguments with which the method was invoked
            # @return [Object] the target to which this source event applies
            def determine_target source_node, source_data, args
              source_target = source_node.targets[0]
              case source_target
              when Contrast::Utils::ObjectShare::RETURN_KEY
                source_data.ret
              when Contrast::Utils::ObjectShare::OBJECT_KEY
                source_data.object
              else
                args[source_target]
              end
            end

            # Simple helper method to flip the type from value to key when the source is the key of a Hash
            #
            # @param source_type [String] the original value source type
            # @return [String] the key form of the source type, if one exists, else the original source type
            def key_type source_type
              case source_type
              when PARAMETER_TYPE
                PARAMETER_KEY_TYPE
              when HEADER_TYPE
                HEADER_KEY_TYPE
              when COOKIE_TYPE
                COOKIE_KEY_TYPE
              else
                source_type
              end
            end

            # Should a source be excluded because it matches one of the input exclusion rules?
            #
            # @return [Boolean]
            def excluded_by_input? source_type, source_name
              context = Contrast::Agent::REQUEST_TRACKER.current
              Contrast::SETTINGS.excluder.assess_excluded_by_input?(context.request, source_type, source_name)
            end

            # Should a source be excluded because it matches one of the url exclusion rules?
            #
            # @return [Boolean]
            def excluded_by_url?
              context = Contrast::Agent::REQUEST_TRACKER.current
              Contrast::SETTINGS.excluder.assess_excluded_by_url?(context.request)
            end
          end
        end
      end
    end
  end
end