nostr 0.45.1

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

//! NIP49: Private Key Encryption
//!
//! <https://github.com/nostr-protocol/nips/blob/master/49.md>

use alloc::string::String;
use alloc::vec::Vec;

use chacha20poly1305::XChaCha20Poly1305;
use chacha20poly1305::aead::{Aead, KeyInit, Payload};
#[cfg(all(feature = "std", feature = "os-rng"))]
use rand::rand_core::UnwrapErr;
#[cfg(all(feature = "std", feature = "os-rng"))]
use rand::rngs::SysRng;
#[cfg(feature = "rand")]
use rand::{CryptoRng, Rng};
use scrypt::Params as ScryptParams;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use unicode_normalization::UnicodeNormalization;
use zeroize::Zeroizing;

use super::nip19::{FromBech32, ToBech32};
use crate::error::{Error, ErrorKind};
use crate::key::SecretKey;

const SALT_SIZE: usize = 16;
const NONCE_SIZE: usize = 24;
const CIPHERTEXT_SIZE: usize = 48;
const KEY_SIZE: usize = 32;

fn unknown_version(version: u8) -> Error {
    Error::new(
        ErrorKind::Unsupported,
        format!("unknown version: {version}"),
    )
}

fn unknown_key_security(key_security: u8) -> Error {
    Error::new(
        ErrorKind::Unsupported,
        format!("unknown key security: {key_security}"),
    )
}

#[inline]
fn version_not_found() -> Error {
    Error::with_static_message(ErrorKind::Missing, "version not found")
}

#[inline]
fn log2_round_not_found() -> Error {
    Error::with_static_message(ErrorKind::Missing, "log2 round not found")
}

#[inline]
fn salt_not_found() -> Error {
    Error::with_static_message(ErrorKind::Missing, "salt not found")
}

#[inline]
fn nonce_not_found() -> Error {
    Error::with_static_message(ErrorKind::Missing, "nonce not found")
}

#[inline]
fn key_security_not_found() -> Error {
    Error::with_static_message(ErrorKind::Missing, "key security not found")
}

#[inline]
fn cipher_text_not_found() -> Error {
    Error::with_static_message(ErrorKind::Missing, "cipher text not found")
}

/// Encrypted Secret Key version (NIP49)
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Version {
    /// V2
    #[default]
    V2 = 0x02,
}

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

    fn try_from(version: u8) -> Result<Self, Self::Error> {
        match version {
            // 0x01 => deprecated,
            0x02 => Ok(Self::V2),
            v => Err(unknown_version(v)),
        }
    }
}

/// Key security
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum KeySecurity {
    /// The key has been known to have been handled insecurely (stored unencrypted, cut and paste unencrypted, etc)
    Weak = 0x00,
    /// The key has NOT been known to have been handled insecurely (stored encrypted, cut and paste encrypted, etc)
    Medium = 0x01,
    /// The client does not track this data
    #[default]
    Unknown = 0x02,
}

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

    fn try_from(key_security: u8) -> Result<Self, Self::Error> {
        match key_security {
            0x00 => Ok(Self::Weak),
            0x01 => Ok(Self::Medium),
            0x02 => Ok(Self::Unknown),
            v => Err(unknown_key_security(v)),
        }
    }
}

/// Encrypted Secret Key
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct EncryptedSecretKey {
    version: Version,
    log_n: u8,
    salt: [u8; SALT_SIZE],
    nonce: [u8; NONCE_SIZE],
    key_security: KeySecurity,
    ciphertext: [u8; CIPHERTEXT_SIZE],
}

impl EncryptedSecretKey {
    /// Encrypted Secret Key len
    pub const LEN: usize = 1 + 1 + SALT_SIZE + NONCE_SIZE + 1 + CIPHERTEXT_SIZE; // 91;

    /// Maximum scrypt `log_n` accepted by [`Self::decrypt`].
    ///
    /// This limits the memory requested by untrusted encrypted key payloads to
    /// approximately 256 MiB. Use [`Self::decrypt_with_max_log_n`] to explicitly
    /// accept a different limit.
    pub const MAX_LOG_N: u8 = 18;

