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

require 'monitor'

# base
require 'contrast/agent/patching/policy/patch_status'
require 'contrast/agent/patching/policy/method_policy'
require 'contrast/agent/patching/policy/module_policy'
require 'contrast/components/logger'
require 'contrast/components/scope'
require 'contrast/utils/class_util'
require 'contrast/utils/patching/policy/patcher_utils'

# assess
require 'contrast/agent/assess/policy/policy'
require 'contrast/agent/assess/policy/policy_scanner'
require 'contrast/agent/assess/policy/source_method'
require 'contrast/agent/assess/policy/trigger_method'

# deadzone
require 'contrast/agent/deadzone/policy/policy'

# inventory
require 'contrast/agent/inventory/policy/policy'

# protect
require 'contrast/agent/protect/policy/policy'

# patch
require 'contrast/agent/patching/policy/patch'

# after load patched
require 'contrast/agent/patching/policy/after_load_patcher'

module Contrast
  module Agent
    module Patching
      module Policy
        # This is how we patch into our customer's code. It provides a way to
        # track which classes we need to patch into and, once we've woven,
        # provides a map for which methods our renamed functions need to call
        # and how.
        module Patcher
          extend Contrast::Agent::Patching::Policy::AfterLoadPatcher
          extend Contrast::Utils::Patching::PatcherUtils
          extend Contrast::Components::Logger::InstanceMethods
          extend Contrast::Components::Scope::InstanceMethods

          class << self
            # Hook to install the Contrast changes needed to allow for the
            # instrumentation of the application - this only occurs once, during
            # startup to catchup on everything we didn't see get loaded
            def patch
              catchup_after_load_patches
              catchup_loaded_methods
            end

            # Hook to only monkeypatch Contrast. This will not trigger any
            # other functions, like rewriting or scanning. Exposed for those
            # situations, like ActiveRecord dynamically defining functions,
            # where only a subset of the Assess changes are needed.
            PATCH_MONITOR = Monitor.new

            # Iterate over and patch those Modules and Methods which were loaded before the Agent was enabled.
            def catchup_loaded_methods
              PATCH_MONITOR.synchronize do
                t = Contrast::Agent::Thread.new do
                  synchronized_catchup_loaded_methods
                end
                # aborting on exception makes exceptions propagate.
                t.abort_on_exception = true
                t.priority = Thread.current.priority + 1
                t.join
              end
            end

            private

            POLICIES = [
              Contrast::Agent::Assess::Policy::Policy,
              Contrast::Agent::Inventory::Policy::Policy,
              Contrast::Agent::Protect::Policy::Policy,
              Contrast::Agent::Deadzone::Policy::Policy
            ].cs__freeze

            def status_type
              Contrast::Agent::Patching::Policy::PatchStatus
            end

            # Find, and return, the constant with the given name. If none
            # exists, return nil.
            #
            # @param name [String] the constant whose value should be retrieved
            # @return [Object, nil]
            def patchable name
              return unless Contrast::Utils::ClassUtil.truly_defined?(name)

              Module.cs__const_get(name)
            end

            # @return [Array<String>] the names of all the Modules for which
            #   there patches in our policies
            def all_module_names
              @_all_module_names ||=
                POLICIES.each_with_object(Set.new) { |policy, set| set.merge(policy.instance.module_names) }
            end

            # Hook to only monkeypatch Contrast. This will not trigger any
            # other functions, like rewriting or scanning. This method should
            # only be invoked by the patch_methods method above in order to
            # ensure it is wrapped in a synchronize call
            def synchronized_catchup_loaded_methods
              logger.trace_with_time('Running patching') do
                patched = []
                all_module_names.each do |patchable_name|
                  next if ::Contrast::AGENT.skip_instrumentation?(patchable_name)

                  patchable_mod = patchable(patchable_name)
                  next unless patchable_mod

                  module_data = Contrast::Agent::ModuleData.new(patchable_mod, patchable_name)
                  patch_into_module(module_data)
                  patched << patchable_name if status_type.get_status(patchable_mod).patched?
                end
                all_module_names.subtract(patched)
              end
            end

            # Given the patchers that apply to this class that may apply, patch Contrast method calls into the methods
            # for which we have rules.
            #
            # @param module_data [Contrast::Agent::ModuleData] the module, and its name, that's being patched into
            # @param redo_patch [Boolean] a trigger to force patching regardless of the state of the
            #   Contrast::Agent::Patching::Policy::PatchStatus status on the Module
            def patch_into_module module_data, redo_patch = false
              status = Contrast::Agent::Patching::Policy::PatchStatus.get_status(module_data.mod)
              return if (status&.patched? || status&.patching?) && !redo_patch

              # Begin patching our sources into the given module. Any patcher that has the name of the module will be
              # evaluated for patching. Find all the patchers that apply to this class, sorted by type.
              module_policy = Contrast::Agent::Patching::Policy::ModulePolicy.create_module_policy(module_data.mod_name)

              # If there's nothing to match, then set that status and exit
              if module_policy.empty?
                status.no_patch!
                return
              end

              patch_methods(status, module_data, module_policy)
            rescue StandardError => e
              status&.failed_patch!
              logger.warn('Patching failed', e, module: module_data.mod_name)
            ensure
              logger.trace('Patching complete',
                           module: module_data.mod_name,
                           result:
                             Contrast::Agent::Patching::Policy::PatchStatus.get_status(module_data.mod).patch_status)
            end

            def patch_methods status, module_data, module_policy
              status.patching!
              num_applied_patches = patch_into_instance_methods(module_data, module_policy)
              num_applied_patches += patch_into_singleton_methods(module_data, module_policy)

              if adjust_for_prepend(module_data) || module_policy.num_expected_patches == num_applied_patches
                status.patched!
              else
                status.partial_patch!
              end
            end

            # Get all of the instance methods on the given module, excluding
            # those from super classes. this list will always include the
            # initialize method
            #
            # @param mod [Module] The module from which to retrieve instance
            #   methods.
            # @param private [Boolean] Indicate if the query should include
            #   private, as well as public, instance methods
            # @return [Array<Symbol>]
            def all_instance_methods mod, private = false
              instance_methods = mod.instance_methods(false)
              # C magic rb_define_global_function creates private instance
              # methods. We need to instrument those dudes
              # https://www.thecodingforums.com/threads/object-private_instance_methods-kernel-singleton_methods.843013/
              instance_methods.concat(mod.private_instance_methods(false)) if private
              # initialize is an exception to the inheritance rule.
              instance_methods << :initialize unless instance_methods.include?(:initialize)
              # Object owns these, but we don't want to patch every :dup or
              # :clone, just the ones on String, so we'll special case this.
              # This breaks the norm, where we usually instrument the lowest
              # method possible.
              if mod == String
                instance_methods << :clone
                instance_methods << :dup
              end
              instance_methods
            end

            # Patch into the Instance Methods, including private, of the given Module that match the ModulePolicy
            # provided.
            #
            # @param module_data [Contrast::Agent::ModuleData] the module, and its name, that's being patched into
            # @param module_policy [Contrast::Agent::Patching::Policy::ModulePolicy] All the patchers that apply to
            #   this module, sorted by type.
            def patch_into_instance_methods module_data, module_policy
              mod = module_data.mod
              methods = all_instance_methods(mod, true)
              methods.delete(:initialize) if mod.to_s.starts_with?('RSpec') && mod.to_s.include?('Matchers')
              patch_into_methods(mod, methods, module_policy, true)
            end

            # Patch into the Singleton Methods of the given Module that match the ModulePolicy provided.
            #
            # @param module_data [Contrast::Agent::ModuleData] the module, and its name, that's being patched into
            # @param module_policy [Contrast::Agent::Patching::Policy::ModulePolicy] All the patchers that apply to
            #   this module, sorted by type.
            def patch_into_singleton_methods module_data, module_policy
              mod = module_data.mod
              methods = mod.singleton_methods(false)
              patch_into_methods(mod, methods, module_policy, false)
            end

            # We've found the patchers that apply to this class (or module). Now we'll filter on the given method.
            #
            # @param mod [Module] The module from which to retrieve instance methods.
            # @param methods [Array<Symbol>] The names of all the methods in in this module
            # @param module_policy [Contrast::Agent::Patching::Policy::ModulePolicy] All the patchers that apply to
            #   this class, sorted by type.
            # @param is_instance_method [Boolean] Indicates if these methods are the instance or singleton methods for
            #   this module.
            # @return [Integer] number of methods patched.
            def patch_into_methods mod, methods, module_policy, is_instance_method
              count = 0
              methods.each do |method|
                method_policy = Contrast::Agent::Patching::Policy::MethodPolicy.build_method_policy(method,
                                                                                                    module_policy,
                                                                                                    is_instance_method)
                next if method_policy.empty?

                count += 1 if patch_method(mod, methods, method_policy)
              end
              count
            end

            # CGI falls pray to the weirdness that is the history of a
            # previously included Module being prepended - the class that has
            # already included the Module does not learn about the prepend and
            # it has to be reapplied.
            # TODO: RUBY-620 should remove the need for this
            #
            # @param module_data [Contrast::Agent::ModuleData] the module, and its name, that's being patched into
            # @return [Boolean] if an adjustment was made or not
            def adjust_for_prepend module_data
              return false unless module_data.mod.cs__name == 'CGI::Util'

              CGI.include(CGI::Util)
              CGI.extend(CGI::Util)
              true
            end
          end
        end
      end
    end
  end
end

# core extensions
require 'contrast/extension/module'
require 'contrast/extension/assess'
require 'contrast/extension/inventory'
require 'contrast/extension/protect'

require 'cs__contrast_patch/cs__contrast_patch'