/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::{
    ffi::{c_char, c_void},
    ptr, slice,
};

use base64::prelude::*;

use moz_http::Client;
use nserror::nsresult;
use nsstring::{nsCString, nsString};
use url::Url;
use xpcom::{getter_addrefs, interfaces::nsIAuthModule, RefPtr};

unsafe extern "C" {
    /// Defined and documented in `mailnews_ffi_glue.h`.
    unsafe fn new_auth_module(
        auth_method: *const c_char,
        out_module: *mut *const nsIAuthModule,
    ) -> nsresult;
}

/// The outcome of [`authenticate`].
pub enum NTLMAuthOutcome {
    /// We've successfully managed to authenticate over NTLM.
    Success,

    /// We've failed to authenticate over NTLM.
    Failure,
}

/// Retrieve the next token from the provided [`nsIAuthModule`] in accordance
/// with the in the authentication flow we're currently at.
///
/// If an `input` is provided, it's expected to be a base64-encoded string of
/// the input to feed [`nsIAuthModule::GetNextToken`].
///
/// Calling [`nsIAuthModule::GetNextToken`] on the provided module is expected
/// to return an output buffer that can be cast as `*mut c_char` (i.e. `char*`
/// in C++).
fn next_b64_token_from_auth_module(
    auth_module: &RefPtr<nsIAuthModule>,
    input: Option<String>,
) -> Result<String, nsresult> {
    let input = match input {
        Some(input) => {
            let input = BASE64_STANDARD
                .decode(input)
                .or(Err(nserror::NS_ERROR_FAILURE))?;
            Some(input)
        }
        None => None,
    };

    // We don't shadow `input` here so that the vector stays in scope and isn't
    // deallocated before `GetNextToken` has read the value from its raw
    // pointer.
    let (input_ptr, input_len) = if let Some(ref input) = input {
        let input_len: u32 = input.len().try_into().or(Err(nserror::NS_ERROR_FAILURE))?;
        (input.as_ptr() as *const c_void, input_len)
    } else {
        (ptr::null(), 0u32)
    };

    let mut out_ptr: *mut c_void = ptr::null_mut();
    let mut out_len = 0u32;

    // SAFETY: The memory for `out_ptr` is allocated by the `GetNextToken`
    // implementation, with the size returned in `out_len`. We've also ensured
    // `input_len` matches the length of `input_ptr`. The string that
    // `input_ptr` points to isn't nul-terminated, but this should be fine since
    // we also provide the amount of bytes to read from it. `GetNextToken`'s
    // documentation says the consumer is in charge of ensuring `out_buf` gets
    // eventually freed, but Rust's memory management should ensure this happens
    // automatically.
    unsafe { auth_module.GetNextToken(input_ptr, input_len, &mut out_ptr, &mut out_len) }
        .to_result()?;

    // SAFETY: Considering the existing implementations of `GetNextToken`,
    // although this function's call contract states we expect the module to
    // output a `*c_char`, it might not be nul-terminated. Therefore, we can't
    // just call `CStr::from_ptr`, but instead we need to gather the bytes into
    // a slice with the value written into `out_len` (which `GetNextToken`'s
    // call contract guarantees represents the length of `out_buf`). In here we
    // convert from `c_char` (which might be `i8`) to `u8`, though this should
    // be fine as the binary data should not be affected, only the integer
    // representation of characters (this is also what `CStr::from_ptr` does).
    let out_len: usize = out_len.try_into().or(Err(nserror::NS_ERROR_FAILURE))?;
    let out_buf: &[u8] = unsafe { slice::from_raw_parts((out_ptr as *mut c_char).cast(), out_len) };
    let out_buf = BASE64_STANDARD.encode(out_buf);

    Ok(out_buf)
}

