# 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'
cs__scoped_require 'contrast/core_extensions/module'
cs__scoped_require 'contrast/utils/object_share'

module Contrast
  module Agent
    module Assess
      # This is used to revert, or undo, the patches that we've placed into
      # modules. This is necessary for those cases where the original method
      # was supposed to be removed, but wasn't b/c we had renamed it -- looking
      # at you, FactoryBot
      class ClassReverter
        include Contrast::Components::Interface
        access_component :logging, :scope

        class << self
          def unpatch!
            Contrast::Agent::FeatureState.instance.uninstrument_namespaces.
                each { |mod_sym| revert_module mod_sym }
          end

          private

          def revert_module mod
            with_contrast_scope do
              revert_child_modules(mod, [])
            end
          rescue StandardError => e
            logger.error(e, "Unable to remove patches from the module #{ mod }")
          end

          def revert_child_modules mod, reverted_modules, parent_mod = nil
            return if parent_mod == mod
            return if reverted_modules.include?(mod)

            reverted_modules << mod

            immediate_constants = mod.cs__constants(false).collect! { |k| mod.cs__const_get(k) }
            immediate_constants.select! { |k| k.is_a?(Module) }
            immediate_constants.flatten!
            if immediate_constants.any?
              immediate_constants.each do |const|
                revert_aliases(const)
                revert_child_modules(const, reverted_modules, mod)
              end
            else
              revert_aliases(mod)
            end
          end

          # in order to fully uninstrument classes we must use true when getting the singleton/instance methods
          # that way if a class makes use of instance_eval for instance(defined on Kernel) we are able to revert this
          # specific instance of that method to use the default behavior while leaving the remainder of
          # objects using the patched behavior
          def revert_aliases clazz
            marker = Contrast::Utils::ObjectShare::CONTRAST_PATCHED_METHOD_START
            instance_methods = (clazz.instance_methods(true) + clazz.private_instance_methods(true)).select { |method| method.to_s.start_with?(marker) }
            singleton_methods = clazz.cs__singleton_class.instance_methods(true).select { |method| method.to_s.start_with?(marker) }
            instance_methods.each { |i_method| revert_alias(clazz, i_method, instance_methods) }
            singleton_methods.each { |s_method| revert_alias(clazz.cs__singleton_class, s_method, singleton_methods) }
          end

          def revert_alias clazz, current_method_name, methods
            original_method_name = clazz.instance_method(current_method_name).original_name

            is_private = clazz.private_method_defined?(original_method_name)

            # revert aliasing only for those methods currently defined on the original
            if is_private || clazz.method_defined?(original_method_name)
              clazz.send(:alias_method, original_method_name, current_method_name)
              clazz.send(:private, original_method_name) if is_private
            end
            clazz.send(:undef_method, current_method_name) if methods.include?(current_method_name)
          end
        end
      end
    end
  end
end