    /// Encrypt secret key
    #[inline]
    #[cfg(all(feature = "std", feature = "os-rng"))]
    pub fn new(
        secret_key: &SecretKey,
        password: &str,
        log_n: u8,
        key_security: KeySecurity,
    ) -> Result<Self, Error> {
        Self::new_with_rng(
            secret_key,
            password,
            log_n,
            key_security,
            &mut UnwrapErr(SysRng),
        )
    }

    /// Encrypt secret key
    #[cfg(feature = "rand")]
    pub fn new_with_rng<R>(
        secret_key: &SecretKey,
        password: &str,
        log_n: u8,
        key_security: KeySecurity,
        rng: &mut R,
    ) -> Result<Self, Error>
    where
        R: Rng + CryptoRng,
    {
        // Generate salt
        let salt: [u8; SALT_SIZE] = {
            let mut salt: [u8; SALT_SIZE] = [0u8; SALT_SIZE];
            rng.fill_bytes(&mut salt);
            salt
        };

        // Generate nonce
        let mut nonce: [u8; NONCE_SIZE] = [0u8; NONCE_SIZE];
        rng.fill_bytes(&mut nonce);

        Self::new_with_salt_and_nonce(secret_key, password, log_n, key_security, salt, nonce)
    }

    /// Encrypt secret key with custom salt and nonce
    ///
    /// **Use with caution**: improper usage can catastrophically compromise security.
    ///
    /// * **Nonce**: Must be unique for every encryption with the same key. Reusing a nonce
    ///   with the same derived key destroys the security of the stream cipher, potentially
    ///   leaking the secret key.
    /// * **Salt**: Should be random. Using a non-random salt weakens protection against
    ///   pre-computation attacks and causes the same password to always derive the same
    ///   encryption key.
    pub fn new_with_salt_and_nonce(
        secret_key: &SecretKey,
        password: &str,
        log_n: u8,
        key_security: KeySecurity,
        salt: [u8; SALT_SIZE],
        nonce: [u8; NONCE_SIZE],
    ) -> Result<Self, Error> {
        // Derive key
        let key: Zeroizing<[u8; KEY_SIZE]> = derive_key(password, &salt, log_n)?;

        // Compose cipher
        let cipher =
            XChaCha20Poly1305::new_from_slice(key.as_ref()).map_err(Error::crypto_display)?;

        // Compose payload
        let payload = Payload {
            msg: secret_key.as_secret_bytes(),
            aad: &[key_security as u8],
        };

        // Encrypt
        let ciphertext: Vec<u8> = cipher
            .encrypt(&nonce.into(), payload)
            .map_err(Error::crypto_display)?;
        let ciphertext: [u8; CIPHERTEXT_SIZE] =
            ciphertext.as_slice().try_into().map_err(Error::malformed)?;

        Ok(Self {
            version: Version::default(),
            log_n,
            salt,
            nonce,
            key_security,
            ciphertext,
        })
    }

    /// Parse encrypted secret key from bytes
    pub fn from_slice(slice: &[u8]) -> Result<Self, Error> {
        if slice.len() != Self::LEN {
            return Err(Error::with_static_message(
                ErrorKind::Invalid,
                "invalid length",
            ));
        }

        // Version
        let version: u8 = slice.first().copied().ok_or(version_not_found())?;
        let version: Version = Version::try_from(version)?;

        // Log 2 rounds
        let log_n: u8 = slice.get(1).copied().ok_or(log2_round_not_found())?;

        // Salt
        let salt: &[u8] = slice.get(2..2 + SALT_SIZE).ok_or(salt_not_found())?;
        let salt: [u8; SALT_SIZE] = salt.try_into().map_err(Error::malformed)?;

        // Nonce
        let nonce: &[u8] = slice
            .get(2 + SALT_SIZE..2 + SALT_SIZE + NONCE_SIZE)
            .ok_or(nonce_not_found())?;
        let nonce: [u8; NONCE_SIZE] = nonce.try_into().map_err(Error::malformed)?;

        // Key security
        let key_security: u8 = slice
            .get(2 + SALT_SIZE + NONCE_SIZE)
            .copied()
            .ok_or(key_security_not_found())?;
        let key_security: KeySecurity = KeySecurity::try_from(key_security)?;

        // Ciphertext
        let ciphertext: &[u8] = slice
            .get(2 + SALT_SIZE + NONCE_SIZE + 1..)
            .ok_or(cipher_text_not_found())?;
        let ciphertext: [u8; CIPHERTEXT_SIZE] = ciphertext.try_into().map_err(Error::malformed)?;

        Ok(Self {
            version,
            log_n,
            salt,
            nonce,
            key_security,
            ciphertext,
        })
    }