/// Authenticate with the given credentials using NTLM.
///
/// NTLM is performed over HTTP(S) using the following flow:
///  * the client (us) first sends a request with an `Authorization` header that
///    includes a token (generated by an instance of [`nsIAuthModule`]).
///  * the server responds with a `WWW-Authenticate` header that includes a
///    challenge to solve.
///  * the client (us) then sends another request, again with with an
///    `Authorization` header that includes a token generated by
///    [`nsIAuthModule`], this time in order to solve the challenge and
///    communicate the user's credentials.
///  * the server responds with either a 200 OK or 401 Unauthorized status.
///
/// Ideally, we could include the original SOAP request body in the second
/// request, and the final response would also include the server's response if
/// successful. However, at the point of sending this request, we're not
/// entirely sure whether we're authenticating with the right credentials (the
/// user might have supplied e.g. an invalid password). Considering the requests
/// body might be large (e.g. if we're sending a message with a lot of
/// attachments, or operating on a large number of messages or folders), sending
/// a smaller request first slightly increases the number of requests made to
/// perform a single operation, but likely helps reduce the network traffic
/// overall.
///
/// See <https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-ntht/f09cf6e1-529e-403b-a8a5-7368ee096a6a>
pub async fn authenticate(
    username: &str,
    password: &str,
    ews_url: &Url,
) -> Result<NTLMAuthOutcome, nsresult> {
    // SAFETY: `new_auth_module`'s call contract states it always
    // returns a valid `nsIAuthModule` if the function succeeds.
    let ntlm_module = getter_addrefs(|p| unsafe { new_auth_module(c"ntlm".as_ptr(), p) })?;

    let username = nsString::from(username);
    let password = nsString::from(password);

    // SAFETY: All the references/pointers we pass are valid. This
    // call replicates the way the authentication module gets
    // initialized in `nsMsgProtocol::DoNtlmStep1`.
    unsafe {
        ntlm_module.Init(
            &*nsCString::new(),
            0,
            &*nsString::new(),
            &*username,
            &*password,
        )
    }
    .to_result()?;

    // Generate the first message for our NTLM negotiation, and submit
    // it to the server.
    let msg_1 = next_b64_token_from_auth_module(&ntlm_module, None)?;

    log::debug!("ntlm: sending message 1");

    let client = Client::new();
    let resp = client
        .get(ews_url)?
        .header("Authorization", format!("NTLM {msg_1}").as_str())
        .send()
        .await?;

    // Extract and validate the `WWW-Authenticate` header from the
    // response, which should include the NTLM challenge to solve.
    let authenticate_header = resp
        .header("WWW-Authenticate".to_owned())?
        .into_iter()
        // We'll likely have two `WWW-Authenticate` headers in the response at
        // this stage: one with the value `Negotiate` and one containing the
        // challenge (`NTLM XXXXX`). We only care about the latter.
        .find(|value| value.to_lowercase().starts_with("ntlm "));

    let challenge = if let Some(hdr) = authenticate_header {
        hdr.split(" ")
            .nth(1)
            .ok_or_else(|| {
                log::error!("ntlm: missing challenge");
                nserror::NS_ERROR_FAILURE
            })?
            .to_string()
    } else {
        log::error!("ntlm: unexpected empty or missing WWW-Authenticate response header");
        return Err(nserror::NS_ERROR_FAILURE);
    };

    log::debug!("ntlm: sending message 2");

    // Use the challenge to generate the second and final message,
    // which we can attach to the `Authorization` header in the
    // final request to authenticate it.
    let msg_2 = next_b64_token_from_auth_module(&ntlm_module, Some(challenge))?;

    let resp = client
        .get(ews_url)?
        .header("Authorization", format!("NTLM {msg_2}").as_str())
        .send()
        .await?;

    // If we still got an error, then we're not using the right credentials.
    // Otherwise, the server should have set a cookie that will authenticate us
    // in further requests.
    let outcome = match resp.status()?.0 {
        401 => NTLMAuthOutcome::Failure,
        _ => NTLMAuthOutcome::Success,
    };

    Ok(outcome)
}