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

require 'contrast/components/interface'

module Contrast
  module Agent
    module Assess
      module Policy
        module Propagator
          # Propagation that results in all the tags of the source being
          # applied to the target. Unlike other propagators, this actually
          # results in new source nodes to track which columns in the database
          # have been tainted.
          class DatabaseWrite < Contrast::Agent::Assess::Policy::Propagator::Base
            include Contrast::Components::Interface
            access_component :analysis

            class << self
              def propagate propagation_node, preshift, target
                class_type = preshift.object.cs__class
                class_name = class_type.cs__name
                tainted_columns = {}

                known_tainted = ASSESS.tainted_columns[class_name]
                propagation_node.sources.each do |source|
                  handle_write(propagation_node, source, preshift, target, known_tainted, tainted_columns)
                end
                return if tainted_columns.empty?

                if known_tainted
                  known_tainted.concat(tainted_columns.keys)
                else
                  ASSESS.tainted_columns[class_name] = tainted_columns.keys
                end

                Contrast::Agent::Assess::Policy::DynamicSourceFactory.create_sources class_type, tainted_columns
              end

              private

              def handle_write propagation_node, source, preshift, target, known_tainted, tainted_columns
                arg = preshift.args[source]
                return unless arg.cs__respond_to?(:each_pair)

                arg.each_pair do |key, value|
                  next unless value
                  next if known_tainted&.include?(key)
                  next unless (properties = Contrast::Agent::Assess::Tracker.properties!(value))

                  # TODO: RUBY-540 handle sanitization, handle nested objects
                  Contrast::Agent::Assess::Policy::PropagationMethod.apply_tags(propagation_node, value)
                  properties.build_event(propagation_node, value, preshift.object, target, preshift.args)
                  next unless tracked_value?(value)

                  tainted_columns[key] = properties
                end
              end
            end
          end
        end
      end
    end
  end
end