cbwaw 0.5.50

use blake2::{
    Blake2bVar,
    digest::{Update, VariableOutput},
};
use bytes::Bytes;
use chacha20poly1305::{
    XChaCha20Poly1305, XNonce,
    aead::{Aead, AeadCore, KeyInit, OsRng},
};
use opaque_ke::{
    ClientLogin, ClientLoginFinishParameters, CredentialFinalization, CredentialRequest,
    CredentialResponse, ServerLogin, ServerLoginParameters, ServerRegistration, ServerSetup,
};
use std::error::Error;
use tracing::instrument;

use crate::DefaultCipherSuite;

/// (state, message)
#[cfg(feature = "core")]
pub fn server_start(
    setup: &ServerSetup<DefaultCipherSuite>,
    account: &[u8],
    password_file: &[u8],
    client_start: &[u8],
) -> Result<(Vec<u8>, Bytes), Box<dyn Error>> {
    let password_file = ServerRegistration::<DefaultCipherSuite>::deserialize(password_file)?;
    let mut rng = OsRng;
    let login_start_result = ServerLogin::start(
        &mut rng,
        setup,
        Some(password_file),
        CredentialRequest::deserialize(client_start)?,
        account,
        ServerLoginParameters::default(),
    )?;

    Ok((
        login_start_result.state.serialize().to_vec(),
        Bytes::copy_from_slice(&login_start_result.message.serialize()[..]),
    ))
}

/// encrypted token
#[cfg(feature = "core")]
#[allow(clippy::type_complexity)]
pub fn server_finish(
    auth: &crate::Auth,
    account: &[u8],
    mfa_hash: &[u8],
    encrypted_mfa_hash: &[u8],
    mfa_hash_nonce: &[u8],
    client_finish: &[u8],
    server_start: &[u8],
) -> Result<(Bytes, [u8; 32], Option<Bytes>), Box<dyn Error>> {
    let start_state = ServerLogin::<DefaultCipherSuite>::deserialize(server_start)?;

    let finish_result = start_state.finish(
        CredentialFinalization::deserialize(client_finish)?,
        ServerLoginParameters::default(),
    )?;

    let mut key = [0u8; 32];

    let mut hasher = match Blake2bVar::new(32) {
        Ok(v) => v,
        Err(err) => return Err(err.to_string().into()),
    };
    hasher.update(&finish_result.session_key);
    if let Err(err) = hasher.finalize_variable(&mut key) {
        return Err(err.to_string().into());
    }

    let cipher = XChaCha20Poly1305::new(&key.into());

    let mfa_nonce = XNonce::from_slice(mfa_hash_nonce);
    let mfa_hash_verifier_key_bytes = match cipher.decrypt(mfa_nonce, encrypted_mfa_hash) {
        Ok(v) => v,
        Err(err) => return Err(err.to_string().into()),
    };

    if mfa_hash != &mfa_hash_verifier_key_bytes[0..32] {
        return Err("MFA codes do not match".into());
    }

    let verifier;

    let token = if mfa_hash_verifier_key_bytes.len() == 64 {
        verifier = Some(Bytes::copy_from_slice(&mfa_hash_verifier_key_bytes[32..]));
        auth.generate_refresh_token(account, Some(&mfa_hash_verifier_key_bytes[32..]))?
    } else {
        verifier = None;
        auth.generate_refresh_token(account, None)?
    };

    let nonce = XChaCha20Poly1305::generate_nonce(&mut OsRng);
    let mut ciphertext = match cipher.encrypt(&nonce, token.as_ref()) {
        Ok(v) => v,
        Err(err) => return Err(err.to_string().into()),
    };

    ciphertext.extend_from_slice(&nonce);

    Ok((Bytes::copy_from_slice(&ciphertext), key, verifier))
}

/// (state, message)
#[instrument(skip(password), err)]
#[cfg(feature = "client")]
pub fn client_start(password: &[u8]) -> Result<(Vec<u8>, Vec<u8>), Box<dyn Error>> {
    let mut rng = OsRng;

    match ClientLogin::<DefaultCipherSuite>::start(&mut rng, password) {
        Ok(login) => Ok((
            login.state.serialize().to_vec(),
            login.message.serialize().to_vec(),
        )),
        Err(err) => Err(err.to_string().into()),
    }
}

/// (message, `session_key`)
#[instrument(skip(password, client_start, server_start), err)]
#[cfg(feature = "client")]
pub fn client_finish(
    password: &[u8],
    client_start: &[u8],
    server_start: &[u8],
) -> Result<(Vec<u8>, Vec<u8>), Box<dyn Error>> {
    let client_login = match ClientLogin::<DefaultCipherSuite>::deserialize(client_start) {
        Ok(l) => l,
        Err(err) => return Err(err.to_string().into()),
    };

    let server_start = match CredentialResponse::deserialize(server_start) {
        Ok(s) => s,
        Err(err) => return Err(err.to_string().into()),
    };

    let mut rng = OsRng;

    match client_login.finish(
        &mut rng,
        password,
        server_start,
        ClientLoginFinishParameters::default(),
    ) {
        Ok(finish) => Ok((
            finish.message.serialize().to_vec(),
            finish.session_key.to_vec(),
        )),
        Err(err) => Err(err.to_string().into()),
    }
}

/// token
#[cfg(feature = "client")]
pub fn decrypt_token(encrypted_token: &[u8], session_key: &[u8]) -> Result<Bytes, Box<dyn Error>> {
    let mut key = [0u8; 32];

    let mut hasher = match Blake2bVar::new(32) {
        Ok(v) => v,
        Err(err) => return Err(err.to_string().into()),
    };
    hasher.update(session_key);
    if let Err(err) = hasher.finalize_variable(&mut key) {
        return Err(err.to_string().into());
    }

    if encrypted_token.len() > 24 {
        let nonce_start = encrypted_token.len() - 24;

        let cipher = XChaCha20Poly1305::new(&key.into());
        let nonce = XNonce::from_slice(&encrypted_token[nonce_start..]);

        let plaintext = match cipher.decrypt(nonce, encrypted_token[..nonce_start].as_ref()) {
            Ok(pt) => pt,
            Err(err) => return Err(err.to_string().into()),
        };

        Ok(Bytes::copy_from_slice(&plaintext))
    } else {
        Err("encrypted token is not long enough to have nonce or claims".into())
    }
}