# -*- coding: binary -*-

#
# This class implements an override for RubySMB's default authentication method to instead
# use a kerberos authenticator
#
module Msf::Exploit::Remote::SMB::Client::KerberosAuthentication
  # @param [Msf::Exploit::Remote::Kerberos::ServiceAuthenticator::SMB] kerberos_authenticator The authenticator to make the required Kerberos requests
  def kerberos_authenticator=(kerberos_authenticator)
    @kerberos_authenticator = kerberos_authenticator
  end

  def authenticate
    raise ::RubySMB::Error::AuthenticationFailure, "Missing negotiation security buffer" if negotiation_security_buffer.nil?

    begin
      gss_api = OpenSSL::ASN1.decode(negotiation_security_buffer)
      mech_types = RubySMB::Gss.asn1dig(gss_api, 1, 0, 0, 0)&.value || []
      has_kerberos_gss_mech_type = mech_types&.any? { |mech_type| mech_type.value == ::Rex::Proto::Gss::OID_MICROSOFT_KERBEROS_5.value }
    rescue OpenSSL::ASN1::ASN1Error
      raise ::Rex::Proto::Kerberos::Model::Error::KerberosDecodingError.new('Invalid GSS Response')
    end

    error = "Unable to negotiate kerberos with the remote host. Expected oid #{::Rex::Proto::Gss::OID_MICROSOFT_KERBEROS_5.value} in #{mech_types.map(&:value).inspect}"
    raise ::RubySMB::Error::AuthenticationFailure, error unless has_kerberos_gss_mech_type
    @mech_type = :kerberos

    if smb1
      smb1_authenticate
    else
      smb2_authenticate
    end
  end

  #
  # SMB1 Methods
  #

  # Handles SMB1 Kerberos Authentication by delegating to a kerberos_authenticator implementation
  # to generate a GSS security blob with an embedded AP_REQ. On success information is stored
  # about the peer/server.
  def smb1_authenticate
    raise ::RubySMB::Error::AuthenticationFailure, 'Missing kerberos authenticator' unless @kerberos_authenticator

    kerberos_result = @kerberos_authenticator.authenticate

    @application_key = @session_key = kerberos_result[:session_key].value[0...16]

    raw_kerberos_response = smb1_kerberos_authenticate(kerberos_result[:security_blob])
    response = smb1_session_setup_response(raw_kerberos_response)
    @kerberos_authenticator.validate_response!(response.data_block.security_blob)

    response_code = response.status_code

    # Store the available OS information before going forward.
    @peer_native_os = response.data_block.native_os.to_s
    @peer_native_lm = response.data_block.native_lan_man.to_s

    @user_id = response.smb_header.uid if response_code == WindowsError::NTStatus::STATUS_SUCCESS

    response_code
  end

  #
  # @param type3_message [String] the NTLM Type 3 message
  # @param user_id [Integer] the temporary user ID from the Type 2 response
  # @return [String] the raw binary response from the server
  def smb1_kerberos_authenticate(security_buffer)
    packet = smb1_kerberos_authenticate_packet(security_buffer)
    send_recv(packet)
  end

  # Generates the {RubySMB::SMB1::Packet::SessionSetupRequest} packet
  # with the NTLM Type 3 (Auth) message in the security_blob field.
  #
  # @param type3_message [String] the NTLM Type 3 message
  # @param user_id [Integer] the temporary user ID from the Type 2 response
  # @return [RubySMB::SMB1::Packet::SessionSetupRequest] the second authentication packet to send
  def smb1_kerberos_authenticate_packet(security_blob)
    packet = RubySMB::SMB1::Packet::SessionSetupRequest.new
    # packet.smb_header.uid = user_id
    packet.set_security_buffer(security_blob)
    packet.parameter_block.max_buffer_size = self.max_buffer_size
    packet.parameter_block.max_mpx_count = 50
    packet.smb_header.flags2.extended_security = 1
    packet
  end

  #
  # SMB 2 Methods
  #

  # Handles SMB2 Kerberos Authentication by delegating to a kerberos_authenticator implementation
  # to generate a GSS security blob with an embedded AP_REQ. On success information is stored
  # about the peer/server.
  def smb2_authenticate
    raise ::RubySMB::Error::AuthenticationFailure, 'Missing kerberos authenticator' unless @kerberos_authenticator

    kerberos_result = @kerberos_authenticator.authenticate

    raw_kerberos_response = smb2_kerberos_authenticate(kerberos_result[:security_blob])
    response = smb2_session_setup_response(raw_kerberos_response)
    @kerberos_authenticator.validate_response!(response.buffer)

    @session_id = response.smb2_header.session_id
    if @encryption_algorithm.present?
      key_len = OpenSSL::Cipher.new(@encryption_algorithm).key_len
    else
      key_len = 16
    end
    @application_key = @session_key = kerberos_result[:session_key].value[0...key_len]

    @session_is_guest = response.session_flags.guest == 1

    if @smb3
      if response.session_flags.encrypt_data == 1
        # if the server indicates that encryption is required, enable it
        @session_encrypt_data = true
      elsif (@session_is_guest && password != '') || (username == '' && password == '')
        # disable encryption when necessary
        @session_encrypt_data = false
      end

      # see: https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/7fd079ca-17e6-4f02-8449-46b606ea289c
      if @dialect == '0x0300' || @dialect == '0x0302'
        @application_key = Rex::Crypto::KeyDerivation::NIST_SP_800_108.counter_hmac(
          @session_key,
          16,
          'SHA256',
          label: "SMB2APP\x00",
          context: "SmbRpc\x00"
        ).first
      else
        @application_key = Rex::Crypto::KeyDerivation::NIST_SP_800_108.counter_hmac(
          @session_key,
          16,
          'SHA256',
          label: "SMBAppKey\x00",
          context: @preauth_integrity_hash_value
        ).first
      end
      # otherwise, leave encryption to the default value that it was initialized to
    end
    ######
    # DEBUG
    #puts "Session ID = #{@session_id.to_binary_s.each_byte.map {|e| '%02x' % e}.join}"
    #puts "Session key = #{@session_key.each_byte.map {|e| '%02x' % e}.join}"
    #puts "PreAuthHash = #{@preauth_integrity_hash_value.each_byte.map {|e| '%02x' % e}.join}" if @preauth_integrity_hash_value
    ######

    response.status_code
  end

  def smb2_kerberos_authenticate_packet(security_blob)
    packet = RubySMB::SMB2::Packet::SessionSetupRequest.new
    packet.set_security_buffer(security_blob)
    packet.security_mode.signing_enabled = 1
    packet
  end

  def smb2_kerberos_authenticate(security_blob)
    packet = smb2_kerberos_authenticate_packet(security_blob)
    response = send_recv(packet)
    if @dialect == '0x0311'
      update_preauth_hash(packet)
    end
    response
  end
end
