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

require 'contrast/components/interface'
require 'contrast/utils/duck_utils'

module Contrast
  module Agent
    module Assess
      module Policy
        module Propagator
          # This class is specifically for String#(g)sub propagation
          #
          # Disclaimer: there may be a better way, but we're
          # in a 'get it work' state. hopefully, we'll be in
          # a 'get it right' state soon.
          class Substitution
            include Contrast::Components::Interface
            access_component :logging

            CAPTURE_GROUP_REGEXP = /\\[[:digit:]]/.cs__freeze
            CAPTURE_NAME_REGEXP = /\\k<[[:alpha:]]/.cs__freeze

            class << self
              # gsub is hard. there are four versions of this method
              # 1) pattern, replacement (my fav)
              # 2) pattern, hash        (not bad)
              # 3) pattern, block       (are you kidding me?)
              # 4) pattern              (plz no)
              #
              # In addition, it requires things from $~ & $1-9, which
              # are method scoped. Rather than fight that, we'll
              # call gsub from C land using a CUSTOM patch.
              def gsub_tagger patcher, preshift, ret, block
                substitution_tagger(patcher, preshift, ret, !block.nil?)
              end

              def sub_tagger patcher, preshift, ret, block
                substitution_tagger(patcher, preshift, ret, !block.nil?, false)
              end

              private

              def substitution_tagger patcher, preshift, ret, block, global = true
                return ret unless ret

                begin
                  source = preshift.object
                  self_tracked = Contrast::Agent::Assess::Tracker.tracked?(source)
                  args = preshift.args[1]
                  incoming_tracked = args && determine_tracked(args)
                  return ret unless self_tracked || incoming_tracked

                  if block
                    block_sub(self_tracked, source, ret)
                  elsif args.is_a?(String)
                    string_sub(self_tracked, preshift, ret, args, incoming_tracked, global)
                  elsif args.is_a?(Hash)
                    hash_sub(self_tracked, source, ret)
                  else # Enumerator, only for gsub
                    pattern_gsub(preshift, ret)
                  end
                  string_build_event(patcher, preshift, ret)
                rescue StandardError => e
                  logger.error('Unable to apply gsub propagator', e)
                end
                ret
              end

              def determine_tracked args
                # if there's no arg, it can't be tracked
                return false unless args

                # if it's a string, just ask if it's tracked
                case args
                when String
                  Contrast::Agent::Assess::Tracker.tracked?(args)
                  # if it's a hash, ask if it has a tracked string
                when Hash
                  args.values.any? { |value| value.is_a?(String) && Contrast::Agent::Assess::Tracker.tracked?(value) }
                  # this should never happen
                else
                  false
                end
              end

              def string_sub self_tracked, preshift, ret, incoming, incoming_tracked, global
                properties = Contrast::Agent::Assess::Tracker.properties(ret)
                return unless properties

                pattern = preshift.args[0]
                source = preshift.object

                # We can't efficiently find the places that things were
                # copied from regexp / captures. Trading accuracy for
                # performance
                if incoming.match?(CAPTURE_GROUP_REGEXP) || incoming.match?(CAPTURE_NAME_REGEXP)
                  properties.splat_from(source, ret) if self_tracked
                  return
                end

                # if it's just a straight insert, that we can do
                # Copy the tags from us to the return
                properties.copy_from(source, ret)
                # Figure out where inserts occurred
                last_idx = 0
                ranges = []
                # For each insert, move the tag ranges
                while last_idx
                  idx = source.index(pattern, last_idx)
                  break unless idx

                  last_idx = idx ? idx + 1 : nil
                  start_index = idx
                  end_index = idx + incoming.length
                  ranges << (start_index...end_index)
                  break unless global
                end
                properties.delete_tags_at_ranges(ranges)
                properties.shift_tags(ranges)
                return unless incoming_tracked

                incoming_properties = Contrast::Agent::Assess::Tracker.properties(incoming)
                tags = incoming_properties.tag_keys
                ranges.each do |range|
                  tags.each do |tag|
                    properties.add_tag(tag, range)
                  end
                end
              end

              def block_sub self_tracked, source, ret
                properties = Contrast::Agent::Assess::Tracker.properties(ret)
                properties&.splat_from(source, ret) if self_tracked
              end

              def hash_sub self_tracked, source, ret
                properties = Contrast::Agent::Assess::Tracker.properties(ret)
                properties&.splat_from(source, ret) if self_tracked
              end

              def pattern_gsub preshift, ret
                properties = Contrast::Agent::Assess::Tracker.properties(ret)
                return unless properties

                source = preshift.object
                source_properties = Contrast::Agent::Assess::Tracker.properties(source)
                return unless source_properties

                source_properties.tag_keys.each do |key|
                  properties.add_tag(key, 0...1)
                end
              end

              def string_build_event patcher, preshift, ret
                return unless Contrast::Agent::Assess::Tracker.tracked?(ret)

                properties = Contrast::Agent::Assess::Tracker.properties(ret)
                args = preshift.args
                if args.length > 1
                  arg = args[1]
                  arg_properties = Contrast::Agent::Assess::Tracker.properties(arg)
                  arg_properties&.events&.each do |event|
                    properties.events << event
                  end
                end
                properties.build_event(
                    patcher,
                    ret,
                    preshift.object,
                    ret,
                    args,
                    2)
              end
            end
          end
        end
      end
    end
  end
end