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

# This class is responsible for the origination of traces. A Source is any method
# that returns an untrusted value. In general, these sources are request methods
# as the request object is user controlled.
#
# Going forward, we may add in Dynamic sources (determined at runtime) based on
# database configs or other variables (used for Stored XSS or other persisted
# vulnerability detection rules)
cs__scoped_require 'set'

cs__scoped_require 'contrast/utils/object_share'
cs__scoped_require 'contrast/utils/sha256_builder'
cs__scoped_require 'contrast/agent/assess/adjusted_span'
cs__scoped_require 'contrast/agent/assess/policy/source_validation/source_validation'

cs__scoped_require 'contrast/components/interface'

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
          include Contrast::Components::Interface
          access_component :logging, :analysis

          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
            # @return [Object, nil] the tracked Return or nil if no changes
            #   were made
            def source_patchers method_policy, object, ret, args
              return if method_policy.source_node.nil?

              current_context = Contrast::Agent::REQUEST_TRACKER.current
              return unless current_context&.analyze_request? && ASSESS.enabled?

              replaced_return = nil
              source_node = method_policy.source_node

              target = determine_target(source_node, object, ret, args)

              # We don't propagate to frozen things that haven't been tracked
              # before. By definition, something that is a source has not
              # previously been tracked; therefore, we can break out early.
              if target.cs__frozen?
                # That being said, we don't have enough context to know if we
                # can make this assumption and still function, so we'll allow for
                # source tracking of frozen things by a common config setting.
                #
                # Rails' StrongParameters make a case for this to be default
                # behavior
                return replaced_return unless ASSESS.track_frozen_sources?

                # If we're tracking the frozen target, we need to unfreeze
                # (dup) it to track and then freeze that result. For
                # simplicities sake, we ONLY do this if the return is the
                # target (I don't want to have to deal with unfreezing self)
                return replaced_return unless source_node.targets[0] == Contrast::Utils::ObjectShare::RETURN_KEY

                restore_frozen_state = true
                ret = Contrast::Utils::FreezeUtil.unfreeze_dup(ret)
                target = ret
              end

              apply_source(current_context, source_node, target, object, ret, source_node.type, nil, *args)

              ret.cs__freeze if restore_frozen_state
              ret
            end

            private

            # This is our method that actually taints the object our
            # source_node targets.
            #
            # @param context [Contrast::Utils::ThreadTracker] the current request
            #   context
            # @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 object [Object] the Object on which the method was invoked
            # @param ret [Object] the Return of the invoked method
            # @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_source context, source_node, target, object, ret, source_type, source_name = nil, *args
              return unless context && source_node && target

              source_name ||= determine_source_name(source_node, object, ret, *args)
              # We know we only work on certain things.
              # Skip if this isn't one of them
              if Contrast::Utils::DuckUtils.quacks_to?(target, :cs__properties)
                apply_tags(source_node, target, object, ret, source_type, source_name, *args)
                # 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
              elsif Contrast::Utils::DuckUtils.iterable_hash?(target)
                to_replace = []
                target.each_pair do |key, value|
                  # We only do this for Strings b/c of the way Hash lookup works.
                  # To replace another object would break hash lookup and,
                  # therefore, the application
                  if can_track_hash_key?(key, target)
                    key = Contrast::Utils::FreezeUtil.unfreeze_dup(key)
                    to_replace << key
                  end
                  apply_source(context, source_node, key, object, ret, key_type(source_type), key, *args)
                  apply_source(context, source_node, value, object, ret, source_type, key, *args)
                end

                # Hash is designed to keep one instance of the string key in it.
                # We need to remove the existing one and replace it with our new
                # tracked one.
                to_replace.each do |key|
                  key.cs__freeze
                  value = target[key]
                  target.delete(key)
                  target[key] = value
                end
                # 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 { |value| apply_source(context, source_node, value, object, ret, source_type, source_name, *args) }
              end
            rescue StandardError => e
              logger.warn('Unable to apply source', e, node_id: source_node.id)
            end

            def can_track_hash_key? key, target
              ASSESS.track_frozen_sources? &&
                  key.is_a?(String) &&
                  Contrast::Utils::DuckUtils.quacks_to?(target, :delete)
            end

            def apply_tags source_node, target, object, ret, source_type, source_name, *args
              # don't apply second source -- probably needs tuning later if we
              # use more than 'UNTRUSTED' in our sources
              return if target.cs__tracked? || target.cs__frozen?

              # 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)
                target.cs__properties.add_tag(tag, Contrast::Agent::Assess::AdjustedSpan.new(0, length))
                target.cs__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
              target.cs__properties.build_event(source_node, target, object, ret, args, source_type, source_name)
            end

            # Find the name of the source
            #
            # @param source_node [Contrast::Agent::Assess::Policy::SourceNode]
            #   the node to direct applying this source event
            # @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
            # @return [String, nil] the human readable name of the target to
            #   which this source event applies, or nil if none provided by the
            #   node
            def determine_source_name source_node, object, ret, *args
              return source_node.get_property('dynamic_source_name') if source_node.type == 'UNTRUSTED_DATABASE'

              source_node_source = source_node.sources[0]
              case source_node_source
              when nil
                nil
              when Contrast::Utils::ObjectShare::RETURN_KEY
                ret
              when Contrast::Utils::ObjectShare::OBJECT_KEY
                object
              else
                args[source_node_source]
              end
            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 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
            # @return [Object] the target to which this source event applies
            def determine_target source_node, object, ret, args
              source_target = source_node.targets[0]
              case source_target
              when Contrast::Utils::ObjectShare::RETURN_KEY
                ret
              when Contrast::Utils::ObjectShare::OBJECT_KEY
                object
              else
                if source_target.is_a?(Integer)
                  args[source_target]
                  # If this isn't an index param, it's a named one. R.I.P.
                else
                  arg = nil
                  args.each do |search|
                    next unless search.is_a?(Hash)

                    arg = search[source_target]
                    break if arg
                  end
                  arg
                end
              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
          end
        end
      end
    end
  end
end