# 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/no_sqli'
require 'contrast/agent/protect/policy/rule_applicator'

module Contrast
  module Agent
    module Protect
      module Policy
        # This Module is how we apply the NoSQL Injection rule. It is called from
        # our patches of the targeted methods in which the execution of String
        # based NoSQL queries occur. It is responsible for deciding if the
        # infilter methods of the rule should be invoked.
        module AppliesNoSqliRule
          extend Contrast::Agent::Protect::Policy::RuleApplicator
          DATABASE_NOSQL = 'MongoDB'
          class << self
            def invoke method, _exception, properties, _object, args
              return unless valid_input?(args)
              return if skip_analysis?

              database = properties['database']
              operations = args[0]
              context = Contrast::Agent::REQUEST_TRACKER.current
              if operations.is_a?(Array)
                operations.each do |m|
                  handle_operation(context, database, method, m)
                end
              else
                handle_operation(context, database, method, operations)
              end
            end

            protected

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

            private

            def valid_input? args
              return false unless args
              return false unless args.cs__is_a?(Array) && args.any?

              args[0]
            end

            def handle_operation context, database, _action, operation
              data = extract_mongo_data(operation)
              return unless data && !data.empty?

              rule.infilter(context, database, data)
            end

            def extract_mongo_data operation
              if operation.cs__respond_to?(:selector)
                operation.selector
              elsif operation.cs__respond_to?(:documents)
                operation.documents
              end
            end
          end
        end
      end
    end
  end
end