# 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/properties'
require 'contrast/agent/assess/finalizers/finalize'

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
        def cs__transfer_properties dup
          Contrast::Agent::Assess::Finalizers::Finalize::PROPERTIES_HASH[dup] ||= Contrast::Agent::Assess::Finalizers::Finalize::PROPERTIES_HASH[self].dup
        end

        # 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
          Contrast::Agent::Assess::Finalizers::Finalize::PROPERTIES_HASH[self] ||= Contrast::Agent::Assess::Properties.new
        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?
          !!Contrast::Agent::Assess::Finalizers::Finalize::PROPERTIES_HASH[self]&.tracked?
        end

        def cs__reset_properties
          Contrast::Agent::Assess::Finalizers::Finalize::PROPERTIES_HASH[self] = 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 unless cs__properties

          cs__adjust_duplicate(obj)

          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
              ret.cs__properties.add_tag(key, 0...ret_length)
            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

        private

        # Because of how our tracking works now, sometimes the Source and
        # Target are the same, but their IDs in our map will be different due
        # to PreShift duplication. To account for this, we have to ensure that
        # the Object we're copying from does not have the same Properties
        # that the Object we're copying to does. If they are the same, wipe the
        # Target so that the copy method can update events and ranges as
        # necessary.
        # DO NOT TAKE THIS OUT!
        def cs__adjust_duplicate obj
          cs__reset_properties if obj.cs__properties == cs__properties
          cs__reset_properties if obj.cs__properties.__id__ == cs__properties.dupped_from
          cs__reset_properties if obj.cs__properties.dupped_from == cs__properties.__id__
        end
      end
    end
  end
end