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

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
            include Contrast::Components::Interface

            access_component :agent, :logging

            SPLIT_TRACKER = Contrast::Utils::ThreadTracker.new
            class << self
              # Propagate taint from a source as it is split into composite
              # sections. This method MUST return nil, otherwise it risks
              # changing the result of 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 [Array, String] the target to which to propagate.
              # @return [nil]
              def propagate propagation_node, preshift, target
                logger.trace('Propagation detected',
                             node_id: propagation_node.id,
                             target_id: target.__id__)
                unless target.is_a?(Array)
                  Contrast::Agent::Assess::Policy::Propagator::Keep.propagate(propagation_node, preshift, target)
                  return
                end

                source = find_source(propagation_node.sources[0], preshift)

                separator_length = if propagation_node.method_name == :grapheme_clusters
                                     # grapheme_clusters break the string apart based on each "user-perceived" character
                                     0
                                   else
                                     # The default for String#split is to use a single whitespace
                                     preshift&.args&.first&.to_s&.length || $FIELD_SEPARATOR&.to_s&.length || 1
                                   end

                current_index = 0
                target.each do |elem|
                  elem_length = elem.length
                  range = current_index...(current_index + elem_length)
                  tags = source.cs__properties.tags_at_range(range)

                  elem.cs__properties.clear_tags
                  tags.each_pair do |key, value|
                    elem.cs__properties.set_tags(key, value)
                  end
                  source.cs__properties.events.each do |event|
                    elem.cs__properties.add_event(event)
                  end
                  elem.cs__properties.build_event(propagation_node, elem, preshift.object, target, preshift.args, 0)
                  elem.cs__properties.add_properties(propagation_node.properties)
                  current_index = current_index + elem_length + separator_length
                end
                nil
              end

              # Marks the point in which the String#split method is called.
              # Responsible for building the context required to propagate when
              # the results of #split are yielded directly to a block
              #
              # @param string [String] the String on which split is invoked
              # @param args [Array<Object>] the arguments passed to the
              #   original split call
              def begin_split string, args
                save_split_depth!
                depth = SPLIT_TRACKER.get(:split_depth)
                save_split_index!(depth)
                save_split_value!(depth, string, args)
              rescue Exception => e # rubocop:disable Lint/RescueException
                # don't let our errors propagate and disable String#split for
                # this since we're in an error state
                logger.warn('Unable to record split context', e)
                end_split
              end

              # Marks the point in which the String#split method is exited.
              # Responsible for removing the context required to propagate when
              # the results of #split are yielded directly to a block
              def end_split
                depth = SPLIT_TRACKER.get(:split_depth)
                return unless depth

                depth -= 1
                if depth.negative?
                  SPLIT_TRACKER.delete(:split_depth)
                  SPLIT_TRACKER.delete(:split_index)
                  SPLIT_TRACKER.delete(:split_value)
                else
                  SPLIT_TRACKER.set(:split_depth, depth)
                end
              rescue StandardError => e
                logger.warn('Unable to remove split context', e)
              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
                depth, index = nil

                depth = SPLIT_TRACKER.get(:split_depth)
                return unless depth

                source = SPLIT_TRACKER.get(:split_value)&.fetch(depth)
                return unless source

                index = SPLIT_TRACKER.get(:split_index)&.fetch(depth)
                return unless index

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

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

              private

              def save_split_depth!
                depth = SPLIT_TRACKER.get(:split_depth)
                if depth
                  depth += 1
                  SPLIT_TRACKER.set(:split_depth, depth)
                else
                  SPLIT_TRACKER.set(:split_depth, 0)
                end
              end

              def save_split_index! depth
                split_index = SPLIT_TRACKER.get(:split_index)
                unless split_index
                  split_index = []
                  SPLIT_TRACKER.set(:split_index, split_index)
                end
                # save the index to the ThreadLocal; not useless
                split_index[depth] = 0 # rubocop:disable Lint/UselessSetterCall
              end

              def save_split_value! depth, string, args
                preshift = Contrast::Agent::Assess::PreShift.build_preshift(split_node, string, args)
                target = string.split
                propagate(split_node, preshift, target)
                split_value = SPLIT_TRACKER.get(:split_value)
                unless split_value
                  split_value = []
                  SPLIT_TRACKER.set(:split_value, split_value)
                end
                # save the target to the ThreadLocal; not useless
                split_value[depth] = target # rubocop:disable Lint/UselessSetterCall
              end

              # Quick hook to the String#split propagation node in our Assess
              # policy
              #
              # @return [Contrast::Agent::Assess::Policy::PropagationNode]
              #   String#split node
              def split_node
                @_split_node ||= begin
                                   Contrast::Agent::Assess::Policy::Policy.instance.propagators.find do |node|
                                     node.class_name == 'String' && node.method_name == :split && node.instance_method?
                                   end
                                 end
              end
            end
          end
        end
      end
    end
  end
end

if RUBY_VERSION >= '2.6.0'
  # Special class to handle String#split in 2.6 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 2.6 direct
    # yield case.
    #
    # Note: because this patch is applied before our standard propagation, this
    # call 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.begin_split(self, args)
        begin
          cs__patched_string_split_special(*args, &block)
        ensure
          Contrast::Agent::Assess::Policy::Propagator::Split.end_split
        end
      else
        cs__patched_string_split_special(*args, &block)
      end
    end
  end

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