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

cs__scoped_require 'contrast/components/interface'

# Monkeypatch Ruby Regexp with Contrast Security tagging support
class Regexp
  include Contrast::Components::Interface
  include Contrast::CoreExtensions::Assess::AssessExtension

  access_component :analysis, :logging, :scope

  def cs__regexp_tagger info_hash
    return unless ASSESS.enabled?

    # Because we have a special case for this propagation,
    # it falls out of regular scoping. As such, any patch to the `=~` method
    # that goes through normal channels, like that for the redos rule,
    # will force this to be in a scope of 1 (instead of the normal 0).
    # As such, a scope of 1 here indicates that,
    # so we know that we're in the top level call for this method.
    # normal patch [-alias-]> special case patch [-alias-]> original method
    # TODO:  RUBY-686
    return if scope_for_current_ec.instance_variable_get(:@contrast_scope) > 1

    with_contrast_scope do
      result = info_hash[:result]
      return unless result

      string = info_hash[:string] || nil
      return unless string

      regexp = info_hash[:regexp]
      return unless (Contrast::Utils::DuckUtils.quacks_to?(string, :cs__tracked?) && string.cs__tracked?) ||
          (Contrast::Utils::DuckUtils.quacks_to?(regexp, :cs__tracked?) && regexp.cs__tracked?)

      Regexp.cs__splat_tags(info_hash[:back_ref], string)
      Regexp.cs__build_event(ARRAY_NODE, info_hash[:back_ref], result, [string])
    end
  end

  REGEXP_EQUAL_SQUIGGLE_HASH = {
      'id' => 'regexp_100',
      'class_name' => 'Regexp',
      'instance_method' => true,
      'method_visibility' => 'public',
      'method_name' => '=~',
      'source' => 'P0',
      'target' => 'R',
      'action' => 'KEEP'
  }.cs__freeze
  ARRAY_NODE = Contrast::Agent::Assess::Policy::PropagationNode.new(REGEXP_EQUAL_SQUIGGLE_HASH)
  private_constant :REGEXP_EQUAL_SQUIGGLE_HASH
  private_constant :ARRAY_NODE

  REGEXP_MATCH_HASH = {
      'id' => 'regexp_101',
      'class_name' => 'Regexp',
      'instance_method' => true,
      'method_visibility' => 'public',
      'method_name' => 'match',
      'source' => 'P0',
      'target' => 'R',
      'action' => 'SPLAT'
  }.cs__freeze
  MATCH_NODE = Contrast::Agent::Assess::Policy::PropagationNode.new(REGEXP_MATCH_HASH)
  private_constant :REGEXP_MATCH_HASH
  private_constant :MATCH_NODE

  POST_HASH = {
      'class_name' => 'MatchData',
      'instance_method' => true,
      'method_visibility' => 'public',
      'method_name' => 'post_match',
      'source' => 'O',
      'target' => 'R',
      'action' => 'REMOVE'
  }.cs__freeze
  POST_WEAVER = Contrast::Agent::Assess::Policy::PropagationNode.new(POST_HASH)
  private_constant :POST_HASH
  private_constant :POST_WEAVER

  PRE_HASH = {
      'class_name' => 'MatchData',
      'instance_method' => true,
      'method_visibility' => 'public',
      'method_name' => 'pre_match',
      'source' => 'O',
      'target' => 'R',
      'action' => 'SPLAT'
  }.cs__freeze
  PRE_WEAVER = Contrast::Agent::Assess::Policy::PropagationNode.new(PRE_HASH)
  private_constant :PRE_HASH
  private_constant :PRE_WEAVER

  LAST_PAREN_HASH = {
      'class_name' => 'MatchData',
      'instance_method' => true,
      'method_visibility' => 'public',
      'method_name' => 'last_paren',
      'source' => 'O',
      'target' => 'R',
      'action' => 'SPLAT'
  }.cs__freeze
  LAST_PAREN_WEAVER = Contrast::Agent::Assess::Policy::PropagationNode.new(LAST_PAREN_HASH)
  private_constant :LAST_PAREN_HASH
  private_constant :LAST_PAREN_WEAVER

  NTH_HASH = {
      'class_name' => 'MatchData',
      'instance_method' => true,
      'method_visibility' => 'public',
      'method_name' => 'nth_match',
      'source' => 'O',
      'target' => 'R',
      'action' => 'SPLAT'
  }.cs__freeze
  NTH_WEAVER = Contrast::Agent::Assess::Policy::PropagationNode.new(NTH_HASH)
  private_constant :NTH_HASH
  private_constant :NTH_WEAVER

  class << self
    def track_rb_pre_match backref, target
      track_rb_c(PRE_WEAVER, backref, target)
    end

    def track_rb_post_match backref, target
      track_rb_c(POST_WEAVER, backref, target)
    end

    def track_rb_reg_match_last backref, target
      track_rb_c(LAST_PAREN_WEAVER, backref, target)
    end

    def track_rb_n_match backref, target
      track_rb_c(NTH_WEAVER, backref, target)
    end

    # Some propagation occurred, but we're not sure what the
    # exact transformation was. To be safe, we just explode
    # all the tags from the source to the return.
    #
    # If the return already had that tag, the existing tag
    # range is recycled to save us an object.
    def cs__splat_tags ret, source = self
      source.cs__properties.tag_keys.each do |key|
        length = Contrast::Utils::StringUtils.ret_length(ret)
        existing = ret.cs__properties.fetch_tag(key)
        # if the tag already exists, drop all but the first range
        # then change that range to cover the entire return
        if existing&.any?
          existing.drop(existing.length - 1)
          range = existing[0]
          range.repurpose(0, length)
        else
          span = Contrast::Agent::Assess::AdjustedSpan.new(0, length)
          ret.cs__properties.add_tag(key, span)
        end
      end
    end

    def cs__build_event node, target, ret, args
      return unless Contrast::Utils::DuckUtils.quacks_to?(target, :cs__tracked?)
      return unless target.cs__tracked?

      target.cs__properties.build_event(
          node,
          target,
          self,
          ret,
          args)
    end

    def instrument_regexp_track
      @_instrument_regexp_variables ||= begin
          cs__scoped_require 'cs__assess_regexp_track/cs__assess_regexp_track'
          true
        end
    rescue StandardError => e
      logger.error(e, 'Error loading regexp track patch')
      false
    end

    private

    def track_rb_c weaver, backref, target
      return target unless ASSESS.enabled?
      return target unless backref&.respond_to?(:cs__properties)
      return target unless target.is_a?(String) && !target.empty?

      with_contrast_scope do
        cs__splat_tags(target, backref)
        target.cs__properties.build_event(
            weaver,
            target,
            self,
            target,
            [])
        target
      end
    end
  end
end

cs__scoped_require 'cs__assess_regexp/cs__assess_regexp'
Regexp.instrument_regexp_track