note-to-self-lib 0.1.0

use aes_gcm::aead::{Aead, KeyInit};
use aes_gcm::{Aes256Gcm, Nonce};
use anyhow::{anyhow, Context, Result};
use argon2::{Algorithm, Argon2, Params, Version};
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine;
use hkdf::Hkdf;
use rand::{rngs::OsRng, RngCore};
use sha2::{Digest, Sha256};
use zeroize::Zeroize;

const ARGON2_MEMORY_KIB: u32 = 16 * 1024;
const ARGON2_ITERATIONS: u32 = 1;
const ARGON2_PARALLELISM: u32 = 4;
const KEY_LEN: usize = 32;
const NONCE_LEN: usize = 12;

#[derive(Clone, Zeroize)]
#[zeroize(drop)]
pub struct DerivedKeys {
    pub encryption_key: [u8; KEY_LEN],
    pub auth_key: [u8; KEY_LEN],
}

pub fn derive_keys(username: &str, password: &str) -> Result<DerivedKeys> {
    let username = username.trim().to_ascii_lowercase();
    if username.is_empty() {
        return Err(anyhow!("username cannot be empty"));
    }
    if password.is_empty() {
        return Err(anyhow!("password cannot be empty"));
    }

    let salt = Sha256::digest(username.as_bytes());
    let params = Params::new(
        ARGON2_MEMORY_KIB,
        ARGON2_ITERATIONS,
        ARGON2_PARALLELISM,
        Some(KEY_LEN),
    )
    .map_err(|err| anyhow!("invalid Argon2 parameters: {err}"))?;
    let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);
    let mut master = [0u8; KEY_LEN];
    argon2
        .hash_password_into(password.as_bytes(), &salt, &mut master)
        .map_err(|err| anyhow!("key derivation failed: {err}"))?;

    let hkdf = Hkdf::<Sha256>::new(None, &master);
    let mut encryption_key = [0u8; KEY_LEN];
    let mut auth_key = [0u8; KEY_LEN];
    hkdf.expand(b"note-to-self-encryption", &mut encryption_key)
        .map_err(|_| anyhow!("failed to derive encryption key"))?;
    hkdf.expand(b"note-to-self-auth", &mut auth_key)
        .map_err(|_| anyhow!("failed to derive auth key"))?;
    master.zeroize();

    Ok(DerivedKeys {
        encryption_key,
        auth_key,
    })
}

pub fn derive_journal_key(username: &str, journal: &str, password: &str) -> Result<[u8; KEY_LEN]> {
    let username = username.trim().to_ascii_lowercase();
    if username.is_empty() {
        return Err(anyhow!("username cannot be empty"));
    }
    if journal.is_empty() {
        return Err(anyhow!("journal name cannot be empty"));
    }
    if password.is_empty() {
        return Err(anyhow!("journal password cannot be empty"));
    }

    let salt = Sha256::digest(format!("note-to-self-locked-journal:{username}:{journal}"));
    let params = Params::new(
        ARGON2_MEMORY_KIB,
        ARGON2_ITERATIONS,
        ARGON2_PARALLELISM,
        Some(KEY_LEN),
    )
    .map_err(|err| anyhow!("invalid Argon2 parameters: {err}"))?;
    let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);
    let mut master = [0u8; KEY_LEN];
    argon2
        .hash_password_into(password.as_bytes(), &salt, &mut master)
        .map_err(|err| anyhow!("journal key derivation failed: {err}"))?;

    let hkdf = Hkdf::<Sha256>::new(None, &master);
    let mut journal_key = [0u8; KEY_LEN];
    hkdf.expand(b"note-to-self-journal-lock", &mut journal_key)
        .map_err(|_| anyhow!("failed to derive journal key"))?;
    master.zeroize();

    Ok(journal_key)
}

pub fn auth_token(auth_key: &[u8; KEY_LEN]) -> String {
    URL_SAFE_NO_PAD.encode(auth_key)
}

pub fn encode_bytes(bytes: &[u8]) -> String {
    URL_SAFE_NO_PAD.encode(bytes)
}

pub fn decode_bytes(encoded: &str) -> Result<Vec<u8>> {
    URL_SAFE_NO_PAD
        .decode(encoded)
        .context("invalid base64-encoded bytes")
}

pub fn seal(encryption_key: &[u8; KEY_LEN], plaintext: &[u8]) -> Result<Vec<u8>> {
    let cipher = Aes256Gcm::new_from_slice(encryption_key).context("invalid encryption key")?;
    let mut nonce = [0u8; NONCE_LEN];
    OsRng.fill_bytes(&mut nonce);
    let ciphertext = cipher
        .encrypt(Nonce::from_slice(&nonce), plaintext)
        .map_err(|_| anyhow!("encryption failed"))?;

    let mut sealed = Vec::with_capacity(NONCE_LEN + ciphertext.len());
    sealed.extend_from_slice(&nonce);
    sealed.extend_from_slice(&ciphertext);
    Ok(sealed)
}

pub fn open(encryption_key: &[u8; KEY_LEN], sealed: &[u8]) -> Result<Vec<u8>> {
    if sealed.len() < NONCE_LEN {
        return Err(anyhow!("sealed blob is too short"));
    }
    let (nonce, ciphertext) = sealed.split_at(NONCE_LEN);
    let cipher = Aes256Gcm::new_from_slice(encryption_key).context("invalid encryption key")?;
    cipher
        .decrypt(Nonce::from_slice(nonce), ciphertext)
        .map_err(|_| anyhow!("decryption failed"))
}

pub fn checksum(bytes: &[u8]) -> String {
    blake3::hash(bytes).to_hex().to_string()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn derivation_is_deterministic() {
        let a = derive_keys("User@example.com", "correct horse").unwrap();
        let b = derive_keys("user@example.com", "correct horse").unwrap();
        assert_eq!(a.encryption_key, b.encryption_key);
        assert_eq!(a.auth_key, b.auth_key);
    }

    #[test]
    fn journal_key_depends_on_user_journal_and_password() {
        let key = derive_journal_key("alice", "private", "secret").unwrap();
        assert_eq!(
            key,
            derive_journal_key("Alice", "private", "secret").unwrap()
        );
        assert_ne!(key, derive_journal_key("alice", "work", "secret").unwrap());
        assert_ne!(key, derive_journal_key("bob", "private", "secret").unwrap());
        assert_ne!(
            key,
            derive_journal_key("alice", "private", "different").unwrap()
        );
    }

    #[test]
    fn seal_round_trips_and_uses_random_nonce() {
        let keys = derive_keys("alice", "password").unwrap();
        let a = seal(&keys.encryption_key, b"hello").unwrap();
        let b = seal(&keys.encryption_key, b"hello").unwrap();
        assert_ne!(a, b);
        assert_eq!(open(&keys.encryption_key, &a).unwrap(), b"hello");
        assert_eq!(open(&keys.encryption_key, &b).unwrap(), b"hello");
    }
}