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

require 'monitor'
require 'contrast/components/interface'

require 'contrast/agent'
require 'contrast/logger/log'
require 'contrast/agent/patching/policy/method_policy'
require 'contrast/agent/patching/policy/patch_status'
require 'contrast/agent/patching/policy/trigger_node'

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

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

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

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 Patch
          class << self
            include Contrast::Agent::Assess::Policy::SourceMethod
            include Contrast::Agent::Assess::Policy::PropagationMethod
            include Contrast::Agent::Assess::Policy::TriggerMethod

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

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

            def enter_method_scope! method_policy
              method_policy.scopes_to_enter.each do |scope|
                enter_scope!(scope)
              end
            end

            def exit_method_scope! method_policy
              method_policy.scopes_to_exit.each do |scope|
                exit_scope!(scope)
              end
            end

            # THIS IS CALLED FROM C. Do not change the signature lightly.
            #
            # This method functions to call the infilter methods from our
            # patches, allowing for analysis and reporting at the point just
            # before the patched code is invoked.
            #
            # @param method_policy [Contrast::Agent::Patching::Policy::MethodPolicy]
            #   Mapping of the triggers on the given method.
            # @param method [Symbol] The method into which we're patching
            # @param exception [StandardError] Any exception raised during the
            #   call of the patched method.
            # @param object [Object] The object on which the method is invoked,
            #   typically what would be returned by self.
            # @param args [Array<Object>] The arguments passed to the method
            #   being invoked.
            def apply_pre_patch method_policy, method, exception, object, args
              apply_protect(method_policy, method, exception, object, args)
              apply_inventory(method_policy, method, exception, object, args)
            rescue Contrast::SecurityException => e
              # We were told to block something, so we gotta. Don't catch this
              # one, let it get back to our Middleware or even all the way out to
              # the framework
              raise e
            rescue StandardError => e
              # Anything else was our bad and we gotta catch that to allow for
              # normal application flow
              logger.error('Unable to apply pre patch to method.', e)
            rescue Exception => e # rubocop:disable Lint/RescueException
              # This is something like NoMemoryError that we can't
              # hope to handle.  Nonetheless, shouldn't leak scope.
              exit_contrast_scope!
              raise e
            end

            # THIS IS CALLED FROM C. Do not change the signature lightly.
            #
            # This method functions to call the infilter methods from our
            # patches, allowing for analysis and reporting at the point just
            # after the patched code is invoked
            #
            # @param method_policy [Contrast::Agent::Patching::Policy::MethodPolicy]
            #   Mapping of the triggers on the given method.
            # @param preshift [Contrast::Agent::Assess::PreShift] The capture
            #   of the state of the code just prior to the invocation of the
            #   patched method.
            # @param object [Object] The object on which the method was
            #   invoked, typically what would be returned by self.
            # @param ret [Object] The return of the method that was invoked.
            # @param args [Array<Object>] The arguments passed to the method
            #   being invoked.
            # @param block [Proc] The block passed to the method that was
            #   invoked.
            def apply_post_patch method_policy, preshift, object, ret, args, block
              apply_assess(method_policy, preshift, object, ret, args, block)
            rescue StandardError => e
              logger.error('Unable to apply post patch to method.', e)
            end

            # Apply the Protect patch which applies to the given method.
            #
            # @param method_policy [Contrast::Agent::Patching::Policy::MethodPolicy]
            #   Mapping of the triggers on the given method.
            # @param method [Symbol] The method into which we're patching
            # @param exception [StandardError] Any exception raised during the
            #   call of the patched method.
            # @param object [Object] The object on which the method is invoked,
            #   typically what would be returned by self.
            # @param args [Array<Object>] The arguments passed to the method
            #   being invoked.
            def apply_protect method_policy, method, exception, object, args
              return unless AGENT.enabled?
              return unless PROTECT.enabled?

              apply_trigger_only(method_policy&.protect_node,
                                 method,
                                 exception,
                                 object,
                                 args)
            end

            # Apply the Inventory patch which applies to the given method.
            #
            # @param method_policy [Contrast::Agent::Patching::Policy::MethodPolicy]
            #   Mapping of the triggers on the given method.
            # @param method [Symbol] The method into which we're patching
            # @param exception [StandardError] Any exception raised during the
            #   call of the patched method.
            # @param object [Object] The object on which the method is invoked,
            #   typically what would be returned by self.
            # @param args [Array<Object>] The arguments passed to the method
            #   being invoked.
            def apply_inventory method_policy, method, exception, object, args
              return unless INVENTORY.enabled?

              apply_trigger_only(method_policy&.inventory_node,
                                 method,
                                 exception,
                                 object,
                                 args)
            end

            # Apply the Assess patches which apply to the given method.
            #
            # @param method_policy [Contrast::Agent::Patching::Policy::MethodPolicy]
            #   Mapping of the triggers on the given method.
            # @param preshift [Contrast::Agent::Assess::PreShift] The capture
            #   of the state of the code just prior to the invocation of the
            #   patched method.
            # @param object [Object] The object on which the method was
            #   invoked, typically what would be returned by self.
            # @param ret [Object] The return of the method that was invoked.
            # @param args [Array<Object>] The arguments passed to the method
            #   being invoked.
            # @param block [Proc] The block passed to the method that was
            #   invoked.
            def apply_assess method_policy, preshift, object, ret, args, block
              source_ret = nil
              propagated_ret = nil
              return ret unless method_policy && ASSESS.enabled?

              current_context = Contrast::Agent::REQUEST_TRACKER.current
              return ret unless current_context.analyze_request?

              trigger_node = method_policy.trigger_node
              Contrast::Agent::Assess::Policy::TriggerMethod.apply_trigger_rule(trigger_node, object, ret, args) if trigger_node
              if method_policy.source_node
                # If we were given a frozen return, and it was the target of a
                # source, and we have frozen sources enabled, we'll need to
                # replace the return. Note, this is not the default case.
                source_ret = Contrast::Agent::Assess::Policy::SourceMethod.source_patchers(method_policy, object, ret, args)
              end
              if method_policy.propagation_node
                propagated_ret = Contrast::Agent::Assess::Policy::PropagationMethod.apply_propagation(
                    method_policy,
                    preshift,
                    object,
                    source_ret || ret,
                    args,
                    block)
              end
              handle_return(propagated_ret, source_ret, ret)
            rescue StandardError => e
              logger.error('Unable to assess method call.', e)
              handle_return(propagated_ret, source_ret, ret)
            rescue Exception => e # rubocop:disable Lint/RescueException
              logger.error('Unable to assess method call.', e)
              handle_return(propagated_ret, source_ret, ret)
              raise e
            end

            # Generic invocation of the Inventory or Protect patch which apply
            # to the given method.
            #
            # @param trigger_node [Contrast::Agent::Inventory::Policy::TriggerNode]
            #   Mapping of the specific trigger on the given method.
            # @param method [Symbol] The method into which we're patching
            # @param exception [StandardError] Any exception raised during the
            #   call of the patched method.
            # @param object [Object] The object on which the method is invoked,
            #   typically what would be returned by self.
            # @param args [Array<Object>] The arguments passed to the method
            #   being invoked.
            def apply_trigger_only trigger_node, method, exception, object, args
              return unless trigger_node

              # If that rule only applies in the case of an exception being
              # thrown and there's no exception here, move along, or vice versa
              return if trigger_node.on_exception && !exception
              return if !trigger_node.on_exception && exception

              # Each patch has an applicator that handles logic for it. Think
              # of this as being similar to propagator actions, most closely
              # resembling CUSTOM - they all have a common interface but their
              # own logic based on what's in the method(s) they've been patched
              # into.
              # Each patch also knows the method of its applicator. Some
              # things, like AppliesXxeRule, have different methods depending
              # on the library patched. This lets us handle the boilerplate of
              # patching while still allowing for custom handling of the
              # methods.
              applicator_method = trigger_node.applicator_method
              # By calling send like this, we can reuse all the patching.
              # We `send` to the given method of the given class
              # (applicator) since they all accept the same inputs
              trigger_node.applicator.send(applicator_method, method, exception, trigger_node.properties, object, args)
            end

            # Method to choose which replaced return from the post_patch to
            # actually return
            #
            # @param propagated_ret [Object, nil] The replaced return from the
            #   propagation patch.
            # @param source_ret [Object, nil] The replaced return from the
            #   source patch.
            # @param ret [Object, nil] The original return of the patched
            #   method.
            # @return [Object, nil] The thing to return from the post patch.
            def handle_return propagated_ret, source_ret, ret
              safe_return = propagated_ret || source_ret || ret
              safe_return.rewind if Contrast::Utils::IOUtil.should_rewind?(safe_return)
              safe_return
            end

            # Given a module and method, construct an expected name for the
            # alias by which Contrast will reference the original.
            #
            # @param patched_class [Module] the module being patched
            # @param patched_method [Symbol] the method being patched
            # @return [Symbol]
            def build_method_name patched_class, patched_method
              (Contrast::Utils::ObjectShare::CONTRAST_PATCHED_METHOD_START +
                  patched_class.cs__name.gsub('::', '_').downcase +
                  Contrast::Utils::ObjectShare::UNDERSCORE +
                  patched_method.to_s).to_sym
            end

            # Given a method, return a symbol in the format
            # <method_start>_unbound_<method_name>
            def build_unbound_method_name patcher_method
              (Contrast::Utils::ObjectShare::CONTRAST_PATCHED_METHOD_START +
                  'unbound' +
                  Contrast::Utils::ObjectShare::UNDERSCORE +
                  patcher_method.to_s).to_sym
            end

            # @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 method to be patched.
            # @return [Boolean] if patched, either by this invocation or a
            #   previous, or not
            def instrument_with_alias mod, methods, method_policy
              cs_method_name = build_method_name(mod, method_policy.method_name)
              # we've already patched this class, don't do it again
              return true if methods.include?(cs_method_name)

              begin
                contrast_define_method(mod, method_policy, cs_method_name)
              rescue NameError => e
                # This shouldn't happen anymore, but just in case calling alias
                # results in a NameError, we'll be safe here.
                logger.error(
                    'Attempted to alias a method on a Module that doesn\'t respond to it.',
                    e,
                    module: mod.cs__name,
                    method: method_policy.method_name)
                return false
              end
              true
            end

            # @param mod [Module] the module in which the patch should be
            #   placed.
            # @param method_policy [Contrast::Agent::Patching::Policy::MethodPolicy]
            #   the policy that applies to the method to be patched.
            # @return [Boolean] if patched, either by this invocation or a
            #   previous, or not
            def instrument_with_prepend mod, method_policy
              contrast_prepend_method(mod, method_policy)
            end

            # @param unbound_method [UnboundMethod] An unbound method, that doesn't reference its binding.
            #   This method executes C hooking code.
            def register_c_hook unbound_method
              # current binding is as meaningless as any other.  but we need something
              unbound_method.bind_call(self)
            end

            # @param target_module_name [String] Fully-qualified module name, as string, to which the C patch applies.
            # @param unbound_method [UnboundMethod] An unbound method, to be patched into target_module.
            # @param impl [Symbol] Strategy for applying the patch: { :alias_instance, :alias_singleton, or :prepend }:
            #   :alias_instance     -> alias instance method of module
            #   :alias_singleton    -> alias instance method of singleton class of module
            #                          (equivalent to :alias, where `module = module.singleton class`)
            #                          (this is a.k.a. "class-method patch")
            #   :prepend            -> prepend instance method of module
            #   [prepending singleton is easily supported too, just not implemented yet.]
            # @return [Symbol] new alias for the underlying method (presumably, so the patched method can call it)
            def register_c_patch target_module_name, unbound_method, impl = :alias_instance
              # These could be set as AfterLoadPatches.
              method_name = unbound_method.name.to_sym
              underlying_method_name = build_unbound_method_name(method_name).to_sym

              target_module = Module.cs__const_get(target_module_name)

              target_module = target_module.cs__singleton_class if %i[alias_singleton prepend].include? impl

              visibility = if target_module.private_instance_methods(false).include?(method_name)
                             :private
                           elsif target_module.protected_instance_methods(false).include?(method_name)
                             :protected
                           elsif target_module.public_instance_methods(false).include?(method_name)
                             :public
                           else
                             raise NoMethodError,
                                   <<~ERR
                                     Tried to register a C-defined #{ impl } patch for \
                                     #{ target_module_name }##{ method_name }, but can't find :#{ method_name }.
                                   ERR
                           end

              case impl
              when :alias_instance, :alias_singleton
                unless target_module.instance_methods(false).include? underlying_method_name
                  # alias_method may be private
                  target_module.send(:alias_method, underlying_method_name, method_name)
                  # TODO: RUBY-1052
                  # rubocop:disable Performance/Kernel/DefineMethod
                  target_module.send(:define_method, method_name, unbound_method.bind(target_module))
                  # rubocop:enable Performance/Kernel/DefineMethod
                end
                target_module.send(visibility, method_name) # e.g., module.private(:my_method)
              when :prepend
                prepending_module = Module.new
                # TODO: RUBY-1052
                # rubocop:disable Performance/Kernel/DefineMethod
                prepending_module.send(:define_method, method_name, unbound_method.bind(target_module))
                # rubocop:enable Performance/Kernel/DefineMethod
                prepending_module.send(visibility, method_name)
                # This prepends to the singleton class (it patches a class method)
                target_module.prepend prepending_module
              end
              # Ougai::Logger.create_item_with_2args calls Hash#[]=, so we
              # can't invoke this logging method or we'll seg fault as we'd
              # change the method definition mid-call
              # if method_name != :[]=
              #   logger.trace(
              #       'Registered C-defined patch',
              #       implementation: impl,
              #       module: target_module_name,
              #       method: method_name,
              #       visibility: visibility)
              # end
              underlying_method_name.to_sym
            end

            # @return [Boolean]
            def skip_contrast_analysis?
              return true if in_contrast_scope?
              return true unless defined?(Contrast::Agent::REQUEST_TRACKER)
              return true unless Contrast::Agent::REQUEST_TRACKER.current&.analyze_request?

              false
            end

            # Skip if we should skip_contrast_analysis?, sampling says to ignore this
            # request, or assess has been disabled.
            #
            # @return [Boolean]
            def skip_assess_analysis?
              return true if skip_contrast_analysis?

              !ASSESS.enabled?
            end
          end
        end
      end
    end
  end
end

require 'cs__contrast_patch/cs__contrast_patch'