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

require 'contrast/agent/protect/rule/xxe'
require 'contrast/agent/protect/policy/rule_applicator'
require 'contrast/utils/object_share'

module Contrast
  module Agent
    module Protect
      module Policy
        # This Module is how we apply the XXE rule. It is called from our patches
        # of the targeted methods in which XML parsing and entity resolution
        # occurs. It is responsible for deciding if the infilter methods of the
        # rule should be invoked.
        module AppliesXxeRule
          extend Contrast::Agent::Protect::Policy::RuleApplicator

          class << self
            def apply_rule method, _exception, _properties, object, args
              xml = args[0]
              xxe_check(method, xml, object)
            end

            # IO is tricky. If we can't rewind it, we can't fix it back to the
            # original state. To be safe, we'll skip non-rewindable IO objects.
            def apply_rule__io method, _exception, _properties, object, args
              need_rewind = false
              potential_xml = args[0]
              return unless potential_xml.cs__respond_to?(:rewind)

              xml = potential_xml.read
              need_rewind = true
              xxe_check(method, xml, object)
            ensure
              potential_xml.rewind if need_rewind
            end

            # Oga's Lexer is a special case b/c the information we need is on the
            # object itself -- specifically in the @data instance variable
            def apply_rule__lexer method, _exception, _properties, object, _args
              return unless valid_data_input?(object)

              data = object.instance_variable_get(DATA_KEY)
              xxe_check(method, data, object)
            ensure
              data.rewind if data&.cs__respond_to?(:rewind)
            end

            protected

            def rule_name
              Contrast::Agent::Protect::Rule::Xxe::NAME
            end

            private

            DATA_KEY = :@data
            def valid_data_input? object
              object.instance_variable_defined?(DATA_KEY) && object.instance_variable_get(DATA_KEY)
            end

            NOKOGIRI_MARKER = 'Nokogiri::'
            PARSER_NOKOGIRI = 'Nokogiri'
            OX_MARKER = 'Ox' # breaks marker pattern b/c Ox is entire classname
            PARSER_OX = 'Ox'
            OGA_MARKER = 'Oga::'
            PARSER_OGA = 'Oga'
            # Given an object, determine the XML parser type that it represents.
            #
            # @param object[Object] the parsing instance or Module
            # @return [String] the name of the parser
            def determine_parser object
              clazz = object.is_a?(Module) ? object : object.cs__class
              name = clazz.cs__name

              if name.start_with?(NOKOGIRI_MARKER)
                PARSER_NOKOGIRI
              elsif name.start_with?(OX_MARKER)
                PARSER_OX
              elsif name.start_with?(OGA_MARKER)
                PARSER_OGA
              end
            end

            # Given an xml, convert it to a String and pass it to the rule for
            # analysis.
            #
            # @param method [Symbol] the name of the method doing this work.
            # @param xml [Object] the container of the XML to be checked.
            # @param potential_parser [Object] the entity that may be an XML
            #   parser.
            # @raise [Contrast::SecurityException] Security exception if an XXE
            #   attack is found and the rule is in block mode.
            def xxe_check method, xml, potential_parser
              return if skip_analysis?
              return unless xml

              parser = determine_parser(potential_parser)
              return unless parser

              if xml.cs__is_a?(String)
                rule.infilter(Contrast::Agent::REQUEST_TRACKER.current, parser, xml)
              elsif xml.cs__respond_to?(:each_line)
                xml.each_line do |line|
                  rule.infilter(Contrast::Agent::REQUEST_TRACKER.current, parser, line)
                end
              elsif xml.cs__respond_to?(:each)
                xml.each do |chunk|
                  rule.infilter(Contrast::Agent::REQUEST_TRACKER.current, parser, chunk)
                end
              end
            rescue Contrast::SecurityException => e
              raise(e)
            rescue StandardError => e
              parser ||= Contrast::Utils::ObjectShare::UNKNOWN
              logger.error('Error applying xxe', e, module: potential_parser.cs__class.cs__name, method: method,
                                                    parser: parser)
            end
          end
        end
      end
    end
  end
end