nostr 0.45.0-alpha.6

// Copyright (c) 2022-2023 Yuki Kishimoto
// Copyright (c) 2023-2025 Rust Nostr Developers
// Distributed under the MIT software license

use alloc::string::String;
use alloc::vec::Vec;
use core::fmt::Debug;

use base64::engine::{Engine, general_purpose};
#[cfg(feature = "rand")]
use rand::Rng;
#[cfg(all(feature = "std", feature = "os-rng"))]
use rand::rand_core::UnwrapErr;
#[cfg(all(feature = "std", feature = "os-rng"))]
use rand::rngs::SysRng;

use super::v2::{self, ConversationKey};
use crate::error::{Error, ErrorKind};
use crate::key::{PublicKey, SecretKey};
#[cfg(feature = "rand")]
use crate::util;

/// Payload version
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(u8)]
pub enum Version {
    /// V2 - Secp256k1 ECDH, HKDF, padding, ChaCha20, HMAC-SHA256 and base64
    #[default]
    V2 = 0x02,
}

impl Version {
    /// Get [`Version`] as `u8`
    #[inline]
    pub fn as_u8(&self) -> u8 {
        *self as u8
    }
}

impl TryFrom<u8> for Version {
    type Error = Error;

    fn try_from(version: u8) -> Result<Self, Self::Error> {
        match version {
            0x02 => Ok(Self::V2),
            _ => Err(Error::new(
                ErrorKind::Unsupported,
                format!("unknown version: {version}"),
            )),
        }
    }
}

/// NIP-44 nonce
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Nonce {
    /// V2 - 32-byte nonce
    V2([u8; 32]),
}

/// Encrypt
#[inline]
#[cfg(all(feature = "std", feature = "os-rng"))]
pub fn encrypt<T>(
    secret_key: &SecretKey,
    public_key: &PublicKey,
    content: T,
    version: Version,
) -> Result<String, Error>
where
    T: AsRef<[u8]>,
{
    encrypt_with_rng(
        secret_key,
        public_key,
        content,
        version,
        &mut UnwrapErr(SysRng),
    )
}

/// Encrypt
#[cfg(feature = "rand")]
pub fn encrypt_with_rng<R, T>(
    secret_key: &SecretKey,
    public_key: &PublicKey,
    content: T,
    version: Version,
    rng: &mut R,
) -> Result<String, Error>
where
    R: Rng,
    T: AsRef<[u8]>,
{
    let nonce: Nonce = match version {
        Version::V2 => {
            let nonce: [u8; 32] = util::random_32_bytes(rng);
            Nonce::V2(nonce)
        }
    };

    encrypt_with_nonce(secret_key, public_key, content, nonce)
}

/// Encrypt
pub fn encrypt_with_nonce<T>(
    secret_key: &SecretKey,
    public_key: &PublicKey,
    content: T,
    nonce: Nonce,
) -> Result<String, Error>
where
    T: AsRef<[u8]>,
{
    let payload: Vec<u8> = encrypt_to_bytes_with_nonce(secret_key, public_key, content, nonce)?;
    Ok(general_purpose::STANDARD.encode(payload))
}

/// Encrypt to bytes (**not base64 encoded!**)
pub fn encrypt_to_bytes_with_nonce<T>(
    secret_key: &SecretKey,
    public_key: &PublicKey,
    content: T,
    nonce: Nonce,
) -> Result<Vec<u8>, Error>
where
    T: AsRef<[u8]>,
{
    match nonce {
        Nonce::V2(nonce) => {
            let conversation_key: ConversationKey =
                ConversationKey::derive(secret_key, public_key)?;
            let payload: Vec<u8> =
                v2::encrypt_to_bytes_with_nonce(&conversation_key, content.as_ref(), nonce)?;
            Ok(payload)
        }
    }
}

/// Decrypt
#[inline]
pub fn decrypt<T>(
    secret_key: &SecretKey,
    public_key: &PublicKey,
    payload: T,
) -> Result<String, Error>
where
    T: AsRef<[u8]>,
{
    let bytes: Vec<u8> = decrypt_to_bytes(secret_key, public_key, payload)?;
    String::from_utf8(bytes).map_err(Error::malformed)
}

/// Decrypt **without** converting bytes to UTF-8 string
pub fn decrypt_to_bytes<T>(
    secret_key: &SecretKey,
    public_key: &PublicKey,
    payload: T,
) -> Result<Vec<u8>, Error>
where
    T: AsRef<[u8]>,
{
    // Decode base64 payload
    let payload: Vec<u8> = general_purpose::STANDARD
        .decode(payload)
        .map_err(Error::malformed_display)?;

    // Get version byte
    let version: u8 = *payload
        .first()
        .ok_or_else(|| Error::with_static_message(ErrorKind::Missing, "version not found"))?;

    match Version::try_from(version)? {
        Version::V2 => {
            let conversation_key: ConversationKey =
                ConversationKey::derive(secret_key, public_key)?;
            v2::decrypt_to_bytes(&conversation_key, &payload)
        }
    }
}

#[cfg(test)]
#[cfg(all(feature = "std", feature = "os-rng"))]
mod tests {
    use core::str::FromStr;

    use super::*;
    use crate::Keys;

    #[test]
    fn test_nip44_encryption_decryption() {
        // Alice keys
        let alice_sk =
            SecretKey::from_str("5c0c523f52a5b6fad39ed2403092df8cebc36318b39383bca6c00808626fab3a")
                .unwrap();
        let alice_keys = Keys::new(alice_sk);
        let alice_pk = alice_keys.public_key();

        // Bob keys
        let bob_sk =
            SecretKey::from_str("4b22aa260e4acb7021e32f38a6cdf4b673c6a277755bfce287e370c924dc936d")
                .unwrap();
        let bob_keys = Keys::new(bob_sk);
        let bob_pk = bob_keys.public_key();

        let content = String::from("hello");
        let encrypted_content =
            encrypt(alice_keys.secret_key(), &bob_pk, &content, Version::V2).unwrap();
        assert_eq!(
            decrypt(bob_keys.secret_key(), &alice_pk, encrypted_content).unwrap(),
            content
        );
    }
}