    /// Get encrypted secret key as bytes
    pub fn as_vec(&self) -> Vec<u8> {
        let mut bytes: Vec<u8> = Vec::with_capacity(Self::LEN);
        bytes.push(self.version as u8);
        bytes.push(self.log_n);
        bytes.extend_from_slice(&self.salt);
        bytes.extend_from_slice(&self.nonce);
        bytes.push(self.key_security as u8);
        bytes.extend_from_slice(&self.ciphertext);
        bytes
    }

    /// Get the encrypted secret key version
    #[inline]
    pub fn version(&self) -> Version {
        self.version
    }

    /// Get encryption log_n value
    #[inline]
    pub fn log_n(&self) -> u8 {
        self.log_n
    }

    /// Get encrypted secret key security
    #[inline]
    pub fn key_security(&self) -> KeySecurity {
        self.key_security
    }

    /// Decrypt secret key
    pub fn decrypt(&self, password: &str) -> Result<SecretKey, Error> {
        self.decrypt_with_max_log_n(password, Self::MAX_LOG_N)
    }

    /// Decrypt secret key with a custom maximum scrypt `log_n` value.
    ///
    /// The memory required by scrypt doubles for every increment of `log_n`.
    /// Only increase this limit for trusted payloads and after checking the
    /// resources available to the application.
    pub fn decrypt_with_max_log_n(
        &self,
        password: &str,
        max_log_n: u8,
    ) -> Result<SecretKey, Error> {
        if self.log_n > max_log_n {
            return Err(Error::new(
                ErrorKind::Invalid,
                format!("scrypt log_n {} exceeds maximum {max_log_n}", self.log_n),
            ));
        }

        let bytes = self.decrypt_secret_bytes(password)?;
        SecretKey::from_slice(&bytes)
    }

    fn decrypt_secret_bytes(&self, password: &str) -> Result<Zeroizing<Vec<u8>>, Error> {
        // Derive key
        let key = derive_key(password, &self.salt, self.log_n)?;

        // Compose cipher
        let cipher =
            XChaCha20Poly1305::new_from_slice(key.as_ref()).map_err(Error::crypto_display)?;

        // Compose payload
        let payload = Payload {
            msg: &self.ciphertext,
            aad: &[self.key_security as u8],
        };

        // Decrypt
        let bytes: Vec<u8> = cipher
            .decrypt(&self.nonce.into(), payload)
            .map_err(Error::crypto_display)?;
        Ok(Zeroizing::new(bytes))
    }
}

impl Serialize for EncryptedSecretKey {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let cryptsec: String = self.to_bech32().map_err(serde::ser::Error::custom)?;
        serializer.serialize_str(&cryptsec)
    }
}

impl<'de> Deserialize<'de> for EncryptedSecretKey {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let cryptsec: String = String::deserialize(deserializer)?;
        Self::from_bech32(&cryptsec).map_err(serde::de::Error::custom)
    }
}

#[inline]
fn normalize_password(password: &str) -> Zeroizing<String> {
    Zeroizing::new(password.nfkc().collect())
}

fn derive_key(
    password: &str,
    salt: &[u8; SALT_SIZE],
    log_n: u8,
) -> Result<Zeroizing<[u8; KEY_SIZE]>, Error> {
    // Unicode Normalization
    let password: Zeroizing<String> = normalize_password(password);

    // Compose params
    let params: ScryptParams = ScryptParams::new(log_n, 8, 1).map_err(Error::invalid)?;

    // Derive key
    let mut key: Zeroizing<[u8; KEY_SIZE]> = Zeroizing::new([0u8; KEY_SIZE]);
    scrypt::scrypt(password.as_bytes(), salt, &params, key.as_mut()).map_err(Error::invalid)?;
    Ok(key)
}

