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

require 'contrast/components/logger'
require 'contrast/utils/lru_cache'

module Contrast
  module Agent
    module Assess
      # In order to properly shift tags to account for the changes this method
      # caused, we'll need to store the state before the change occurred.
      class PreShift
        include Contrast::Components::Logger::InstanceMethods
        extend Contrast::Components::Logger::InstanceMethods

        @lru_cache = Contrast::Utils::LRUCache.new
        UNDUPLICABLE_MODULES = [
          Enumerator # dup'ing results in 'can't copy execution context'
        ].cs__freeze

        attr_accessor :object, :object_length, :args, :arg_lengths

        class << self
          # Save the state before we do the propagation. This is one of the
          # few things that we'll call BEFORE calling the original method.
          # Unfortunately, we may waste some time here if the method then
          # throws an exception, but hey, it threw an exception. These few
          # extra objects aren't the real problem in that case.
          #
          # @param propagation_node [Contrast::Agent::Assess::Policy::PropagationNode]
          #   The node responsible for the propagation action required by this
          #   method.
          # @param object [Object] the object on which the method is to be
          #   called.
          # @param args [Array<Object>] the arguments to be passed to the
          #   method.
          # @return [Contrast::Agent::Assess::PreShift, nil] a holder saving
          #   the state of the object and arguments just prior to the method
          #   being called or nil if one is not required.
          def build_preshift propagation_node, object, args
            return unless propagation_node
            return unless ::Contrast::ASSESS.enabled?

            initializing = propagation_node.method_name == :initialize
            return if unsafe_io_object?(object, initializing)

            needs_object = propagation_node.needs_object?
            needs_args = propagation_node.needs_args?

            preshift = Contrast::Agent::Assess::PreShift.new
            append_object_details(preshift, initializing, object) if needs_object
            append_arg_details(preshift, args) if needs_args

            preshift
          rescue StandardError => e
            logger.error('Unable to build preshift for method.', e)
            nil
          end

          private

          def unsafe_io_object? object, initializing
            Contrast::Utils::DuckUtils.closable_io?(object) && !initializing && object.closed?
          end

          def can_dup? initializing, object
            return false if initializing

            check = object.is_a?(Module) ? object : object.cs__class
            !UNDUPLICABLE_MODULES.include?(check)
          end

          def append_object_details preshift, initializing, object
            can = can_dup?(initializing, object)
            preshift.object = if @lru_cache.key?(object.__id__) && !Contrast::Agent::Assess::Tracker.tracked?(object)
                                @lru_cache[object.__id__]
                              else
                                can ? object.dup : object
                              end
            preshift.object_length = if Contrast::Utils::DuckUtils.quacks_to?(preshift.object, :length)
                                       object.length
                                     else
                                       0
                                     end

            return unless can

            return unless Contrast::Agent::Assess::Tracker.tracked?(object)

            Contrast::Agent::Assess::Tracker.copy(object, preshift.object)
            @lru_cache[object.__id__] = object
          end

          def append_arg_details preshift, args
            args_length = args.length
            preshift.args = Array.new(args_length)
            preshift.arg_lengths = Array.new(args_length)
            idx = 0
            while idx < args_length
              or_arg = args[idx]
              p_arg = if @lru_cache.key?(or_arg.__id__)
                        @lru_cache[or_arg.__id__]
                      else
                        can_dup?(false, or_arg) ? or_arg.dup : or_arg
                      end
              preshift.args[idx] = p_arg
              preshift.arg_lengths[idx] = Contrast::Utils::DuckUtils.quacks_to?(p_arg, :length) ? p_arg.length : 0
              idx += 1
              next if p_arg.__id__ == or_arg.__id__

              Contrast::Agent::Assess::Tracker.copy(or_arg, p_arg)
              @lru_cache[p_arg.__id__] = p_arg
            end
          end
        end
      end
    end
  end
end