# Copyright (c) 2020 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/interface'
require 'contrast/utils/class_util'

# assess
require 'contrast/agent/assess/policy/policy'
require 'contrast/agent/assess/policy/policy_scanner'
require 'contrast/agent/assess/policy/rewriter_patch'
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

          include Contrast::Components::Interface
          access_component :agent, :analysis, :logging, :scope

          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
              patch_methods
              Contrast::Agent::Assess::Policy::RewriterPatch.rewrite_interpolations
            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

            def patch_methods
              PATCH_MONITOR.synchronize do
                t = Contrast::Agent::Thread.new do
                  synchronized_patch_methods
                end
                # aborting on exception makes exceptions propagate.
                t.abort_on_exception = true
                t.priority = Thread.current.priority + 1
                t.join
              end
            end

            # This method is called by TracePointHook to instrument a specific class during a require
            # or eval of dynamic class definition
            def patch_specific_module mod
              mod_name = mod.cs__name
              return unless Contrast::Utils::ClassUtil.truly_defined?(mod_name)
              return if AGENT.skip_instrumentation?(mod_name)

              load_patches_for_module(mod_name)

              return unless all_module_names.any? { |name| name == mod_name }

              module_data = Contrast::Agent::ModuleData.new(mod, mod_name)
              patch_into_module(module_data)
              all_module_names.delete(mod_name) if status_type.get_status(mod).patched?
            rescue StandardError => e
              logger.error('Unable to patch module', e, module: mod_name)
            end

            # We did it, team. We found a patcher(s) that applies to the given
            # class (or module) and the given method. Time to do some tracking.
            #
            # @param mod [Module] the module in which the patch should be
            #   placed.
            # @param methods [Array(Symbol)] all the instance or singleton
            #   methods in this clazz.
            # @param method_policy [Contrast::Agent::Patching::Policy::MethodPolicy]
            #   the policy that applies to the given method_name.
            # @return [Boolean] if patched, either by this invocation or a
            #   previous, or not
            def patch_method mod, methods, method_policy
              return false unless methods&.any? # don't even build the name if there are no methods

              if Contrast::Utils::ClassUtil.prepended_method?(mod, method_policy)
                Contrast::Agent::Patching::Policy::Patch.instrument_with_prepend(mod, method_policy)
              else
                Contrast::Agent::Patching::Policy::Patch.instrument_with_alias(mod, methods, method_policy)
              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_patch_methods
              logger.trace_with_time('Running patching') do
                patched = []
                all_module_names.each do |patchable_name|
                  next if 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 = status_type.get_status(module_data.mod)
              return if (status&.patched? || status&.patching?) && !redo_patch

              # Begin patching our sources into the given clazz (or module)
              # Any patcher that has the name of the clazz 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.name)

              clazz = module_data.mod

              status.patching!
              patched = include_module(module_data)

              counts = 0
              # Monkey patch any methods in this class that have matching nodes in the policy
              unless module_policy.empty?
                instance_methods = all_instance_methods(clazz, true)
                singleton_methods = clazz.singleton_methods(false)
                counts += patch_into_methods(clazz, instance_methods, module_policy, true)
                counts += patch_into_methods(clazz, singleton_methods, module_policy, false)
                counts = module_policy.num_expected_patches if adjust_for_prepend(clazz)
                patched = true
              end

              if patched
                if module_policy.num_expected_patches == counts
                  status.patched!
                else
                  status.partial_patch!
                end
              else
                status.no_patch!
              end
            rescue StandardError => e
              status ||= status_type.get_status(module_data.mod)
              status.failed_patch!
              logger.warn('Patching failed', e, module: module_data.name)
            ensure
              logger.trace('Patching complete',
                           module: module_data.name,
                           result: Contrast::Agent::Patching::Policy::PatchStatus.get_status(
                               module_data.mod).patch_status)
            end

            # Includes the given module with the
            # Contrast::Extension::Assess::AssessExtension
            # @param module_data [Contrast::Agent::ModuleData] the module, and
            #   its name, that's being patched into
            def include_module module_data
              return false unless Contrast::Agent::Assess::Policy::Policy.instance.tracked_classes.include?(module_data.name)

              module_data.mod.send(:include, Contrast::Extension::Assess::AssessExtension)
              true
            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

            # 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?

                patched = patch_method(mod, methods, method_policy)
                count += 1 if patched
              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 mod[Module] the Module for which a prepend action needs to
            #   be accounted
            # @return [Boolean] if an adjustment was made or not
            def adjust_for_prepend mod
              return false unless 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 'contrast/extension/protect/kernel'

require 'cs__contrast_patch/cs__contrast_patch'