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

require 'contrast/agent/assess/policy/preshift'
require 'contrast/components/agent'
require 'contrast/components/logger'
require 'contrast/components/scope'
require 'contrast/utils/thread_tracker'
require 'contrast/utils/assess/split_utils'
require 'contrast/agent/assess/events/event_data'

module Contrast
  module Agent
    module Assess
      module Policy
        module Propagator
          # This class is specifically for String#split & String#grapheme_clusters propagation
          # it propagates tag ranges from a string to elements within an untracked array
          class Split < Contrast::Agent::Assess::Policy::Propagator::Base
            extend Contrast::Components::Scope::InstanceMethods
            extend Contrast::Components::Logger::InstanceMethods
            extend Contrast::Utils::Assess::SplitUtils

            SPLIT_TRACKER = Contrast::Utils::ThreadTracker.new

            class << self
              # Propagate taint from a source as it is split into composite sections.
              #
              # @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 [Array, String] the target to which to propagate.
              # @return [nil] so as not to risk changing the result of the propagation.
              def propagate propagation_node, preshift, target
                return unless target.is_a?(Array) # apply_post_patch is called, but split with block returns a string.

                logger.trace('Propagation detected', node_id: propagation_node.id, target_id: target.__id__)

                source = find_source(propagation_node.sources[0], preshift)
                return unless (source_properties = Contrast::Agent::Assess::Tracker.properties(source))

                update_element_properties propagation_node, target, preshift, source_properties
                nil
              end

              # Context for block split execution.
              #
              # @param string [String] the String on which split is invoked.
              # @param args [Array<Object>] the arguments passed to the original split call.
              def wrap_split string, args
                # String#split start. Build context and yield.
                begin
                  enter_split_scope!
                  save_split_index!
                  save_split_value!(string, args)
                rescue Exception => e # rubocop:disable Lint/RescueException
                  logger.warn('Unable to record split context', e)
                end

                yield
              ensure
                # String#split exit. Remove propagation context.
                begin
                  exit_split_scope!
                  unless in_split_scope?
                    SPLIT_TRACKER.delete(:split_index)
                    SPLIT_TRACKER.delete(:split_value)
                  end
                rescue StandardError => e
                  logger.warn('Unable to remove split context', e)
                end
              end

              # This method is called whenever an rb_yield is called.
              # We need to leave it as soon as possible with as little work as possible.
              #
              # @param target [String] the entity being passed to the yield block
              def propagate_yield target
                return unless (source = SPLIT_TRACKER.get(:split_value)&.fetch(split_scope_depth))
                return unless (index = SPLIT_TRACKER.get(:split_index)&.fetch(split_scope_depth))
                return unless (properties = Contrast::Agent::Assess::Tracker.properties!(target))

                true_source = source[index]
                properties.copy_from(true_source, target)
              rescue StandardError => e
                logger.warn('Unable to track within split context', e)
              ensure
                if in_split_scope? && index
                  idx = SPLIT_TRACKER.get(:split_index)
                  idx[split_scope_depth] = index + 1 if defined?(idx) && idx.is_a?(Array)
                end
              end

              # Load patch.
              def instrument_string_split
                @_instrument_string_split ||= begin
                  if ::Contrast::AGENT.patch_yield? && Funchook.available?
                    require 'cs__assess_yield_track/cs__assess_yield_track'
                  end
                  true
                rescue StandardError => e
                  logger.error('Error loading split rb_yield patch', e)
                  false
                end
              end

              private

              # grapheme_clusters break the string apart based on each "user-perceived" character. Otherwise, the
              # default for String#split is to use a single whitespace.
              #
              # @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.
              def find_separator_length propagation_node, preshift
                return 0 if propagation_node.method_name == :grapheme_clusters

                preshift&.args&.first&.to_s&.length || $FIELD_SEPARATOR&.to_s&.length || 1
              end

              # Save index of the current split object.
              # Create index tracking array as needed.
              def save_split_index!
                unless (split_index = SPLIT_TRACKER.get(:split_index))
                  split_index = []
                  SPLIT_TRACKER.set(:split_index, split_index)
                end
                # save the index to the ThreadLocal; not useless.
                split_index[split_scope_depth] = 0 # rubocop:disable Lint/UselessSetterCall
              end

              # Save value of the current split object.
              # Create value tracking array as needed.
              def save_split_value! string, args
                preshift = Contrast::Agent::Assess::PreShift.build_preshift(split_node, string, args)
                target = string.split
                propagate(split_node, preshift, target)
                unless (split_value = SPLIT_TRACKER.get(:split_value))
                  split_value = []
                  SPLIT_TRACKER.set(:split_value, split_value)
                end
                # Save the target to the ThreadLocal; not useless.
                split_value[split_scope_depth] = target # rubocop:disable Lint/UselessSetterCall
              end

              def update_element_properties propagation_node, target, preshift, source_properties
                separator_length = find_separator_length(propagation_node, preshift)
                current_index = 0
                target.each do |target_elem|
                  next unless (elem_properties = Contrast::Agent::Assess::Tracker.properties!(target_elem))

                  # Get tags for element from source by element range.
                  range = current_index...(current_index + target_elem.length)
                  tags = source_properties.tags_at_range(range)

                  # Set element properties accordingly.
                  elem_properties.clear_tags
                  tags.each_pair { |key, value| elem_properties.set_tags(key, value) }
                  event_data = Contrast::Agent::Assess::Events::EventData.new(propagation_node,
                                                                              target_elem,
                                                                              preshift.object,
                                                                              target,
                                                                              preshift.args)
                  elem_properties.build_event(event_data, 0)
                  elem_properties.add_properties(propagation_node.properties)

                  current_index = range.end + separator_length
                end
              end
            end
          end
        end
      end
    end
  end
end

# Special class to handle String#split in which, when given a block, propagates each split piece directly.
class String
  alias_method :cs__patched_string_split_special, :split

  # Override of the the standard split method to handle the direct yield case.
  #
  # Note: because this patch is applied before our standard propagation, this call is wrapped in it. As such, any call
  # here happens in scope, so there is no need to manage it on our own.
  def split *args, &block
    if block
      Contrast::Agent::Assess::Policy::Propagator::Split.wrap_split(self, args) do
        cs__patched_string_split_special(*args, &block)
      end
    else
      cs__patched_string_split_special(*args, &block)
    end
  end
end

Contrast::Agent::Assess::Policy::Propagator::Split.instrument_string_split