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

require 'contrast/agent/protect/rule/base'
require 'contrast/agent/reporting/details/xxe_details'
require 'contrast/agent/reporting/details/xxe_match'
require 'contrast/agent/reporting/details/xxe_wrapper'
require 'contrast/utils/timer'
require 'contrast/components/logger'

module Contrast
  module Agent
    module Protect
      module Rule
        # Implementation of the XXE Protect Rule used to evaluate XML calls for exploit
        # of unsafe external entity resolution.
        class Xxe < Contrast::Agent::Protect::Rule::Base
          include Contrast::Components::Logger::InstanceMethods
          INPUT_NAME = 'XML Prolog'

          NAME = 'xxe'
          BLOCK_MESSAGE = 'XXE rule triggered. Response blocked.'
          EXTERNAL_ENTITY_PATTERN = /<!ENTITY\s+[a-zA-Z0-f]+\s+(?:SYSTEM|PUBLIC)\s+(.*?)>/.cs__freeze

          class << self
            # @param attack_sample [Contrast::Api::Dtm::RaspRuleSample]
            # @return [Hash] the details for this specific rule
            def extract_details attack_sample
              {
                  xml: attack_sample.xxe.xml,
                  declaredEntities: attack_sample.xxe.declared_entities.map do |entity|
                                      {
                                          start: entity.start_idx,
                                          end: entity.end_idx
                                      }
                                    end,
                  entitiesResolved: attack_sample.xxe.entities_resolved.map do |entity|
                                      {
                                          systemId: entity.system_id,
                                          publicId: entity.public_id
                                      }
                                    end
              }
            end
          end

          def rule_name
            NAME
          end

          # Given an xml, evaluate it for an XXE attack. There's no return here
          # as this method handles appending the evaluation to the request
          # context, connecting it to the reporting mechanism at request end.
          #
          # @param context [Contrast::Agent::RequestContext] the context of the
          #   request in which this input is evaluated.
          # @param framework [Object] the name of the Parser being used.
          # @param xml [Object] the container of the XML to be checked.
          # @raise [Contrast::SecurityException] Security exception if an XXE
          #   attack is found and the rule is in block mode.
          def infilter context, framework, xml
            return if protect_excluded_by_url?(context)

            result = find_attacker(context, xml, framework: framework)
            return unless result

            append_to_activity(context, result)
            return unless blocked?

            cef_logging(result, :successful_attack, value: xml)
            raise(Contrast::SecurityException.new(self, BLOCK_MESSAGE))
          end

          protected

          # Given an XML, find any externally resolved entities and create an
          # Attack Result for them
          #
          # @param context [Contrast::Agent::RequestContext] the context of the
          #   request in which this input is evaluated.
          # @param xml [String] the literal value of the XML being checked for
          #   external entity resolution
          # @param _kwargs [Hash]
          # @return [Contrast::Api::Dtm::AttackResult, nil] the determination
          #   as to whether or not this XML has an XXE attack in it.
          def find_attacker context, xml, **_kwargs
            return unless xml
            return if protect_excluded_by_code?

            xxe_details = build_details(xml)
            return unless xxe_details

            ia_result = build_evaluation(xxe_details.xml)
            build_attack_with_match(context, ia_result, nil, nil, details: xxe_details)
          end

          # Given an XML determined to be unsafe, build out the details of the
          # attack. The details will include a substring of the given XML up to
          # the end of the prolog, where the external entities are declared.
          #
          # @param xml [String] the literal value of the XML being checked for
          #   external entity resolution
          # @return [Contrast::Api::Dtm::XxeDetails] The details of
          #   the XXE attack and the index of the last entity discovered
          def build_details xml
            last_idx = 0
            ss = StringScanner.new(xml)
            while ss.scan_until(EXTERNAL_ENTITY_PATTERN)
              last_idx = ss.pos
              entity_wrapper = Contrast::Agent::Protect::Rule::Xxe::EntityWrapper.new(ss.matched)
              next unless entity_wrapper.external_entity?

              xxe_details ||= Contrast::Agent::Reporting::Details::XxeDetails.new
              xxe_details.declared_entities << build_match(ss)
              xxe_details.entities_resolved << build_wrapper(entity_wrapper)
            end
            # For our definition, the prolog goes from the start of the XML
            # string to the end of the last entity declaration.
            xxe_details.xml = Contrast::Utils::StringUtils.protobuf_safe_string(xml[0, last_idx]) if xxe_details

            xxe_details
          end

          def build_sample context, ia_result, _url, **kwargs
            sample = build_base_sample(context, ia_result)
            sample.user_input = build_user_input(ia_result)
            sample.details = kwargs[:details]
            sample
          end

          def build_user_input ia_result
            input = Contrast::Agent::Reporting::UserInput.new
            input.key = INPUT_NAME
            input.input_type = :UNKNOWN
            input.document_type = :XML
            input.value = ia_result.value
            input
          end

          private

          # We know that this attack happened, so the result is always matched
          # and the level is always critical. Only variable is the XML value
          # supplied by the attacker.
          def build_evaluation xml
            ia_result = Contrast::Api::Settings::InputAnalysisResult.new
            ia_result.rule_id = rule_name
            ia_result.input_type = :UNKNOWN
            ia_result.value = Contrast::Utils::StringUtils.protobuf_safe_string(xml)
            ia_result
          end

          def build_match string_scanner
            match = Contrast::Agent::Reporting::Details::XxeMatch.new
            match.end_idx = string_scanner.pos.to_i
            match.start_idx = match.end_idx - string_scanner.matched_size
            match
          end

          def build_wrapper entity_wrapper
            wrapper = Contrast::Agent::Reporting::Details::XxeWrapper.new
            wrapper.system_id = Contrast::Utils::StringUtils.protobuf_safe_string(entity_wrapper.system_id)
            wrapper.public_id = Contrast::Utils::StringUtils.protobuf_safe_string(entity_wrapper.public_id)
            wrapper
          end
        end
      end
    end
  end
end