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

module Contrast
  module Agent
    module Protect
      module Rule
        # This class handles our implementation of the Untrusted
        # Deserialization Protect rule.
        class Deserialization < Contrast::Agent::Protect::Rule::Base
          # The TeamServer recognized name of this rule
          NAME = 'untrusted-deserialization'

          # The rule specific reason for raising a security exception.
          BLOCK_MESSAGE = 'Untrusted Deserialization rule triggered. Deserialization blocked.'

          # Gadgets that map to ERB modules
          ERB_GADGETS = %w[
            object:ERB
          ].cs__freeze

          # Gadgets that map to ActionDispatch modules
          ACTION_DISPATCH_GADGETS = %w[
            object:ActionDispatch::Routing::RouteSet::NamedRouteCollection
          ].cs__freeze

          # Gadgets that map to Arel Modules
          AREL_GADGETS = %w[
            string:Arel::Nodes::SqlLiteral
            object:Arel::Nodes
          ].cs__freeze

          # Used to indicate to TeamServer the gadget is an ERB module
          ERB = 'ERB'
          # Used to indicate to TeamServer the gadget is an ActionDispatch
          # module
          DISPATCH = 'ActionDispatch'
          # Used to indicate to TeamServer the gadget is an Arel module
          AREL = 'Arel'

          # Return the TeamServer understood id / name of this rule.
          # @return [String] the TeamServer understood id / name of this rule.
          def name
            NAME
          end

          # Return the specific blocking message for this rule.
          # @return [String] the reason for the raised security exception.
          def block_message
            BLOCK_MESSAGE
          end

          # Per the spec, this rule applies regardless of input. Only the mode
          # of the rule and code exclusions apply at this point.
          # @return [Boolean] should the rule apply to this call.
          def infilter? _context
            return false unless enabled?
            return false if protect_excluded_by_code?

            true
          end

          # Determine if the input being deserialized is an attack and take
          # appropriate actions.
          # @param context [Contrast::Agent::RequestContext] the request
          #   context in which this attack is occurring.
          # @param serialized_input [String] the string being deserialized.
          # @raise [Contrast::SecurityException] if attack detected while in
          #   block mode.
          def infilter context, serialized_input
            return unless infilter?(context)

            gadget = find_gadget(serialized_input)
            # If there isn't a gadget, this isn't an attack, so we have nothing
            # to do here. Let the application carry on business as usual.
            return unless gadget

            ia_result = build_evaluation(serialized_input)
            kwargs = { GADGET_TYPE: gadget }
            result = build_attack_with_match(context, ia_result, nil, serialized_input, **kwargs)
            append_to_activity(context, result)

            raise Contrast::SecurityException.new(self, block_message) if blocked?
          end

          # Determine if the issued command was called while we're
          # deserializing. If we are, treat this as an attack.
          # @param gadget_command [String] the command being executed during
          #   the deserialization call.
          # @raise [Contrast::SecurityException] if attack detected while in
          #   block mode.
          def check_command_scope gadget_command
            return unless deserializing?

            context = Contrast::Agent::REQUEST_TRACKER.current
            kwargs = { COMMAND_SCOPE: true }
            ia_result = build_evaluation(gadget_command)
            result = build_attack_with_match(context, ia_result, nil, gadget_command, **kwargs)
            append_to_activity(context, result)
            raise Contrast::SecurityException.new(self, BLOCK_MESSAGE) if blocked?
          end

          # I don't know a better way to do this without introducing another
          # scope.
          # The policy files are designed around using the Module names, not
          # the file names, but stack is built using filenames. These are the
          # files and methods we patch for this rule.
          #
          # There's not real test for this since I can't figure out how to get
          # a command to execute from within the .load methods
          # Assuming someone else does, this should just work (tested with
          # Screener and the combination "%w[erb.rb `result']")
          DESERIALIZER_STACK = [
            %w[/psych.rb: `load'],
            %w[/marshal.rb: `load']
          ].cs__freeze
          # We're considered deserializing if the call stack includes a
          # reference to the file in which a serializer is defined and
          # the method of that serializer responsible for deserializing.
          # @return [Boolean] if the caller indicates deserialization.
          def deserializing?
            Kernel.caller.product(DESERIALIZER_STACK).find do |(frame, (class_signature_form, method_signature_form))|
              frame.end_with?(method_signature_form) && frame.include?(class_signature_form)
            end
          end

          protected

          # Build the RaspRuleSample for the detected Deserialization attack.
          # @param context [Contrast::Agent::RequestContext] the request
          #   context in which this attack is occurring.
          # @param input_analysis_result [Contrast::Api::Settings::InputAnalysisResult]
          #   the result of the analysis done by this rule.
          # @param _candidate_string [nil] unused.
          # @param kwargs [Hash, nil] Hash of inputs used by this rule to flesh
          #   out the report to TeamServer in order to provide specific
          #   information for this rule.
          # @return [Contrast::Api::Dtm::RaspRuleSample] the information needed
          #   to render this attack event in TeamServer.
          def build_sample context, input_analysis_result, _candidate_string, **kwargs
            sample = build_base_sample(context, input_analysis_result)
            sample.untrusted_deserialization = Contrast::Api::Dtm::UntrustedDeserializationDetails.new

            deserializer = kwargs[:GADGET_TYPE]
            sample.untrusted_deserialization.deserializer = Contrast::Utils::StringUtils.protobuf_safe_string(deserializer)

            command = !!kwargs[:COMMAND_SCOPE]
            sample.untrusted_deserialization.command = command

            sample
          end

          private

          # Used to name this input since input analysis isn't done for this
          # rule
          INPUT_NAME = 'Serialized Gadget'
          # We know that this attack happened, so the result is always matched
          # and the level is always critical. Only variable is the Gadget
          # supplied by the attacker.
          # @param gadget_string [String] the input to be deserialized in which
          #   the gadget exists or the command that resulted from deserializing
          #   an input not detected in the initial infilter.
          # @return [Contrast::Api::Settings::InputAnalysisResult] the result
          #   of the analysis done by this rule.
          def build_evaluation gadget_string
            ia_result = Contrast::Api::Settings::InputAnalysisResult.new
            ia_result.rule_id = name
            ia_result.input_type = :UNKNOWN
            ia_result.key = INPUT_NAME
            ia_result.value = Contrast::Utils::StringUtils.protobuf_safe_string(gadget_string)
            ia_result
          end

          # Find the gadget within the serialized input, if such a gadget
          # exists.
          # @param serialized_input [String] the string being deserialized in
          #   which the gadget may exist.
          # @return [String, nil] the gadget, if found.
          def find_gadget serialized_input
            if ERB_GADGETS.any? { |gadget| serialized_input.index(gadget) }
              ERB
            elsif ACTION_DISPATCH_GADGETS.any? { |gadget| serialized_input.index(gadget) }
              DISPATCH
            elsif AREL_GADGETS.any? { |gadget| serialized_input.index(gadget) }
              AREL
            end
          end
        end
      end
    end
  end
end