# Copyright (c) 2020 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/components/interface'
require 'contrast/utils/object_share'
require 'contrast/utils/sha256_builder'

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 :analysis, :logging

          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?

              source_node = method_policy.source_node
              target = determine_target(source_node, object, ret, args)
              restore_frozen_state = false
              if target.cs__frozen? && !Contrast::Agent::Assess::Tracker.trackable?(target)
                return unless ASSESS.track_frozen_sources?
                return unless source_node.targets[0] == Contrast::Utils::ObjectShare::RETURN_KEY

                dup = safe_dup(ret)
                return unless dup

                restore_frozen_state = true
                ret = dup
                target = ret
                Contrast::Agent::Assess::Tracker.pre_freeze(ret)
                ret.cs__freeze
                # double check that we were able to finalize the replaced return
                return unless Contrast::Agent::Assess::Tracker.trackable?(target)
              end
              apply_source(current_context, source_node, target, object, ret, source_node.type, nil, *args)
              restore_frozen_state ? ret : nil
            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::Agent::Assess::Tracker.trackable?(target)
                apply_tags(source_node, target, object, ret, source_type, source_name, *args)
              elsif Contrast::Utils::DuckUtils.iterable_hash?(target)
                apply_hash_tags(context, source_node, target, object, ret, 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 { |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

            # 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 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 args [Array<Object>] the Arguments with which the method
            #   was invoked
            def apply_hash_tags context, source_node, target, object, ret, source_type, *args
              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 replace_hash_key?(key, target)
                  key = key.dup
                  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
              handle_hash_key(target, to_replace)
            end

            # Given an unfrozen hash, if the key is a String, we should replace
            # it with one that we can finalize, allowing us to track that key.
            # This method handles checking if that replace can and should
            # occur.
            #
            # @param key [Object] the key in the hash that may need replacing.
            # @param hash [Hash] the hash to which the key belongs.
            # @return [Boolean] whether replace the key in the hash or not.
            def replace_hash_key? key, hash
              ASSESS.track_frozen_sources? &&
                  !hash.cs__frozen? &&
                  key.is_a?(String) &&
                  !Contrast::Agent::Assess::Tracker.trackable?(key)
            end

            # Safely duplicate the target, or return nil
            #
            # @param target [Object] the thing to check for duplication
            def safe_dup target
              target.dup
            rescue StandardError => _e
              nil
            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.
            def handle_hash_key target, to_replace
              to_replace.each do |key|
                Contrast::Agent::Assess::Tracker.pre_freeze(key)
                key.cs__freeze
                value = target.delete(key)
                target[key] = value
              end
            end

            def apply_tags source_node, target, object, ret, 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
              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
                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
          end
        end
      end
    end
  end
end