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

require 'contrast/agent/patching/policy/after_load_patch'
require 'contrast/components/logger'
require 'contrast/framework/manager'
require 'contrast/extension/extension'
require 'contrast/extension/assess/kernel'
require 'contrast/extension/assess/marshal'

module Contrast
  module Agent
    module Patching
      module Policy
        # Some modules diverge from our generic instrumentation and require custom instrumentation
        # after they've been loaded
        module AfterLoadPatcher
          include Contrast::Components::Logger::InstanceMethods

          # After initialization run a catchup check to instrument any already loaded modules we care about
          def catchup_after_load_patches
            apply_require_patches!
            apply_direct_patches!
            apply_load_patches!
          end

          private

          # These patches need to be applied directly, not from our policy, so
          # do so and do so only once. This should be the new standard so that
          # there are no require time side effects of loading our core
          # extensions.
          def apply_direct_patches!
            @_apply_direct_patches ||= begin
              paths = %w[
                array
                basic_object
                module
                fiber_track
                hash
                kernel
                marshal_module
                regexp
                string
                string_interpolation
              ].cs__freeze
              paths.each do |p|
                path_part = "cs__assess_#{ p }"
                Contrast::Extension::Assess::InstrumentHelper.instrument("#{ path_part }/#{ path_part }")
              end
              # apply Kernel#exec alias patch:
              unless Contrast::Agent::Assess.cs__object_method_prepended?(Kernel, :exec, false)
                apply_kernel_exec_alias_patch
              end
              # apply Marshal load alias patch:
              unless Contrast::Agent::Assess.cs__object_method_prepended?(Marshal, :load, false)
                apply_marshal_load_alias_patch
              end
              apply_array_join_prepend_patch if Contrast::Agent::Assess.cs__object_method_prepended?(Array, :join, true)
              true
            end
          end

          def apply_load_patches!
            after_load_patches.each do |after_load_patch|
              next unless after_load_patch.target_defined?
              next if ::Contrast::AGENT.skip_instrumentation?(after_load_patch.module_name)

              logger.trace('Catching up on already loaded afterload patch - applying instrumentation',
                           module: after_load_patch.module_name)
              after_load_patch.instrument!
            rescue NameError => e
              logger.error('Method undefined in afterload patch', e, module: after_load_patch.module_name,
                                                                     method: after_load_patch.method_to_instrument)
            rescue StandardError => e
              logger.error('Afterload patch failed to apply', e, module: after_load_patch.module_name,
                                                                 method: after_load_patch.method_to_instrument)
            end
            after_load_patches.delete_if(&:applied?)
          end

          # These patches need to be applied directly, not from our policy, and
          # are applied as a result of requiring the file as they alias methods
          # directly, allowing us to control things like scope and exception
          # handling
          def apply_require_patches!
            @_apply_require_patches ||= begin
              require('contrast/extension/thread')
              true
            rescue LoadError, StandardError => e
              logger.error('failed instrumenting apply_require_patches!', e)
              false
            end
          end

          def after_load_patches
            @_after_load_patches ||= Contrast::Agent.framework_manager.find_after_load_patches
          end

          # Use for any checks after we've initialized
          def load_patches_for_module module_name
            return if after_load_patches.empty?

            patch = after_load_patches.find { |after_load_patch| after_load_patch.applies?(module_name) }
            return unless patch

            logger.trace('Detected loading of afterload patch - applying instrumentation', module: module_name)
            patch.instrument!
            after_load_patches.delete_if(&:applied?)
          end

          # Applies the Kernel#exec alias patch
          def apply_kernel_exec_alias_patch
            Kernel.extend(Contrast::Extension::Assess::ContrastKernel)
            Marshal.alias_method(:cs__kernel_exec, :exec)
            Marshal.alias_method(:exec, :cs__kernel_exec)
          end

          # Applies the Marshal#load alias patch
          def apply_marshal_load_alias_patch
            Marshal.extend(Contrast::Extension::Assess::ContrastMarshal)
            Marshal.alias_method(:cs__marshal_load, :load)
            Marshal.alias_method(:load, :cs__marshal_load)
          end

          # Prepend Contrast::Extension::Assess::ContrastArray to pick up the ContrastArray#join method
          def apply_array_join_prepend_patch
            Array.prepend(Contrast::Extension::Assess::ContrastArray)
          end
        end
      end
    end
  end
end