# 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_telemetry_event'
require 'contrast/utils/middleware_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
      include Contrast::Components::Logger::InstanceMethods
      include Contrast::Components::Scope::InstanceMethods
      include Contrast::Utils::MiddlewareUtils

      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
      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
        return app.call(env) unless ::Contrast::AGENT.enabled?

        Contrast::Agent.heapdump_util.start_thread!
        handle_first_request
        call_with_agent(env)
      end
      ::Contrast::Components::Logger.add_trace_log_timing_for(::Contrast::Agent::Middleware, :call)

      private

      # Startup the Agent as part of the initialization process:
      # - start the service sending thread, responsible for sending and processing messages
      # - start the heartbeat thread, which triggers service startup
      # - 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 service') do
          Contrast::Agent.thread_watcher.ensure_running?
        end

        if Contrast::Agent::Telemetry.enabled?
          logger.debug_with_time('middleware: sending startup metrics telemetry event') do
            event = Contrast::Agent::StartupMetricsTelemetryEvent.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]
      def pre_call_with_agent context, request_handler
        with_contrast_scope do
          context.service_extract_request
          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]
      def post_call_with_agent context, env, request_handler, response
        with_contrast_scope do
          context.extract_after(response) # update context with final response information

          # Build and report all collected findings prior response
          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?

          if Contrast::Agent.framework_manager.streaming?(env)
            context.reset_activity
            request_handler.stream_safe_postfilter
          else
            request_handler.ruleset.postfilter
            request_handler.report_activity
            request_handler.send_activity_messages # TODO: RUBY-1438 -- remove
          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