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

require 'contrast/components/logger'
require 'contrast/components/scope'
require 'contrast/agent/assess/events/event_data'

module Contrast
  module Extension
    module Assess
      # This is our patch of the Marshal class
      # Disclaimer: there may be a better way, but we're in a 'get it work' state.
      # Hopefully, we'll be in a 'get it right' state soon.
      # This module is used for our Marshal.load patches
      class MarshalPropagator
        extend Contrast::Components::Scope::InstanceMethods

        class << self
          extend Contrast::Components::Logger::InstanceMethods
          include Contrast::Components::Logger::InstanceMethods

          def cs__load_protect arg
            return if in_contrast_scope?

            with_contrast_scope do
              Contrast::Agent::Protect::Policy::AppliesDeserializationRule.prepended_invoke(arg)
            end
            nil
          end

          def cs__load_assess source, ret
            with_contrast_scope do
              return unless ::Contrast::ASSESS.non_request_tracking? || Contrast::Agent::REQUEST_TRACKER.current

              args = [source]
              # source might not be all the args passed in, but it is the one we care
              # about. we could pass in all the args in the last param here if it
              # becomes an issue in rendering on TS
              Contrast::Agent::Assess::Policy::TriggerMethod.build_finding(trigger_node('Marshal', :load), source,
                                                                           self, ret, *args)
              return unless (properties = Contrast::Agent::Assess::Tracker.properties!(ret))

              properties.copy_from(source, ret)

              node = Contrast::Agent::Assess::Policy::Policy.instance.find_propagator_node('Marshal', :load, false)
              event_data = Contrast::Agent::Assess::Events::EventData.new(node, ret, self, ret, args)
              properties.build_event(event_data)
            rescue StandardError => e
              logger.error('Unable to run Assess for Marshal.load', e)
            end
          end

          def trigger_node clazz, method
            triggers = Contrast::Agent::Assess::Policy::Policy.instance.triggers
            return unless triggers

            triggers.find { |node| node.class_name == clazz && node.method_name == method }
          end
        end
      end
    end
  end
end