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

# intentional -- we're using a << operator here

cs__scoped_require 'contrast/agent/class_reopener'
cs__scoped_require 'contrast/agent/patching/policy/patch_status'
cs__scoped_require 'contrast/components/interface'
cs__scoped_require 'contrast/utils/ruby_ast_rewriter'

module Contrast
  module Agent
    # Used for Ruby 2.4 & 2.5 to allow us to rewrite those methods which have
    # interpolation in them.
    # @deprecated Changes to this class are discouraged as this approach is
    #   being phased out with support for those language versions.
    class Rewriter
      include Contrast::Components::Interface
      access_component :logging, :scope

      SELF_DEFINITION = 'def self.'
      DEFINITION = 'def '

      class << self
        def rewrite_class module_data, redo_rewrite = false
          mod = module_data.mod
          return unless mod

          status = Contrast::Agent::Patching::Policy::PatchStatus.get_status(mod)
          return if (status.rewritten? || status.rewriting?) && !redo_rewrite

          status.rewriting!

          # default-initialize to nil within top-level method scope
          # so that it's available w/in the `ensure` block
          opener = nil
          with_contrast_scope do
            unless should_rewrite?(module_data)
              status.no_rewrite!
              return
            end

            opener = Contrast::Agent::ClassReopener.new(module_data)
            if opener.nil? || (instance_methods.empty? && singleton_methods.empty?)
              status.no_rewrite!
              return
            end

            rewrite_all_methods(opener, mod, mod.public_instance_methods(false), :PUBLIC_INSTANCE)
            rewrite_all_methods(opener, mod, mod.protected_instance_methods(false), :PROTECTED_INSTANCE)
            rewrite_all_methods(opener, mod, mod.private_instance_methods(false), :PRIVATE_INSTANCE)
            rewrite_all_methods(opener, mod, mod.singleton_methods(false), :SINGLETON)
            status.rewritten!
          end
        rescue SyntaxError, StandardError => e
          opener = nil
          mod ||= module_data.mod
          logger.debug(e, "Reopening #{ mod } threw a handled exception - skipping rewriting")
          status ||= Contrast::Agent::Patching::Policy::PatchStatus.get_status(mod)
          status.failed_rewrite!
        ensure
          with_contrast_scope do
            opener&.commit_patches
          end
          logger.debug(
              nil,
              "Rewriting #{ module_data.name } resulted in #{ Contrast::Agent::Patching::Policy::PatchStatus.get_status(module_data.mod).rewrite_status }")
        end

        private

        def rewrite_all_methods opener, clazz, methods, type
          methods.each do |method|
            # Skip contrast woven methods.
            # There should be a better way to do this
            next if method.to_s.start_with?('cs__')

            original_source_code = source_code(opener, clazz, method, type)
            next if original_source_code.nil?
            next if unrepeatable?(original_source_code)
            next unless interpolations?(original_source_code)

            replace_method_definition(opener, clazz, method, original_source_code, type)
          end
        end

        def source_code opener, clazz, method_name, type
          return unless method_exists?(clazz, method_name, type)

          method_instance = method_instance(clazz, method_name, type)
          return nil if method_instance.nil?

          location = method_instance.source_location
          return nil if location.nil?
          return nil if location.empty? || location[0].empty? || location[0].include?('eval')
          return nil if opener.written_from_location?(location)

          opener.written_from_location!(location)
          opener.source_code(location, method_name)
        rescue SyntaxError
          logger.debug(nil, "SyntaxError: Can't parse method source from #{ clazz }##{ method_name }")
        rescue StandardError => e
          if defined?(MethodSource) && defined?(MethodSource::SourceNotFoundError)
            logger.debug(nil, "SourceNotFoundError: Can't parse method source from #{ clazz }##{ method_name }")
          else
            logger.debug(e, "Method source lookup of #{ clazz }##{ method_name } failed")
          end
        end

        def method_exists? clazz, method_name, type
          case type
          when :SINGLETON
            clazz.cs__singleton_class.public_method_defined?(method_name)
          when :PUBLIC_INSTANCE
            clazz.public_method_defined?(method_name)
          when :PROTECTED_INSTANCE
            clazz.protected_method_defined?(method_name)
          when :PRIVATE_INSTANCE
            clazz.private_method_defined?(method_name)
          end
        end

        def method_instance clazz, method_name, type
          case type
          when :SINGLETON
            clazz.singleton_method(method_name)
          else
            clazz.instance_method(method_name)
          end
        end

        def replace_method_definition opener, clazz, method, original_source_code, type
          new_method_source = rewrite_method(original_source_code)
          return unless valid_code?(new_method_source)

          case type
          when :SINGLETON
            opener.public_singleton_methods << new_method_source
          when :PUBLIC_INSTANCE
            opener.public_instance_methods << new_method_source
          when :PROTECTED_INSTANCE
            opener.protected_instance_methods << new_method_source
          when :PRIVATE_INSTANCE
            opener.private_instance_methods << new_method_source
          end
        rescue StandardError => e
          logger.debug(e, "Error in rewriter class_eval for of #{ clazz }##{ method }")
        end

        def rewrite_method original_source_code
          new_code = rewrite(original_source_code)
          new_code = remove_self_definition(new_code)
          new_code
        end

        def rewrite source
          @rewriter ||= Contrast::Utils::RubyAstRewriter.new
          @rewriter.rewrite(source)
        end

        def remove_self_definition new_code
          new_code.gsub(SELF_DEFINITION, DEFINITION)
        end

        # attempt to generate an AST for the rewritten code, if one can be generated we can assume it is at least valid ruby
        # if it returns nil that means that the AST was unable to build due to an error and we should use the original source code
        def valid_code? new_content
          !Ripper.sexp(new_content).nil?
        end

        def interpolations? method_source
          method_source.match?(/^.*".*#(\$\w+|@\w+|@@\w+|\{.*\}).*".*$/)
        end

        # These functions cannot be repeated -- when doing replace things, we
        # should comment these lines out.
        # cannot repeat class foo <
        #                         bar
        UNREPEATABLE_FUNCTIONS = /(^|\s)(class\s*\W.*<|included(\s+do|\s*{))/.cs__freeze

        # these methods dynamically create functions. replacing them in rails
        # resulted in a lot of undefined_ errors. i know we'll have to figure
        # out how to support them, but baby steps
        UNREPEATABLE_ACCESSOR_FUNCTIONS =
          /(^|\s)((cattr_accessor|cattr_reader|cattr_writer|attr_reader|attr_writer|attr_accessor|mattr_accessor|mattr_reader|mattr_writer|public) )/.cs__freeze
        # evals are scary and terrify me. we shouldn't rewrite them if we don't
        # have to
        POTENTIALLY_UNREPEATABLE_FUNCTIONS =
          /(class_eval|instance_eval|method_eval|eval|module_eval|define_method|define_method_attribute|require_dependency|instance_predicate)/.cs__freeze
        # Binding a Constant, ie starting with ::, affects the lookup of that
        # Constant. Since we're effectively moving the file in which it would
        # be declared, we cannot repeat this lookup, so we cannot replace this
        # file
        BOUND_CONSTANTS = /(^|\s)::\w/.cs__freeze

        DISABLE_FUNCTIONS_REGEXP = /(^|\s)((undef|include|extend|alias|alias_method|require|require_relative|send) )/.cs__freeze
        def unrepeatable? original_source_code
          (original_source_code.index('__dir__') ||
            original_source_code.index('File.dirname') ||
            original_source_code.index('__FILE__') ||
            original_source_code.match?(UNREPEATABLE_FUNCTIONS) ||
            original_source_code.match?(UNREPEATABLE_ACCESSOR_FUNCTIONS) ||
            original_source_code.match?(POTENTIALLY_UNREPEATABLE_FUNCTIONS) ||
            original_source_code.match?(BOUND_CONSTANTS) ||
            original_source_code.match?(DISABLE_FUNCTIONS_REGEXP))
        end

        UNTOUCHABLE_MODULES = %w[
          AbstractController
          ActionCable ActionController ActionDispatch ActionMailer ActionPack ActionView
          ActiveModel ActiveRecord ActiveSupport
          Brakeman
          Contrast
          Debase
          Debugger
          Erubis
          GraphQL
          NewRelic
          Pry
          Puma
          PhusionPassenger
          Rails::Application Rails::Engine Rails::Railtie
          Rake
          RSpec
          Warden
          WEBrick
        ].cs__freeze
        def should_rewrite? module_data
          clazz = module_data.mod
          name = module_data.name
          return false unless clazz

          # Name can be nil for anonymous modules. We won't work on them.
          return false unless name
          return false if name == Contrast::Utils::ObjectShare::CLASS
          return false if name == Contrast::Utils::ObjectShare::MODULE

          # We've ran into issues rewriting these Modules. Ideally, we'll solve
          # the issues and this check will go away
          UNTOUCHABLE_MODULES.none? { |untouchable| name.include?(untouchable) }
        end
      end
    end
  end
end