# Copyright (c) 2022 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/propagator'
require 'contrast/components/logger'
require 'contrast/utils/object_share'
require 'contrast/utils/sha256_builder'
require 'contrast/utils/assess/propagation_method_utils'
require 'contrast/agent/assess/events/event_data'

module Contrast
  module Agent
    module Assess
      module Policy
        # This class is responsible for the continuation of traces. A Propagator is any method that transforms an
        # untrusted value. In general, these methods work on the String class or a holder of Strings.
        # rubocop:disable Metrics/ModuleLength
        module PropagationMethod
          extend Contrast::Components::Logger::InstanceMethods
          extend Contrast::Utils::Assess::PropagationMethodUtils

          class << self
            # @param method_policy [Contrast::Agent::Patching::Policy::MethodPolicy] the policy that governs the
            #   patches to this method
            # @param preshift [Contrast::Agent::Assess::PreShift] The capture of the state of the code just prior to
            #   the invocation of the patched method.
            # @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
            # @param block [Block] the Block passed to the original method
            # @return [Object, nil] the tracked Return or nil if no changes were made; will replace the return of the
            #   original function if not nil
            def apply_propagation method_policy, preshift, object, ret, args, block
              return unless method_policy.propagation_node
              return unless preshift

              propagation_node = method_policy.propagation_node

              target = determine_target(propagation_node, ret, object, args)
              propagation_data = Contrast::Agent::Assess::Events::EventData.new(nil, nil, object, ret, args)
              PropagationMethod.apply_propagator(propagation_node, preshift, target, propagation_data, block)
            end

            # I lied above. We had to figure out what the target of the propagation was. Now that we know, we'll
            # actually do things to it. Note that the return of this method will replace the original return of the
            # patched function unless it is nil, so be sure you're returning what you intend.
            #
            # @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 target [Object] the Target to which to propagate.
            # @param propagation_data [Contrast::Agent::Assess::Events::EventData] this will hold the
            #                         object [Object] the Object on which the method was invoked
            #                         ret [Object] the Return of the invoked method
            #                         args [Array<Object>] the Arguments with which the method was invoked
            # @param block [Block] the Block passed to the original method
            # @return [Object, nil] the tracked Return or nil if no changes were made; will replace the return of the
            #   original function if not nil
            def apply_propagator propagation_node, preshift, target, propagation_data, block
              return unless propagation_possible?(propagation_node, target)

              if propagation_node.action == DB_WRITE_ACTION
                Contrast::Agent::Assess::Policy::Propagator::DatabaseWrite.propagate(propagation_node,
                                                                                     preshift,
                                                                                     propagation_data.ret)
              elsif propagation_node.action == CUSTOM_ACTION
                Contrast::Agent::Assess::Policy::Propagator::Custom.propagate(propagation_node,
                                                                              preshift,
                                                                              propagation_data.ret,
                                                                              block)
              elsif propagation_node.action == SPLIT_ACTION
                Contrast::Agent::Assess::Policy::Propagator::Split.propagate(propagation_node, preshift, target)
              elsif Contrast::Utils::DuckUtils.iterable_hash?(target)
                handle_hash_propagation(propagation_node, preshift, target, propagation_data, block)
              elsif Contrast::Utils::DuckUtils.iterable_enumerable?(target)
                handle_enumerable_propagation(propagation_node, preshift, target, propagation_data, block)
              else
                handle_cs_properties_propagation(propagation_node, preshift, target, propagation_data, block)
              end
            rescue StandardError => e
              logger.warn('Unable to apply propagation', e, node_id: propagation_node.id)
              nil
            end

            # If this patcher has tags, apply them to the entire target
            # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode] the node that governs this
            #   propagation event.
            # @param target [Object] the Target to which to propagate.
            def apply_tags propagation_node, target
              return unless propagation_node.tags
              return unless (properties = Contrast::Agent::Assess::Tracker.properties(target))

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

            def context_available?
              !!Contrast::Agent::REQUEST_TRACKER.current
            end

            # If this patcher has tags, remove them from the entire target
            #
            # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode] the node that governs this
            #   propagation event.
            # @param target [Object] the Target to which to propagate.
            def apply_untags propagation_node, target
              return unless propagation_node.untags
              return unless (properties = Contrast::Agent::Assess::Tracker.properties(target))

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

            private

            # This is checked right before actual propagation
            # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode] the node that governs this
            #   propagation event.
            # @param target [Object] the Target to which to propagate.
            # @return [Boolean]
            def propagation_possible? propagation_node, target
              return false unless propagation_node && valid_target?(target, propagation_node)
              return false unless context_available? || Contrast::ASSESS.non_request_tracking?
              return false unless valid_length?(target, propagation_node.action)

              true
            end

            # Safely duplicate the target, or return nil
            #
            # @param target [Object] the thing to check for duplication
            # @return [Object, nil]
            def safe_dup target
              target.dup
            rescue StandardError => _e
              nil
            end

            # Iterate over each key and value in a hash to allow for propagation to each.
            #
            # @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 target [Object] the Target to which to propagate.
            # @param propagation_data [Contrast::Agent::Assess::Events::EventData] this will hold the
            #                         object [Object] the Object on which the method was invoked
            #                         ret [Object] the Return of the invoked method
            #                         args [Array<Object>] the Arguments with which the method was invoked
            # @param block [Block] the Block passed to the original method
            def handle_hash_propagation propagation_node, preshift, target, propagation_data, block
              target.each_pair do |key, value|
                apply_propagator(propagation_node, preshift, key, propagation_data, block)
                apply_propagator(propagation_node, preshift, value, propagation_data, block)
              end
            end

            # Iterate over each value in an enumerable to allow for propagation to each.
            #
            # @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 target [Object] the Target to which to propagate.
            # @param propagation_data [Contrast::Agent::Assess::Events::EventData] this will hold the
            #                         object [Object] the Object on which the method was invoked
            #                         ret [Object] the Return of the invoked method
            #                         args [Array<Object>] the Arguments with which the method was invoked
            # @param block [Block] the Block passed to the original method
            def handle_enumerable_propagation propagation_node, preshift, target, propagation_data, block
              target.each do |value|
                next if target == value

                apply_propagator(propagation_node, preshift, value, propagation_data, block)
              end
            end

            # Move the properties from the source(s) to the target of the propagation.
            #
            # @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 target [Object] the Target to which to propagate.
            # @param propagation_data [Contrast::Agent::Assess::Events::EventData] this will hold the
            #                         object [Object] the Object on which the method was invoked
            #                         ret [Object] the Return of the invoked method
            #                         args [Array<Object>] the Arguments with which the method was invoked
            # @param _block [Block] the Block passed to the original method
            def handle_cs_properties_propagation propagation_node, preshift, target, propagation_data, _block
              return if propagation_node.action == NOOP_ACTION
              return unless can_propagate?(propagation_node, preshift, target)
              return unless (propagation_class = find_propagation_class(propagation_node))

              restore_frozen_state = false
              if target.cs__frozen? && !Contrast::Agent::Assess::Tracker.trackable?(target)
                return unless can_handle_frozen?(propagation_node)
                return unless (dup = safe_dup(propagation_data.ret))

                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
              propagation_class.propagate(propagation_node, preshift, target)
              update_properties restore_frozen_state, propagation_node, target, propagation_data, ret
            end

            def update_properties restore_frozen_state, propagation_node, target, propagation_data, ret
              # Once we've propagated, attempt to tag the target if there is a tag(s) to be applied
              apply_tags(propagation_node, target)
              # Even though we skipped propagating tags from the source if they were included in untags, the target may
              # have already had some on it. Let's go ahead and remove them. In this order, untags takes precedent over
              # tags; but we control both and there should never be a propagator that has a tag in its untag.
              apply_untags(propagation_node, target)
              return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))

              properties.add_properties(propagation_node.properties)
              event_data = Contrast::Agent::Assess::Events::EventData.new(propagation_node,
                                                                          target,
                                                                          propagation_data.object,
                                                                          ret,
                                                                          propagation_data.args)
              properties.build_event(event_data)
              logger.trace('Propagation detected', node_id: propagation_node.id, target_id: target.__id__)
              restore_frozen_state ? ret : nil
            end

            # Find the propagation class from the given node, if one exists.
            #
            # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode] the node that governs a
            #   propagation event.
            # @return [Contrast::Agent::Assess::Policy::Propagator, nil]
            def find_propagation_class propagation_node
              unless (propagation_class = PROPAGATION_ACTIONS.fetch(propagation_node.action, nil))
                logger.warn('Unknown propagation action received. Unable to propagate.',
                            node_id: propagation_node.id,
                            action: propagation_node.action)
              end
              propagation_class
            end

            # We can handle frozen propagation iff we're allowed to, as determined by configuration, and the target of
            # the propagation is a return, as that's a replaceable value.
            #
            # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode] the node that governs a
            #   propagation event.
            # @return [Boolean]
            def can_handle_frozen? propagation_node
              ::Contrast::ASSESS.track_frozen_sources? &&
                  propagation_node.targets[0] == Contrast::Utils::ObjectShare::RETURN_KEY
            end
          end
        end
        # rubocop:enable Metrics/ModuleLength
      end
    end
  end
end