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

require 'ffi'
# require the gem
require 'contrast-agent-lib'
require 'contrast/components/logger'

module Contrast
  module AgentLib
    # This module is defined in Rust as external, we used it here.
    # Initializes the AgentLib. Here will be all methods from
    # the C bindings  contrast_c::input_tracing module.
    module InputTracing
      # Struct could be used as parameter to C functions,
      # and also to save the return data. This is like
      # specifying a type for pointer.
      #
      # Expected struct:
      # pub struct CCheckQuerySinkResult {
      #     pub start_index: u64,
      #     pub end_index: u64,
      #     pub boundary_overrun_index: u64,
      #     pub input_boundary_index: u64,
      # }
      #
      class QuerySinkResult < FFI::ManagedStruct
        layout :start_index,            :ulong_long,
               :end_index,              :ulong_long,
               :boundary_overrun_index, :ulong_long,
               :input_boundary_index,   :ulong_long

        # when QuerySinkResult object gets out of scope it will be GCed.
        def self.release ptr
          Contrast::AgentLib::InputTracing.dl__free_check(ptr)
        end
      end

      # This class will be used to hold the return type from input tracings
      # done with the AgentLib. Using the ManagedStruct will automatically
      # call the release function, and since the memory is allocated by the
      # AgentLib is should be GC also there.
      class CEvalResult < FFI::ManagedStruct
        layout :rule_id, :ulong_long, :input_type, :ulong_long, :score, :double

        def self.release ptr
          Contrast::AgentLib::InputTracing.dl__free_eval(ptr)
        end
      end

      extend FFI::Library
      include Contrast::Components::Logger::InstanceMethods
      ffi_lib ContrastAgentLib::CONTRAST_C
      ffi_convention :stdcall

      # Returns 0 if evaluation was successful, -1 if there was an unexpected error.
      attach_function :check_cmd_injection_query, %i[int int string pointer], :int
      # Return 0 if evaluation was successful, -1 iif there was an unexpected error.
      attach_function :check_sql_injection_query, %i[int int int string pointer], :int
      # Free function for all input tracing results pass to check functions.
      attach_function :free_check_query_sink_result, [:pointer], :int
      # Evaluate header input:
      attach_function :evaluate_header_input, %i[string string uint64 uint64 pointer pointer], :int
      # Evaluate input:
      attach_function :evaluate_input, %i[string uint64 uint64 uint64 pointer pointer], :int
      # Free input evaluation:
      attach_function :free_eval_result, [:pointer], :int
      # function init
      attach_function :init, [], :int

      # Once we have pointer or returned data the lib
      # expect us to free it, so it could GCd.
      # Most of the time we would receive result pointer.
      # Used with dl__check_cmdi_query and dl__check_sqli_query.
      #
      # @param result [Pointer]
      def self.dl__free_check result
        free_check_query_sink_result(result)
      end

      # Free the input evaluation result received from
      # dl__eval_input, and dl__eval_header_input
      #
      # @param result [Pointer]
      def self.dl__free_eval result
        free_eval_result(result)
      end

      private

      # Evaluate a subset of executing shell command for token boundaries being crossed.
      # This is used by the cmd-injection rule during sink time. If a previously
      # worth-watching input is contained in the shell command it crosses token
      # boundaries then an injection should be raised.
      #
      # @param input_index [Integer] index in the cmdText string where user input was found.
      # @param input_length [Integer] length of the user input.
      # @return result [Hash, nil] output parameter which will contain the result of the evaluation.
      def dl__check_cmdi_query input_index, input_length, input_cmd
        # The memory for the result is allocated by AgentLib/Rust.  So we need to create just an empty pointer
        # This pointer will be populated by the output parameter of check_cmd_injection_query#
        result_ptr = FFI::MemoryPointer.new(:pointer)
        result_code = check_cmd_injection_query(input_index, input_length, input_cmd, result_ptr)

        unless result_code == -1
          struct = QuerySinkResult.new(result_ptr.read_pointer)
          result_hash = dl__struct_to_result(struct)

          return result_hash
        end

        nil
      end

      # This function will call the check sql injection query
      # SAFETY DISCLAIMER
      # The sql_query parameter must point to an allocated C-string in UTF-8 encoding
      # The result MUST be free'd using free_check_query_sink_result as soon as result has been processed.
      #
      # @param input_index[Integer] index in the sqlQuery string where user input was found
      # @param input_length[Integer] length of the user input
      # @param sql_query[String] full SQL query being evaluated
      # @param db_type[Integer] database engine being evaluated. Must be one of the DbType enum values
      # @return [Hash, nil] Returns 0 if evaluation was successful, -1 if there was an unexpected error.
      def dl__check_sql_injection_query input_index, input_length, db_type, sql_query
        db_type = :Sqlite if db_type.to_s.casecmp('sqlite3').zero?
        dbtype = Contrast::AGENT_LIB.db_set[db_type.to_s.upcase.to_sym]
        pointer = FFI::MemoryPointer.new(:pointer)
        check = check_sql_injection_query(input_index, input_length, dbtype, sql_query, pointer)

        return if check == -1

        struct = QuerySinkResult.new(pointer.read_pointer)
        dl__struct_to_result(struct)
      rescue RuntimeError => e
        # silence the runtime error from the agent lib
        logger.debug('Following RuntimeError was recording during input tracing: ', e: e)
        nil
      ensure
        pointer.free unless pointer.nil? && pointer.address
      end

      # Evaluates a header for input tracing rules.
      #
      # @param header_name [String] key/name of the header to be evaluated.
      # NOTE: the only header names used are "custom", "accept", and "user-agent".
      # Header-name-used is part of the output because the accept and user-agent
      # headers get special handling by agent-lib, so they are evaluated twice -
      # once exactly like all other headers and once with the specific header that
      # gets special handling.
      # @param header_value [String] header value
      # @param rule_set [Integer] one or more bit flag of input tracing rules to be evaluated.
      # The bit flags must match the RuleType enum. Currently supported types:
      # [RuleType::CmdInjection, RuleType::PathTraversal, RuleType::SqlInjection,
      # RuleType::UnsafeFileUpload, RuleType::ReflectedXss, RuleType::BotBlocker]
      #     public enum RuleType : ulong
      #     {
      #         UnsafeFileUpload = 1 << 0,
      #         PathTraversal = 1 << 1,
      #         ReflectedXss = 1 << 2,
      #         SqlInjection = 1 << 3,
      #         CmdInjection = 1 << 4,
      #         NosqlInjectionMongo = 1 << 5,
      #         BotBlocker = 1 << 6
      #     }
      # @param eval_options [Integer] u64 representation of the EvalOptions enum.
      #     public enum EvalOptions : ulong
      #     {
      #         None = 0,
      #         PreferWorthWatching = 1 << 0,
      #     }
      def dl__eval_header_input header_name, header_value, rule_set, eval_options
        # The memory for the result is allocated by AgentLib/Rust.  So we need to create just an empty pointer
        # This pointer will be populated by the output parameter of check_cmd_injection_query#
        result_ptr = FFI::MemoryPointer.new(:pointer)
        result_size = FFI::MemoryPointer.new(:uint32)
        # @param result_size [Integer] output parameter which will contain the number of Results after evaluation:
        # 0 if no result.
        evaluate_header_input(header_name, header_value, rule_set, eval_options, result_size, result_ptr)

        size = result_size.read_int
        result_size&.free

        unless size.zero?
          struct = CEvalResult.new(result_ptr.read_pointer)
          return dl__eval_struct_to_result(struct)
        end
        nil
      end

      # Evaluates an input part for input tracing rules. This should be used for all input parts except headers.
      #
      # @param input [String] input value to be evaluated.
      # @param input_type [Symbol]  type of the input. This needs to be a u64 representation of the
      # supported InputType enum values:
      #        {
      #         CookieName = 1,
      #         CookieValue = 2,
      #         HeaderKey = 3,
      #         HeaderValue = 4,
      #         JsonKey = 5,
      #         JsonValue = 6,
      #         Method = 7,
      #         ParameterKey = 8,
      #         ParameterValue = 9,
      #         UriPath = 10,
      #         UrlParameter = 11,
      #         MultipartName = 12,
      #         XmlValue = 13
      #        }
      # @param rule_set [Integer] one or more bit flag of input tracing rules to be evaluated.
      # The bit flags must match the RuleType enum. Currently supported types:
      # [RuleType::CmdInjection, RuleType::PathTraversal, RuleType::SqlInjection,
      # RuleType::UnsafeFileUpload, RuleType::ReflectedXss, RuleType::BotBlocker]
      # @param eval_options [Integer] u64 representation of the EvalOptions enum.
      #     public enum EvalOptions : ulong
      #     {
      #         None = 0,
      #         PreferWorthWatching = 1 << 0,
      #     }
      def dl__eval_input input, input_type, rule_set, eval_options
        # The memory for the result is allocated by AgentLib/Rust.  So we need to create just an empty pointer
        # This pointer will be populated by the output parameter of check_cmd_injection_query#
        result_ptr = FFI::MemoryPointer.new(:pointer)
        result_size = FFI::MemoryPointer.new(:uint32)
        # @param result_size [Integer] output parameter which will contain the number of Results after evaluation:
        # 0 if no result.
        evaluate_input(input, input_type, rule_set, eval_options, result_size, result_ptr)
        size = result_size.read_int
        result_size&.free

        unless size.zero?
          struct = CEvalResult.new(result_ptr.read_pointer)
          return dl__eval_struct_to_result(struct)
        end
        nil
      end

      # Convert c_struct cmdi result to ruby hash
      #
      # @param struct [Contrast::AgentLib::QuerySinkResult]
      # @return hash [Hash]
      def dl__struct_to_result struct
        hash = {}
        hash[:start_index] = struct[:start_index]
        hash[:end_index] = struct[:end_index]
        hash[:boundary_overrun_index] = struct[:boundary_overrun_index]
        hash[:input_boundary_index] = struct[:input_boundary_index]
        hash
      end

      # Convert c_struct evaluated input to ruby hash
      #
      # @param struct [Contrast::AgentLib::QuerySinkResult]
      # @return hash [Hash]
      def dl__eval_struct_to_result struct
        hash = {}
        hash[:rule_id] = struct[:rule_id]
        hash[:input_type] = struct[:input_type]
        hash[:score] = struct[:score]
        hash
      end
    end
  end
end