module Msf::Exploit::SQLi
  class Common
    include Msf::Module::UI

    #
    #   Creates an instance of an SQL Injection object, users should use the create_dbms method of Msf::Exploit::SQLi instead
    #
    #   @param datastore [DataStore]
    #   @param framework [Framework]
    #   @param user_output [Rex::Ui::Text::Output::Stdio] 
    #   @param opts [Hash] a dictionary containing the parameters needed
    #   @option opts [Integer] truncation_length :  [Optional] The number of characters returned, if the query result is truncated
    #   @option opts [String] concat_separator : [Optional] The separator to use when concatenating rows (default ',')
    #   @option opts [String] second_concat_separator : [Optional] The separator to use when concatenating columns (default ';')
    #   @option opts [Boolean] safe : don't use group_concat, safer for large tables if group_concat truncates the result, but more queries will be performed
    #   @option opts [String] null_replacement : a string that will replace NULL values
    #   @option opts [Boolean] hex_encode_strings : encode strings as hex numbers, no quotes in the payload
    #   @option opts [Object] an encoder name, or a hash specifying a custom encoder, see Encoders in DBMS-specific classes
    #   @param query_proc [Proc] a block that will receive the payload, and should send the request to the target,
    #   - if it's a regular SQL injection, it should return the part of the response that is the query result (one row)
    #   - if it's a boolean-based blind SQL injection, it should return `true`, `false`, or a value that evaluates to one of them
    #   `true` if the query returned a result, false otherwise
    #   - if it's a time-based blind SQL injection, the return value does not matter, the time the block takes to run is used to leak information.
    #
    def initialize(datastore, framework, user_output, opts = {}, &query_proc)
      raise ArgumentError, 'Missing the block that does the requests' unless block_given?
      raise ArgumentError, 'Positional arguments can\'t be nil' if [datastore, framework, user_output].any?(&:nil?)

      check_opts(opts)
      @query_proc = query_proc
      @safe = opts[:safe]
      @concat_separator = opts[:concat_separator]
      @second_concat_separator = opts[:second_concat_separator] || ';'
      @null_replacement = opts[:null_replacement] || ''
      @truncation_length = opts[:truncation_length] if opts[:truncation_length] && opts[:truncation_length].is_a?(Integer) && opts[:truncation_length] > 0
      @hex_encode_strings = opts[:hex_encode_strings]
      @encoder = opts[:encoder]
      @datastore = datastore
      @framework = framework
      @user_output = user_output
    end

    #
    #   Queries the block with the given SQL query, without necessarily returning a result (needed for
    #   example when uploading a file using a time-based SQL injection, as it's not necessary to
    #   run multiple queries for that purpose), not to be overridden, it is guaranteed that the query
    #   will run only once.
    #   @param query [String] The SQL query to execute
    #   @return [void]
    #
    def raw_run_sql(query)
      vprint_status "{SQLi} Executing (#{query})"
      if @hex_encode_strings
        query = hex_encode_strings(query)
        vprint_status "{SQLi} Encoded to (#{query})"
      end
      @query_proc.call(query)
    end

    #
    #   Queries the block with the given SQL query, and returns the result, this method is overridden in
    #   blind SQL injection classes, implementing the logic of leaking one bit at a time, and working 
    #   exactly the same as this method.
    #   @param query [String] The SQL query to execute
    #   @return [String] The query results
    #
    def run_sql(query)
      raw_run_sql(query)
    end

    attr_reader :datastore, :framework
    attr_accessor :concat_separator, :second_concat_separator, :null_replacement, :truncation_length, :safe

    private

    #
    #  This method checks that the required options are present, and have valid values
    #  @param opts [Hash] the options provided by the user
    #  @return [void]
    #
    def check_opts(opts)
      raise ArgumentError, 'null_replacement option cannot contain single quotes' if opts[:null_replacement] && opts[:null_replacement].include?("'")
      raise ArgumentError, 'truncation_length must be an integer' if opts[:truncation_length] && !opts[:truncation_length].is_a?(Integer)
    end
  end
end