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

require 'ipaddr'
require 'json'
require 'rack'

require 'contrast/security_exception'
require 'contrast/utils/object_share'
require 'contrast/components/logger'
require 'contrast/components/scope'
require 'contrast/utils/heap_dump_util'
require 'contrast/utils/telemetry'
require 'contrast/agent/request_handler'
require 'contrast/agent/static_analysis'
require 'contrast/agent/telemetry/events/startup_metrics_event'
require 'contrast/agent/protect/input_analyzer/input_analyzer'
require 'contrast/utils/middleware_utils'
require 'contrast/utils/reporting/application_activity_batch_utils'
require 'contrast/utils/timer'

module Contrast
  module Agent
    # This class allows the Agent to plug into the Rack middleware stack. When the application is first started, we
    # initialize ourselves as a rack middleware inside of #initialize. Afterwards, we process each http request and
    # response as it goes through the middleware stack inside of #call.
    class Middleware # rubocop:disable Metrics/ClassLength
      include Contrast::Components::Logger::InstanceMethods
      include Contrast::Components::Scope::InstanceMethods
      include Contrast::Utils::MiddlewareUtils
      include Contrast::Utils::Reporting::ApplicationActivityBatchUtils

      attr_reader :app

      # Allows the Agent to function as a middleware. We perform all our one-time whole-app routines in here since
      # we're only going to be initialized a single time. Our initialization order is:
      # - capture the application
      # - setup the Agent
      # - startup the Agent
      #
      # @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  # THIS MUST BE FIRST AND ALWAYS SET!
        setup_agent # THIS MUST BE SECOND AND ALWAYS CALLED!
        unless ::Contrast::AGENT.enabled?
          logger.error('The Agent was unable to initialize before the application middleware was initialized. '\
                       'Disabling permanently.')
          ::Contrast::AGENT.disable! # ensure the agent is disabled (probably redundant)
          return
        end
        agent_startup_routine
      rescue StandardError => e
        logger.error('Unable to initialize the agent. Disabling permanently.', e)
        ::Contrast::AGENT.disable! # ensure the agent is disabled (probably redundant)
      end

      # This is where we're hooked into the middleware stack. If the agent is enabled, we're ready to do some
      # processing on a per request basis. If not, we just pass the request along to the next middleware in the stack.
      #
      # @param env [Hash] the various variables stored by this and other Middlewares to know the state and values of
      #   this Request
      # @return [Array,Rack::Response] the Response of this and subsequent Middlewares to be passed back to the user up
      #   the Rack framework.
      def call env
        logger.trace_with_time('Elapsed time for Contrast::Agent::Middleware#call') do
          ::Contrast::Agent::ThreadWatcher.check_before_start
          return app.call(env) unless ::Contrast::AGENT.enabled?

          Contrast::Agent.heapdump_util.start_thread!
          handle_first_request
          call_with_agent(env)
        end
      end

      private

      # Startup the Agent as part of the initialization process:
      # - start the TeamServer sending thread, responsible for sending and processing messages
      # - start the heartbeat thread, which handles periodic messages to TeamServer
      # - start instrumenting libraries and do a 'catchup' patch for everything we didn't see get loaded
      # - enable TracePoint, which handles all class loads and required instrumentation going forward
      def agent_startup_routine
        logger.debug_with_time('middleware: starting reporting threads') do
          Contrast::Agent.thread_watcher.ensure_running?
        end

        if Contrast::Agent::Telemetry::Base.enabled?
          logger.debug_with_time('middleware: sending startup metrics telemetry event') do
            event = Contrast::Agent::Telemetry::StartupMetricsEvent.new
            Contrast::Agent.thread_watcher.telemetry_queue.send_event(event)
          end
        end

        logger.debug_with_time('middleware: instrument shared libraries and patch') do
          Contrast::Agent::Patching::Policy::Patcher.patch
        end

        logger.debug_with_time('middleware: enabling tracepoint') do
          ::Contrast::AGENT.enable_tracepoint
        end
        Contrast::Agent::AtExitHook.exit_hook
      end

      # Some things have to wait until first request to happen, either because resolution is not complete or because
      # the framework will preload classes, which confuses some of our instrumentation.
      def handle_first_request
        @_handle_first_request ||= begin
          Contrast::Agent::StaticAnalysis.catchup
          true
        end
      end

      # This is where we process each request we intercept as a middleware. We make the request context available
      # globally so that it can be accessed from anywhere. A RequestHandler object is made for each request, which
      # handles prefilter and postfilter operations.
      # @param env [Hash] the various variables stored by this and other Middlewares to know the state and values of
      #   this Request
      # @return [Array,Rack::Response] the Response of this and subsequent Middlewares to be passed back to the user
      #   up the Rack framework.
      def call_with_agent env
        Contrast::Agent.thread_watcher.ensure_running?
        framework_request = Contrast::Agent.framework_manager.retrieve_request(env)
        context = Contrast::Agent::RequestContext.new(framework_request)
        response = nil

        # Make the context available for the lifecycle of this request.
        Contrast::Agent::REQUEST_TRACKER.lifespan(context) do
          logger.request_start
          request_handler = Contrast::Agent::RequestHandler.new(context)

          pre_call_with_agent(context, request_handler)
          response = application_code(env)
          post_call_with_agent(context, env, request_handler, response)
        ensure
          logger.request_end
        end

        response
      rescue StandardError => e
        handle_exception(e)
      end

      # Handle the operations the Agent needs to accomplish prior to the Application code executing during this
      # request.
      #
      # @param context [Contrast::Agent::RequestContext]
      # @param request_handler [Contrast::Agent::RequestHandler]
      # @raise [StandardError] raises an error if the exception is security concern
      #    which is being triggered when there is a failure within the pre-call with the agent
      def pre_call_with_agent context, request_handler
        with_contrast_scope do
          context.protect_input_analysis
          request_handler.ruleset.prefilter
        end
      rescue StandardError => e
        raise(e) if security_exception?(e)

        logger.error('Unable to execute agent pre_call', e)
      end

      # Handle the operations the Agent needs to accomplish after the Application code executes during this request.
      #
      # @param context [Contrast::Agent::RequestContext]
      # @param env [Hash] the various variables stored by this and other Middlewares to know the state and values of
      #   this Request
      # @param request_handler [Contrast::Agent::RequestHandler]
      # @param response [Array,Rack::Response]
      # @raise [StandardError] raises an error if the exception is security concern
      #    which is being triggered when there is a failure within the post-call with the agent
      def post_call_with_agent context, env, request_handler, response
        with_contrast_scope do
          context.extract_after(response) # update context with final response information

          Contrast::Agent::FINDINGS.report_collected_findings unless Contrast::Agent::FINDINGS.collection.empty?
          # All protect rules, which are trigger but require response to be reported
          Contrast::Agent::EXPLOITS.report_recorded_exploits(context) unless Contrast::Agent::EXPLOITS.collection.empty?
          # Process Worth Watching Inputs for v2 rules
          Contrast::Agent.worth_watching_analyzer&.add_to_queue(context.agent_input_analysis)
          # Now we can build the ia_results only for postfilter rules.
          context.protect_postfilter_ia

          if Contrast::Agent.framework_manager.streaming?(env)
            context.reset_activity
            request_handler.stream_safe_postfilter
          else
            request_handler.ruleset.postfilter
            request_handler.report_observed_route
            add_activity_to_batch(context.activity)
            report_batch
          end
        end
      # unsuccessful attack
      rescue StandardError => e
        raise(e) if security_exception?(e)

        logger.error('Unable to execute agent post_call', e)
      end
    end
  end
end