<commit_before>require 'securerandom'
require 'base64'

module Kafka
  SCRAM_SHA256 = 'SHA-256'.freeze
  SCRAM_SHA512 = 'SHA-512'.freeze
  class SaslScramAuthenticator
    def initialize(username, password, mechanism: SCRAM_SHA256, logger: nil)
      @username = username
      @password = password
      @mechanism = mechanism
      @logger = logger
    end

    def authenticate!(connection)
      @connection = connection
      response = @connection.send_request(Kafka::Protocol::SaslHandshakeRequest.new('SCRAM-' + @mechanism))

      unless response.error_code == 0 && response.enabled_mechanisms.include?('SCRAM-' + @mechanism)
        raise Kafka::SaslScramError, "SCRAM-#{@mechanism} is not supported."
      end

      log_debug "authenticating #{@username} with scram, mechanism: #{@mechanism}"

      @encoder = @connection.encoder
      @decoder = @connection.decoder

      msg = first_message
      log_debug "[scram] Sending client's first message: #{msg}"
      @encoder.write_bytes(msg)

      @server_first_message = @decoder.bytes
      log_debug "[scram] Received server's first message: #{@server_first_message}"

      msg = final_message
      log_debug "[scram] Sending client's final message: #{msg}"
      @encoder.write_bytes(msg)

      response = parse_response(@decoder.bytes)
      log_debug "[scram] Received server's final msg: #{response}"
      log_debug "[scram] Client calculated server signature: #{server_signature}"

      raise FailedScramAuthentication, response['e'] if response['e']
      raise FailedScramAuthentication, 'Invalid server signature' if response['v'] != server_signature
    rescue FailedScramAuthentication
      raise
    rescue StandardError => e
      @logger.error "authentication error #{e.inspect}\n\n#{e.backtrace.join("\n")}"
      raise FailedScramAuthentication, 'Authentication failed: Unknown reason'
    end

    private

    def log_debug(str)
      @logger.debug str if @logger
    end

    def first_message
      "n,,#{first_message_bare}"
    end

    def first_message_bare
      "n=#{encoded_username},r=#{nonce}"
    end

    def final_message_without_proof
      "c=biws,r=#{rnonce}"
    end

    def final_message
      "#{final_message_without_proof},p=#{client_proof}"
    end

    def server_data
      parse_response(@server_first_message)
    end

    def rnonce
      server_data['r']
    end

    def salt
      Base64.strict_decode64(server_data['s'])
    end

    def iterations
      server_data['i'].to_i
    end

    def auth_message
      [first_message_bare, @server_first_message, final_message_without_proof].join(',')
    end

    def salted_password
      hi(@password, salt, iterations)
    end

    def client_key
      hmac(salted_password, 'Client Key')
    end

    def stored_key
      h(client_key)
    end

    def server_key
      hmac(salted_password, 'Server Key')
    end
<commit_msg>No need for connection to be instance variable<commit_after>require 'securerandom'
require 'base64'

module Kafka
  SCRAM_SHA256 = 'SHA-256'.freeze
  SCRAM_SHA512 = 'SHA-512'.freeze
  class SaslScramAuthenticator
    def initialize(username, password, mechanism: SCRAM_SHA256, logger: nil)
      @username = username
      @password = password
      @mechanism = mechanism
      @logger = logger
    end

    def authenticate!(connection)
      connection = connection
      response = connection.send_request(Kafka::Protocol::SaslHandshakeRequest.new('SCRAM-' + @mechanism))

      unless response.error_code == 0 && response.enabled_mechanisms.include?('SCRAM-' + @mechanism)
        raise Kafka::SaslScramError, "SCRAM-#{@mechanism} is not supported."
      end

      log_debug "authenticating #{@username} with scram, mechanism: #{@mechanism}"

      @encoder = connection.encoder
      @decoder = connection.decoder

      begin
        msg = first_message
        log_debug "[scram] Sending client's first message: #{msg}"
        @encoder.write_bytes(msg)

        @server_first_message = @decoder.bytes
        log_debug "[scram] Received server's first message: #{@server_first_message}"

        msg = final_message
        log_debug "[scram] Sending client's final message: #{msg}"
        @encoder.write_bytes(msg)

        response = parse_response(@decoder.bytes)
        log_debug "[scram] Received server's final msg: #{response}"
        log_debug "[scram] Client calculated server signature: #{server_signature}"
        
        raise FailedScramAuthentication, response['e'] if response['e']
        raise FailedScramAuthentication, 'Invalid server signature' if response['v'] != server_signature
      rescue FailedScramAuthentication
        raise
      rescue => e
        @logger.error "authentication error #{e.inspect}\n\n#{e.backtrace.join("\n")}"
        raise FailedScramAuthentication, 'Authentication failed: Unknown reason'
      end
      @logger.debug "[scram] Authenticated"
    end

    private

    def log_debug(str)
      @logger.debug str if @logger
    end

    def first_message
      "n,,#{first_message_bare}"
    end

    def first_message_bare
      "n=#{encoded_username},r=#{nonce}"
    end

    def final_message_without_proof
      "c=biws,r=#{rnonce}"
    end

    def final_message
      "#{final_message_without_proof},p=#{client_proof}"
    end

    def server_data
      parse_response(@server_first_message)
    end

    def rnonce
      server_data['r']
    end

    def salt
      Base64.strict_decode64(server_data['s'])
    end

    def iterations
      server_data['i'].to_i
    end

    def auth_message
      msg = [first_message_bare, @server_first_message, final_message_without_proof].join(',')
    end

    def salted_password
      hi(@password, salt, iterations)
    end

    def client_key
      hmac(salted_password, 'Client Key')
    end

    def stored_key
      h(client_key)
    end

    def server_key
      hmac(salted_password, 'Server Key')
    end

{'fraction_non_alphanumeric': 0.10038805466509196, 'fraction_numerical': 0.007423654462628649, 'mean_word_length': 3.1483554933519944, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5920875', 'n_tokens_mistral': 1885, 'n_tokens_neox': 1771, 'n_words': 414}