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

require 'contrast/agent/patching/policy/method_policy'
require 'contrast/components/logger'

module Contrast
  module Agent
    module Assess
      module Policy
        # This class is used to create dynamic source nodes & source nodes from a db model that receives untrusted data
        class DynamicSourceFactory
          DB_SOURCE_TYPE = 'TAINTED_DATABASE'
          WRITE_QUERY_TIME = 'writeDateTimeUtc'
          WRITE_QUERY_URL = 'writeRequestUrl'
          READ_TABLE = 'readTable'
          READ_COLUMN = 'readColumn'
          class << self
            include Contrast::Components::Logger::InstanceMethods

            # Given a Class representing a table in a Database and a map of
            # methods representing columns, generate sources for each method
            # such that calls to that method will result in a Source Event.
            #
            # @param klass [Class] the Class to taint
            # @param tainted_columns [Hash{String => Contrast::Agent::Assess::Properties}]
            #   the name of the method to taint, mapped to the properties it
            #   should apply
            def create_sources klass, tainted_columns
              return unless Contrast::ASSESS.require_dynamic_sources?

              class_name = klass.cs__name
              instance_methods = klass.instance_methods
              instance_methods.concat(klass.private_instance_methods)
              current_context = Contrast::Agent::REQUEST_TRACKER.current
              current_request = current_context&.request

              tainted_columns.each_pair do |field, properties|
                next unless properties

                method_name = field.to_sym
                # Move on if we already know about this Dynamic Source
                next if Contrast::Agent::Assess::Policy::Policy.instance.find_source_node(class_name, method_name, true)

                dynamic_source_node = create_source_node(class_name, method_name, Set.new(properties.tag_keys),
                                                         current_request)
                Contrast::Agent::Assess::Policy::Policy.instance.add_node(dynamic_source_node, :dynamic_source)
                method_policy = build_source_policy(method_name, dynamic_source_node)
                Contrast::Agent::Patching::Policy::Patcher.patch_method(klass, instance_methods, method_policy)
              end
            end

            private

            # Given a class and method, create a SourceNode for them that will
            # add the given tags as well as properties required to construct a
            # finding with a TAINTED_DATABASE source
            #
            # @param class_name [String] the name of the Class to patch
            # @param method_name [Symbol] the name of the Method to patch
            # @param tags [Set<String>] the tags this event should apply
            # @param request [Contrast::Agent::Request] the request during
            #  which this source is to be created.
            # @return [Contrast::Agent::Assess::Policy::SourceNode]
            def create_source_node class_name, method_name, tags, request
              node = Contrast::Agent::Assess::Policy::SourceNode.new
              node.class_name = class_name
              node.instance_method = true
              node.method_name = method_name
              node.method_visibility = :public
              node.target_string = Contrast::Utils::ObjectShare::RETURN_KEY
              node.type = DB_SOURCE_TYPE
              node.tags = tags
              node.tags << 'DATABASE_WRITE'
              node.add_property('dynamic_source_id', node.id)
              node.add_property('dynamic_source_name', "#{ class_name }.#{ method_name }")
              node.add_property(READ_TABLE, class_name)
              node.add_property(READ_COLUMN, method_name)
              node.add_property(WRITE_QUERY_TIME, Contrast::Utils::Timer.now_ms)
              node.add_property(WRITE_QUERY_URL, request&.normalized_uri)
              node
            end

            # @param method_name [Symbol] the name of the Method to which this
            #   method applies
            # @param dynamic_source_node [Contrast::Api::Dtm::DynamicSource]
            #   the SourceNode that applies to this method
            # @return [Contrast::Agent::Patching::Policy::MethodPolicy] the
            #   policy to apply to the given method
            def build_source_policy method_name, dynamic_source_node
              Contrast::Agent::Patching::Policy::MethodPolicy.new({
                                                                      method_visibility: :public,
                                                                      instance_method: true,
                                                                      method_name: method_name,
                                                                      source_node: dynamic_source_node
                                                                  })
            end
          end
        end
      end
    end
  end
end