#[cfg(test)]
mod tests {
    use alloc::string::ToString;
    use core::str;

    use zeroize::ZeroizeOnDrop;

    use super::*;

    const CRYPTSEC: &str = "ncryptsec1qgg9947rlpvqu76pj5ecreduf9jxhselq2nae2kghhvd5g7dgjtcxfqtd67p9m0w57lspw8gsq6yphnm8623nsl8xn9j4jdzz84zm3frztj3z7s35vpzmqf6ksu8r89qk5z2zxfmu5gv8th8wclt0h4p";
    const SECRET_KEY: &str = "3501454135014541350145413501453fefb02227e449e57cf4d3a3ce05378683";

    #[test]
    fn test_sensitive_intermediates_zeroize_on_drop() {
        fn assert_zeroize_on_drop<T: ZeroizeOnDrop>(_: &T) {}

        let password = normalize_password("ｎｏｓｔｒ");
        assert_eq!(password.as_str(), "nostr");
        assert_zeroize_on_drop(&password);

        let salt = [0u8; SALT_SIZE];
        let key = derive_key(&password, &salt, 1).unwrap();
        assert_zeroize_on_drop(&key);

        let cipher = XChaCha20Poly1305::new_from_slice(key.as_ref()).unwrap();
        assert_zeroize_on_drop(&cipher);
    }

    #[test]
    fn test_decrypted_plaintext_zeroizes_on_drop() {
        fn assert_zeroize_on_drop<T: ZeroizeOnDrop>(_: &T) {}

        let encrypted_secret_key = EncryptedSecretKey::from_bech32(CRYPTSEC).unwrap();
        let bytes = encrypted_secret_key.decrypt_secret_bytes("nostr").unwrap();

        assert_zeroize_on_drop(&bytes);
        let mut out = [0u8; 64];
        faster_hex::hex_encode(&bytes, &mut out).unwrap();
        let hex_out = unsafe { str::from_utf8_unchecked(&out) };
        assert_eq!(hex_out, SECRET_KEY);
    }

    #[test]
    fn test_encrypted_secret_key_decryption() {
        let encrypted_secret_key = EncryptedSecretKey::from_bech32(CRYPTSEC).unwrap();
        let secret_key: SecretKey = encrypted_secret_key.decrypt("nostr").unwrap();
        assert_eq!(secret_key.to_secret_hex(), SECRET_KEY)
    }

    #[test]
    fn test_decryption_rejects_excessive_scrypt_cost() {
        let encrypted_secret_key = EncryptedSecretKey::from_bech32(CRYPTSEC).unwrap();
        let mut bytes = encrypted_secret_key.as_vec();
        bytes[1] = EncryptedSecretKey::MAX_LOG_N + 1;
        let encrypted_secret_key = EncryptedSecretKey::from_slice(&bytes).unwrap();

        let error = encrypted_secret_key.decrypt("nostr").unwrap_err();
        assert_eq!(error.kind(), ErrorKind::Invalid);
        assert_eq!(error.to_string(), "scrypt log_n 19 exceeds maximum 18");
    }

    #[test]
    fn test_encrypted_secret_key_serialization() {
        let encrypted_secret_key = EncryptedSecretKey::from_bech32(CRYPTSEC).unwrap();
        assert_eq!(encrypted_secret_key.to_bech32().unwrap(), CRYPTSEC)
    }

    #[test]
    #[cfg(all(feature = "std", feature = "os-rng"))]
    fn test_encrypted_secret_key_encryption_decryption() {
        let original_secret_key = SecretKey::from_hex(SECRET_KEY).unwrap();
        let encrypted_secret_key =
            EncryptedSecretKey::new(&original_secret_key, "test", 16, KeySecurity::Medium).unwrap();
        let secret_key: SecretKey = encrypted_secret_key.decrypt("test").unwrap();
        assert_eq!(original_secret_key, secret_key);
        assert_eq!(encrypted_secret_key.version(), Version::default());
        assert_eq!(encrypted_secret_key.key_security(), KeySecurity::Medium);
    }
}