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

module Contrast
  module Agent
    module Patching
      module Policy
        # This indicates the status of the Module into which Contrast is being
        # woven
        class PatchStatus
          class << self
            # Gives the current status for the provided module, setting one if
            # one does not exist.
            #
            # @param mod [Module] the Module for which the status is asked
            # @return [Contrast::Agent::Patching::Policy::PatchStatus]
            def get_status mod
              return mod.cs__const_get(status_key) if mod.cs__const_defined?(status_key, false)

              mod.cs__const_set(status_key, new)
            end

            # Allows our C patches to look up the :MethodPolicy for a given
            # Module and method
            #
            # @param mod [Module or Class] the entity on which the patch has
            #   been placed
            # @param method [Symbol] the name of the method to which the patch
            #   applies
            # @param is_instance_method [Boolean] is the method being patched
            #   an instance method or not (not implying class method).
            # @return [Contrast::Agent::Patching::Policy::MethodPolicy]
            def info_for mod, method, is_instance_method
              mod = mod.cs__class unless mod.cs__is_a?(Module)
              method_key = method_name_key(method, is_instance_method)
              # Ideally, we'll get to a point where this is all that is needed
              ret = find_info_for_one(mod, method_key)
              return ret if ret

              # But to start, the lookup will be on the ancestor, not the
              # calling class. We need to traverse the ancestors, included
              # modules, and Module itself the first time a method is called in
              # a given class. Then we save that lookup to the class
              ret = find_info_for(mod.ancestors, method_key)
              ret ||= find_info_for(mod.cs__singleton_class.included_modules, method_key)
              ret ||= find_info_for(MOD_ARRAY, method_key)
              update_holder(mod, method_key, ret) if ret
              ret
            end

            # Allows our C patches to set the :MethodPolicy for a given
            # Module and method
            #
            # @param mod [Module or Class] the entity on which the patch will
            #   be placed
            # @param method [Symbol] the name of the method to which the patch
            #   applies
            # @param method_policy [:MethodPolicy] the policy that applies to
            #   the given mod & method
            # @param is_instance_method [Boolean] is the method being patched
            #   an instance method or not (not implying class method).
            # @param cs_method [Symbol] the name to which the original method
            #   was aliased, cached for performance reasons
            def set_info_for mod, method, method_policy, is_instance_method, cs_method
              mod.instance_variable_set(method_info_key, {}) unless mod.instance_variable_defined?(method_info_key)
              holder = mod.instance_variable_get(method_info_key)
              # if we already have this information, then we don't need to set it as we'll update the info on access
              return if holder.key?(method)

              holder[method_name_key(method, is_instance_method)] = [method_policy, cs_method]
            end

            private

            def method_info_key
              :@cs__patched_method_info
            end

            # holder for each method - instance combination we've seen, done to
            # avoid rebuilding the key on every method invocation.
            def translated_names
              @_translated_names ||= {}
            end

            # @param method [Symbol] the name of the method to which the patch
            #   applies
            # @param is_instance_method [Boolean] is the method being patched
            #   an instance method or not (not implying class method).
            # @return [Symbol] :"method[true|false]"
            def method_name_key method, is_instance_method
              translated_names[method] = [] unless translated_names.key?(method)
              pre_built = translated_names[method]
              idx = is_instance_method ? 0 : 1
              pre_built[idx] = :"#{ method }#{ is_instance_method }" unless pre_built[idx]
              pre_built[idx]
            end

            def status_key
              :CONTRAST_PATCH_POLICY_STATUS
            end

            def update_holder mod, method, ret
              mod.instance_variable_set(method_info_key, {}) unless mod.instance_variable_defined?(method_info_key)
              holder = mod.instance_variable_get(method_info_key)
              holder[method] = ret
            end

            def find_info_for candidates, method
              candidates.each do |candidate|
                new_ret = find_info_for_one(candidate, method)
                return new_ret if new_ret
              end
              nil
            end

            def find_info_for_one mod, method
              return unless mod.instance_variable_defined?(method_info_key)

              holder = mod.instance_variable_get(method_info_key)
              return unless holder&.key?(method)

              holder[method]
            end

            MOD_ARRAY = [Module].cs__freeze
          end

          attr_reader :patch_status, :rewrite_status

          def patching!
            @patch_status = :PATCHING
          end

          def partial_patch!
            @patch_status = :PARTIAL
          end

          def no_patch!
            @patch_status = :NONE
          end

          def patched!
            @patch_status = :PATCHED
          end

          def failed_patch!
            @patch_status = :FAILED
          end

          def patching?
            @patch_status == :PATCHING
          end

          def patched?
            @patch_status == :PATCHED || @patch_status == :NONE || @patch_status == :FAILED
          end

          def rewriting!
            @rewrite_status = :REWRITING
          end

          def no_rewrite!
            @rewrite_status = :NO_REWRITE
          end

          def rewritten!
            @rewrite_status = :REWRITTEN
          end

          def failed_rewrite!
            @rewrite_status = :FAILED_REWRITE
          end

          def rewriting?
            @rewrite_status == :REWRITING
          end

          def rewritten?
            @rewrite_status == :REWRITTEN || @rewrite_status == :NO_REWRITE || @rewrite_status == :FAILED_REWRITE
          end
        end
      end
    end
  end
end