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

cs__scoped_require 'contrast/agent/patching/policy/method_policy'

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
            # 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
              class_name = klass.cs__name
              instance_methods = klass.instance_methods
              instance_methods.concat(klass.private_instance_methods)
              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))
                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)
                current_context = Contrast::Agent::REQUEST_TRACKER.current
                next unless current_context

                dynamic_source = create_dynamic_source(current_context, dynamic_source_node, field, properties)
                node_id = Contrast::Utils::StringUtils.force_utf8(dynamic_source_node.id)
                current_context.activity.dynamic_sources[node_id] = dynamic_source
              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
            # @return [Contrast::Agent::Assess::Policy::SourceNode]
            def create_source_node class_name, method_name, tags
              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)
              url = Contrast::Agent::REQUEST_TRACKER.current&.request&.normalized_uri
              node.add_property(WRITE_QUERY_URL, url)
              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
              method_policy = Contrast::Agent::Patching::Policy::MethodPolicy.new
              method_policy.method_visibility = :public
              method_policy.instance_method = true
              method_policy.method_name = method_name
              method_policy.source_node = dynamic_source_node
              method_policy
            end

            # @param current_context [Contrast::Agent::RequestContext] the
            #   context of the request in which this source is to be created.
            # @param source_node [Contrast::Agent::Assess::Policy::SourceNode]
            #   the SourceNode that applies to this method
            # @param field [String] the name of the method to which this source
            #   applies
            # @param properties [Contrast::Agent::Assess::Properties] the
            #   properties which this source should relate to the data tracked
            #   as a result of this Source Method
            # @return [Contrast::Api::Dtm::DynamicSource] the message to send
            #   to the Service to allow it to report the events leading up to
            #   the creation of the Dynamic Source
            def create_dynamic_source current_context, source_node, field, properties
              dynamic_source = Contrast::Api::Dtm::DynamicSource.new
              dynamic_source.class_name = Contrast::Utils::StringUtils.force_utf8(source_node.class_name)
              dynamic_source.method_name = Contrast::Utils::StringUtils.force_utf8(field)
              dynamic_source.instance_method = source_node.instance_method?
              dynamic_source.target = Contrast::Utils::StringUtils.force_utf8(source_node.target_string)
              properties.events.each do |event|
                dynamic_source.events << event.to_dtm_event
              end
              dynamic_source.properties[READ_TABLE] = Contrast::Utils::StringUtils.force_utf8(source_node.class_name)
              dynamic_source.properties[READ_COLUMN] = Contrast::Utils::StringUtils.force_utf8(field)
              dynamic_source.properties[WRITE_QUERY_TIME] = Contrast::Utils::StringUtils.force_utf8(Contrast::Utils::Timer.now_ms)
              url = current_context.request.normalized_uri
              dynamic_source.properties[WRITE_QUERY_URL] = Contrast::Utils::StringUtils.force_utf8(url)
              dynamic_source
            end
          end
        end
      end
    end
  end
end