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

cs__scoped_require 'ipaddr'
cs__scoped_require 'json'
cs__scoped_require 'rack'
cs__scoped_require 'contrast/core_extensions/object'
cs__scoped_require 'contrast/security_exception'
cs__scoped_require 'contrast/utils/object_share'
cs__scoped_require 'contrast/utils/gemfile_reader'
cs__scoped_require 'contrast/utils/operating_environment'
cs__scoped_require 'contrast/agent/service_heartbeat'
cs__scoped_require 'contrast/components/interface'
cs__scoped_require 'contrast/utils/heap_dump_util'

cs__scoped_require 'contrast/utils/path_util'
cs__scoped_require 'contrast/utils/timer'
cs__scoped_require 'contrast/utils/freeze_util'
cs__scoped_require 'contrast/utils/service_sender_util'
cs__scoped_require 'contrast/utils/service_response_util'

module Contrast
  module Agent
    # This class allows the Agent to plug into the Rack middleware stack,
    # providing hooks to relevant application events, such as request start and
    # application code.
    class Middleware
      include Contrast::Components::Interface
      access_component :contrast_service, :logging, :agent, :analysis, :config, :scope

      attr_reader :app

      LOG_DEBUG_MIDDLEWARE_ENV     = 'middleware: log environment'
      LOG_DEBUG_MIDDLEWARE_LIB     = 'middleware: instrument shared libraries'
      LOG_DEBUG_MIDDLEWARE_START   = 'middleware: startup agent'
      LOG_DEBUG_MIDDLEWARE_SERVICE = 'middleware: starting service'

      # allows the Agent to function as a middleware
      #
      # @param app [Rack::Application] the application to be instrumented
      # @param _legacy_param [nil] was a flag we no longer need, but Sinatra may call it
      def initialize app, _legacy_param = nil
        @app = app

        # TODO: RUBY-545 nomenclature here needs to be updated.
        #   enabled/initialized are really only reflective of whether we were
        #   able to parse the config, which is the bare minimum for the agent
        #   to do anything.
        unless AGENT.enabled?
          logger.error('Contrast middleware initializer detected an early-stage setup failure (likely config parse).  Disabling.')
          # ensure the agent is disabled (probably redundant)
          AGENT.disable!
          return
        end

        logger.debug_with_time(LOG_DEBUG_MIDDLEWARE_ENV) do
          settings.log_environment
          settings.log_configuration
          settings.log_specific_libraries
          settings.log_all_libraries
        end

        # Initialization order should be:
        # - start separate service sending threads
        # - start heartbeat thread, which triggers service startup
        # - start patching to achieve instrumentation

        logger.debug_with_time(LOG_DEBUG_MIDDLEWARE_SERVICE) do
          # get threads ready to poll for messages on the queue
          run_service_sender_thread

          # sends first message to service, which triggers service startup
          run_service_heartbeat_thread
        end

        # Default instrumentation
        logger.debug_with_time(LOG_DEBUG_MIDDLEWARE_LIB) do
          AGENT.run_instrumentation
        end

        Contrast::Agent::AtExitHook.exit_hook
      end

      def settings
        Contrast::Agent::FeatureState.instance
      end

      LOG_DEBUG_MIDDLEWARE_START_CALL = 'middleware: startup agent (call)'
      # def _call env
      def call env
        Contrast::Utils::HeapDumpUtil.run

        if AGENT.enabled?
          response = call_with_agent(env)
          with_contrast_scope { do_static_analysis_catchup }
          response
        else
          call_without_agent(env)
        end
      end

      def call_without_agent env
        app.call(env)
      end

      REQUEST_PATH = 'REQUEST_PATH'
      LOG_DEBUG_REQUEST = 'HTTP request cycle'
      def call_with_agent env
        rack_request = generate_request(env)

        context = Contrast::Agent::RequestContext.new(rack_request, true) # 'true' here is legacy, by the time we're here, we're ready

        # default response
        streaming = false
        # make the context available for the lifecycle of this request
        Contrast::Agent::REQUEST_TRACKER.lifespan(context) do
          # record entire time
          logger.debug_with_time(LOG_DEBUG_REQUEST) do
            # process filters that can short circuit application processing
            with_contrast_scope do
              context.service_extract_request
              prefilter(context)
            end

            # run application by passing request down the Rack chain using
            # the original env
            response = application_code(env)

            # if streaming, allow for early return with application response
            streaming = possibly_streaming?(env)
            if streaming
              with_contrast_scope { postfilter(context, streaming) }
              return response
            end

            # update context with final response information
            context.extract_after(response)

            # process filters that look at response headers and body
            with_contrast_scope { postfilter(context) }
          end
        end

        # return the response stored in the context in case any postfilter rules
        # updated the response data
        context&.response&.rack_response || response

      # handle security exception
      rescue StandardError => e
        handle_exception(e)
      ensure
        begin
          handle_ensure(context, streaming)
        rescue Exception => e # rubocop:disable Lint/RescueException
          logger.error(e, 'Exception raised while flushing messages to Contrast service!')
          raise
        end
      end

      LOG_WARN_FRAMEWORK_PARSE = 'Unable to generate framework specific request - falling back to Rack'
      # Given the rack environment of this call, generate a framework specific
      # request object to bind to this context. In the event that multiple
      # supported frameworks are defined OR we cannot determine the framework
      # currently in use or there is an exception during request generation, we
      # will fall back on Rack::Request object
      def generate_request env
        rails_defined = defined?(Rails)
        sinatra_defined = defined?(Sinatra) && defined?(Sinatra::Request)

        if rails_defined && !sinatra_defined
          # code from /lib/action_cable/connection/base
          environment = Rails.application.env_config.merge(env) if defined?(Rails.application) && Rails.application
          ActionDispatch::Request.new(environment || env)
        # !defined? currently redundant, won't be if we have more frameworks
        elsif sinatra_defined && !rails_defined
          Sinatra::Request.new(env)
        else # many OR none
          Rack::Request.new(env)
        end
      rescue StandardError => e
        logger.warn(e, LOG_WARN_FRAMEWORK_PARSE)
        Rack::Request.new(env)
      end

      LOG_ERROR_STREAM_CHECK = 'Unable to check for streaming'
      ACTION_CONTROLLER_INSTANCE = 'action_controller.instance'
      # First check to see if an action could potentially be streaming a response
      def possibly_streaming? env
        return false unless defined?(ActionController::Live)

        env[ACTION_CONTROLLER_INSTANCE].cs__class.included_modules.include?(ActionController::Live)
      rescue StandardError => e
        logger.warn(LOG_ERROR_STREAM_CHECK, e)
      end

      LOG_ERROR_ENSURE = 'Context not defined in middleware ensure.'
      # Send a server activity message and an application activity message
      # and return response
      def handle_ensure context, streaming
        if context
          logger.debug('Middleware request lifecycle complete; flushing context activity to Contrast service.')
          Contrast::Utils::GemfileReader.instance.generate_library_usage(context.activity)
          [context.server_activity, context.activity, context.observed_route].each do |message|
            CONTRAST_SERVICE.send_message message
          end
        else
          logger.error(LOG_ERROR_ENSURE)
        end
        return unless streaming

        context.reset_activity
      end

      SECURITY_EXCEPTION_MARKER = 'Contrast::SecurityException'
      LOG_RETHROW_ERROR = 'Re-throwing original error'
      # We're only going to suppress SecurityExceptions indicating a blocked attack.
      # And, only if the config.agent.ruby.exceptions.capture? is set
      def handle_exception exception
        if exception.is_a?(Contrast::SecurityException) ||
              exception.message&.include?(SECURITY_EXCEPTION_MARKER)

          exception_control = Contrast::Agent::FeatureState.instance.exception_control
          raise exception unless exception_control[:enable]

          [exception_control[:status], {}, [exception_control[:message]]]
        else
          # Log original re-raise
          logger.debug(exception, LOG_RETHROW_ERROR)
          raise exception
        end
      end

      LOG_DEBUG_PREFILTER = 'prefilter'
      LOG_ERROR_PREFILTER = 'Unexpected exception during prefilter'
      # Iterate through rules that only depend upon the request object.
      def prefilter context
        return unless context.app_loaded?

        logger.debug_with_time(LOG_DEBUG_PREFILTER) do
          prefilter_assess(context) if ASSESS.enabled? && context.analyze_request?
          prefilter_protect(context) if PROTECT.enabled?
        end
      rescue Contrast::SecurityException => e
        logger.warn("RASP threw security exception in prefilter: exception=#{ e.message }")
        raise e
      rescue StandardError => e
        logger.error(LOG_ERROR_PREFILTER, e)
      end

      def prefilter_assess context
        rules = ASSESS.rules
        logger.debug("Assess: Running #{ rules.length } rules in prefilter.")
        prefilter_rules(rules, context)
      end

      def prefilter_protect context
        rules = PROTECT.rules
        logger.debug("Protect: Running #{ rules.length } rules in prefilter.")
        prefilter_rules(rules, context)
      end

      def prefilter_rules rules, context
        rules.each do |_, rule|
          rule.prefilter(context)
        end
      end

      LOG_DEBUG_APPLICATION = 'application'
      # Run the next level of the Rack stack as normal
      def application_code env
        logger.debug_with_time(LOG_DEBUG_APPLICATION) do
          app.call(env)
        end
      rescue Contrast::SecurityException => e
        logger.info("RASP threw security exception in application code: exception=#{ e.message }")
        raise e
      end

      LOG_DEBUG_POSTFILTER = 'postfilter'
      LOG_ERROR_POSTFILTER = 'Unexpected exception during postfilter'
      # Iterate through rules that depend on the full response object
      def postfilter context, streaming = false
        return unless context.app_loaded?

        logger.debug_with_time(LOG_DEBUG_POSTFILTER) do
          if ASSESS.enabled? && context.analyze_response?
            logger.debug("Assess:\tRunning #{ ASSESS.rules.length } rules in postfilter.")
            postfilter_rules(ASSESS.rules, context, streaming)
          end

          if PROTECT.enabled?
            logger.debug("Protect:\tRunning #{ PROTECT.rules.length } rules in postfilter.")
            postfilter_rules(PROTECT.rules, context, streaming)
          end
        end
      rescue Contrast::SecurityException => e
        logger.warn("RASP threw security exception: exception=#{ e.message }")
        raise e
      rescue StandardError => e
        logger.error(LOG_ERROR_POSTFILTER, e)
      end

      # Iterate through a list of rules with the current context and the full response body
      def postfilter_rules rules, context, streaming
        rules.each do |_, rule|
          if !streaming || rule.stream_safe?
            rule.postfilter(context)
          else
            logger.debug("Skipping rule: #{ rule.name } in streamed response.")
          end
        end
      end

      def run_service_heartbeat_thread
        # Rspec stubs over this method for simplicity's sake in testing.
        # Take care if you refactor this back into #initialize.
        Contrast::Agent::ServiceHeartbeat.new.start
      end

      def run_service_sender_thread
        # Rspec stubs over this method for simplicity's sake in testing.
        # Take care if you refactor this back into #initialize.
        Contrast::Utils::ServiceSenderUtil.start
      end

      LOG_DEBUG_MW_INV = 'middleware: send inventory'
      LOG_WARN_STATIC_ANALYSIS = 'Unable to run post-initialization static analysis'
      # this is memoized, should only be meaningful after the first agented
      # HTTP request
      def do_static_analysis_catchup
        @_do_static_analysis_catchup ||= begin
                                           # Everything in here should be asynchronous, as we need to return
                                           # the HTTP response from middleware ASAP.

                                           # Review already-loaded inventory
                                           logger.debug_with_time('initializer: report loaded gemset') do
                                             Contrast::Utils::GemfileReader.instance.map_loaded_classes
                                           end

                                           # Report already-loaded inventory
                                           logger.debug_with_time(LOG_DEBUG_MW_INV) do
                                             settings.send_inventory_message
                                           end

                                           true
                                         end
      rescue StandardError => e
        logger.warn(LOG_WARN_STATIC_ANALYSIS, e)
      end
    end
  end
end