# 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/agent/assess/properties'
cs__scoped_require 'contrast/agent/assess/insulator'

module Contrast
  module Extension
    module Assess
      # This module is responsible for maintaining the data we need to
      # construct a trace event for the object in which it is included. Rather
      # than have this code all over the place, any class that wants to use
      # dataflow features should be sent
      # 'include Contrast::Extension::Assess::AssessExtension'
      module AssessExtension
        # Lazily build properties object. Only objects that have been tracked
        # will have the @_cs__properties, but all will respond to the
        # cs__properties method call. You should only call this method if you
        # either intend to start tracking an object or you have already checked
        # cs__tracked? and it is true.
        def cs__properties
          # If this object was tracked before being frozen, it'll have
          # mutable properties we need inside of the insulator @_cs__properties
          if cs__frozen?
            if instance_variable_defined?(:@_cs__properties)
              @_cs__properties.properties
            else
              Contrast::Agent::Assess::Insulator.generate_frozen.properties
            end
          else
            @_cs__properties ||= Contrast::Agent::Assess::Insulator.generate
            @_cs__properties.properties
          end
        end

        def cs__properties?
          instance_variable_defined?(:@_cs__properties)
        end

        # This is a way to check if we are already tracking an object without
        # adding tracking to it. If the object already has been tracked we will
        # return the tracking state of its properties. If the object hasn't
        # already been tracked we will return false without starting to track
        # it
        def cs__tracked?
          cs__properties? && cs__properties.tracked?
        end

        def cs__reset_properties
          return unless cs__properties?

          @_cs__properties = nil
        end

        # copy tags and info from object to self if object support methods
        # obj: the object from which to copy tags and events
        # shift: how far to shift the tags, negative moves left
        # skip_tags: array of tags to skip copying
        def cs__copy_from obj, shift = 0, skip_tags = nil
          return if obj.equal?(self)
          return unless Contrast::Utils::DuckUtils.quacks_to?(obj,
                                                              :cs__tracked?)
          return unless obj.cs__tracked?
          return if cs__properties == Contrast::Agent::Assess::Insulator.generate_frozen.properties

          # This was fun to find...
          # the clone and dup methods don't apply to instance variables in the
          # cloned/ duped thing, so the arrays in the properties were the same.
          # The most infinite of infinite loops ensued.
          # DO NOT TAKE THIS OUT!
          cs__reset_properties if obj.cs__properties == cs__properties

          obj.cs__properties.events.each do |event|
            cs__properties.events << event
          end

          obj.cs__properties.tag_keys.each do |key|
            next if skip_tags&.include?(key)

            new_tags = []
            value = obj.cs__properties.fetch_tag(key)
            value.each do |tag|
              new_tags << tag.copy_modified(shift)
            end
            existing = cs__properties.fetch_tag(key)
            if existing
              existing.concat(new_tags)
              Contrast::Utils::TagUtil.size_aware_merge(self, existing)
            else
              cs__properties.set_tags(key, new_tags)
            end
          end
        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
          return unless Contrast::Utils::DuckUtils.trackable?(ret)

          length = Contrast::Utils::StringUtils.ret_length(ret)
          return if length.zero?

          cs__splat_from_source(ret, length, source)
          cs__splat_from_ret(ret, length)
        end

        def cs__splat_from_source ret, ret_length, source
          splat_source = Contrast::Utils::DuckUtils.trackable?(source)
          splat_source &&= source.cs__tracked?
          return unless splat_source

          source.cs__properties.tag_keys.each do |key|
            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
              existing.drop(existing.length - 1)
              range = existing[0]
              range.repurpose(0, ret_length)
            else
              span = Contrast::Agent::Assess::AdjustedSpan.new(0, ret_length)
              ret.cs__properties.add_tag(key, span)
            end
          end
        end

        def cs__splat_from_ret ret, length
          return unless ret.cs__tracked?

          ret.cs__properties.tag_keys.each do |key|
            next unless key

            existing = ret.cs__properties.fetch_tag(key)
            next unless existing

            existing.each do |range|
              range.update_end(length) if range.end_idx > length
            end
          end
        end
      end
    end
  end
end