rialo-cdk 0.2.0-alpha.0

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Encryption utilities for securely storing sensitive data.
//!
//! This module provides functionality for password-based encryption and decryption
//! of ED25519 keypairs and arbitrary strings. It uses AES-256-GCM for encryption
//! with PBKDF2 for key derivation.
//!
//! The encrypted data format includes:
//! - A random salt (for PBKDF2)
//! - A random nonce (for AES-GCM)
//! - The encrypted data with authentication tag

#[cfg(feature = "encryption")]
use std::num::NonZeroU32;

#[cfg(feature = "encryption")]
use ed25519_dalek::SigningKey as Keypair;
#[cfg(feature = "encryption")]
use ring::rand::SecureRandom;
#[cfg(feature = "encryption")]
use ring::{aead, pbkdf2};

#[cfg(feature = "encryption")]
use crate::constants::{NONCE_LEN, SALT_LEN, TAG_LEN};
#[cfg(feature = "encryption")]
use crate::error::{Result, RialoError};

/// Encrypts an ED25519 keypair using password-based encryption.
///
/// This function takes a keypair and password, then:
/// 1. Generates a random salt and nonce
/// 2. Derives an encryption key from the password using PBKDF2
/// 3. Encrypts the keypair using AES-256-GCM
/// 4. Returns the encrypted data as a vector containing salt + nonce + ciphertext
///
/// # Arguments
/// * `keypair` - The ED25519 keypair to encrypt
/// * `password` - The password to use for encryption
///
/// # Returns
/// * `Result<Vec<u8>>` - The encrypted keypair or an error
///
/// # Errors
/// Returns an error if:
/// * Salt or nonce generation fails
/// * Key derivation fails
/// * Encryption fails
///
/// # Examples
/// ```
/// # use ed25519_dalek::SigningKey;
/// # use rialo_cdk::wallet::encryption::encrypt_keypair;
/// # fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let keypair = SigningKey::generate(&mut rand::thread_rng());
/// let encrypted = encrypt_keypair(&keypair, "secure_password")?;
/// # Ok(())
/// # }
/// ```
#[cfg(feature = "encryption")]
pub fn encrypt_keypair(keypair: &Keypair, password: &str) -> Result<Vec<u8>> {
    // Generate random salt
    let mut salt = [0u8; SALT_LEN];
    ring::rand::SystemRandom::new()
        .fill(&mut salt)
        .map_err(|_| RialoError::Encryption("Failed to generate salt".to_string()))?;

    // Generate random nonce
    let mut nonce = [0u8; NONCE_LEN];
    ring::rand::SystemRandom::new()
        .fill(&mut nonce)
        .map_err(|_| RialoError::Encryption("Failed to generate nonce".to_string()))?;

    // Derive key from password and salt
    let key = derive_key(password, &salt)?;

    // Get keypair bytes (both private and public parts)
    let keypair_bytes = keypair.as_bytes();

    // Encrypt the keypair
    let nonce_sequence = ring::aead::Nonce::assume_unique_for_key(nonce);
    let mut in_out = keypair_bytes.to_vec();

    key.seal_in_place_append_tag(nonce_sequence, aead::Aad::empty(), &mut in_out)
        .map_err(|_| RialoError::Encryption("Failed to encrypt keypair".to_string()))?;

    // Combine salt, nonce, and encrypted data
    let mut result = Vec::with_capacity(SALT_LEN + NONCE_LEN + in_out.len());
    result.extend_from_slice(&salt);
    result.extend_from_slice(&nonce);
    result.extend_from_slice(&in_out);

    Ok(result)
}

/// Decrypts an ED25519 keypair that was encrypted with `encrypt_keypair`.
///
/// This function takes encrypted data and a password, then:
/// 1. Extracts the salt, nonce, and ciphertext from the input
/// 2. Derives the encryption key from the password and salt
/// 3. Decrypts the keypair using AES-256-GCM
/// 4. Reconstructs and returns the ED25519 keypair
///
/// # Arguments
/// * `encrypted` - The encrypted keypair bytes from `encrypt_keypair`
/// * `password` - The password used for encryption
///
/// # Returns
/// * `Result<Ed25519KeyPair>` - The decrypted ED25519 keypair or an error
///
/// # Errors
/// Returns an error if:
/// * The encrypted data format is invalid
/// * The password is incorrect
/// * Decryption fails
/// * The decrypted data cannot be converted to a valid keypair
///
/// # Examples
/// ```
/// # use rialo_cdk::wallet::encryption::{encrypt_keypair, decrypt_keypair};
/// # use ed25519_dalek::SigningKey;
/// # fn example() -> Result<(), Box<dyn std::error::Error>> {
/// # let original_keypair = SigningKey::generate(&mut rand::thread_rng());
/// # let encrypted = encrypt_keypair(&original_keypair, "secure_password")?;
/// let decrypted_keypair = decrypt_keypair(&encrypted, "secure_password")?;
/// # Ok(())
/// # }
/// ```
#[cfg(feature = "encryption")]
pub fn decrypt_keypair(encrypted: &[u8], password: &str) -> Result<Keypair> {
    // Check if encrypted data has enough bytes
    if encrypted.len() < SALT_LEN + NONCE_LEN + TAG_LEN {
        return Err(RialoError::Encryption("Invalid encrypted data".to_string()));
    }

    // Extract salt, nonce, and encrypted data
    let salt = &encrypted[0..SALT_LEN];
    let nonce = &encrypted[SALT_LEN..SALT_LEN + NONCE_LEN];
    let ciphertext = &encrypted[SALT_LEN + NONCE_LEN..];

    // Derive key from password and salt
    let key = derive_key(password, salt)?;

    // Create a fixed-size array for the nonce
    let mut nonce_array = [0u8; NONCE_LEN];
    nonce_array.copy_from_slice(nonce);

    let nonce_sequence = ring::aead::Nonce::assume_unique_for_key(nonce_array);
    let mut decrypted = ciphertext.to_vec();

    let decrypted_bytes = key
        .open_in_place(nonce_sequence, aead::Aad::empty(), &mut decrypted)
        .map_err(|_| RialoError::Password("Invalid password".to_string()))?;

    // Create keypair from decrypted bytes
    let secret_key_bytes: [u8; 32] = decrypted_bytes[0..32]
        .try_into()
        .map_err(|_| RialoError::Encryption("Failed to extract secret key".to_string()))?;

    // Fix: Keypair::from_bytes doesn't return a Result, so we don't use map_err
    let keypair = Keypair::from_bytes(&secret_key_bytes);

    Ok(keypair)
}

/// Derives an encryption key from a password and salt using PBKDF2.
///
/// This function uses PBKDF2-HMAC-SHA256 with 100,000 iterations to derive
/// a 32-byte key suitable for AES-256-GCM encryption.
///
/// # Arguments
/// * `password` - The password to derive the key from
/// * `salt` - The salt to use for key derivation
///
/// # Returns
/// * `Result<ring::aead::LessSafeKey>` - The derived encryption key or an error
///
/// # Errors
/// Returns an error if:
/// * Key derivation fails
/// * Creating the AES-GCM key fails
#[cfg(feature = "encryption")]
fn derive_key(password: &str, salt: &[u8]) -> Result<ring::aead::LessSafeKey> {
    // Derive key using PBKDF2
    let iterations = NonZeroU32::new(100_000)
        .ok_or_else(|| RialoError::Encryption("Invalid iteration count".to_string()))?;

    let mut key = [0u8; 32]; // AES-256 key

    pbkdf2::derive(
        pbkdf2::PBKDF2_HMAC_SHA256,
        iterations,
        salt,
        password.as_bytes(),
        &mut key,
    );

    // Create AES-GCM key
    let unbound_key = aead::UnboundKey::new(&aead::AES_256_GCM, &key)
        .map_err(|_| RialoError::Encryption("Failed to create encryption key".to_string()))?;

    Ok(aead::LessSafeKey::new(unbound_key))
}

/// Encrypts a string using password-based encryption.
///
/// This function takes a string and password, then:
/// 1. Generates a random salt and nonce
/// 2. Derives an encryption key from the password using PBKDF2
/// 3. Encrypts the string using AES-256-GCM
/// 4. Returns the encrypted data as a vector containing salt + nonce + ciphertext
///
/// # Arguments
/// * `text` - The string to encrypt
/// * `password` - The password to use for encryption
///
/// # Returns
/// * `Result<Vec<u8>>` - The encrypted string or an error
///
/// # Errors
/// Returns an error if:
/// * Salt or nonce generation fails
/// * Key derivation fails
/// * Encryption fails
///
/// # Examples
/// ```
/// # use rialo_cdk::wallet::encryption::encrypt_string;
/// # fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let sensitive_data = "secret information";
/// let encrypted = encrypt_string(sensitive_data, "secure_password")?;
/// # Ok(())
/// # }
/// ```
#[cfg(feature = "encryption")]
pub fn encrypt_string(text: &str, password: &str) -> Result<Vec<u8>> {
    // Generate random salt
    let mut salt = [0u8; SALT_LEN];
    ring::rand::SystemRandom::new()
        .fill(&mut salt)
        .map_err(|_| RialoError::Encryption("Failed to generate salt".to_string()))?;

    // Generate random nonce
    let mut nonce = [0u8; NONCE_LEN];
    ring::rand::SystemRandom::new()
        .fill(&mut nonce)
        .map_err(|_| RialoError::Encryption("Failed to generate nonce".to_string()))?;

    // Derive key from password and salt
    let key = derive_key(password, &salt)?;

    // Get text bytes
    let text_bytes = text.as_bytes();

    // Encrypt the text
    let nonce_sequence = ring::aead::Nonce::assume_unique_for_key(nonce);
    let mut in_out = text_bytes.to_vec();

    key.seal_in_place_append_tag(nonce_sequence, aead::Aad::empty(), &mut in_out)
        .map_err(|_| RialoError::Encryption("Failed to encrypt text".to_string()))?;

    // Combine salt, nonce, and encrypted data
    let mut result = Vec::with_capacity(SALT_LEN + NONCE_LEN + in_out.len());
    result.extend_from_slice(&salt);
    result.extend_from_slice(&nonce);
    result.extend_from_slice(&in_out);

    Ok(result)
}

/// Decrypts a string that was encrypted with `encrypt_string`.
///
/// This function takes encrypted data and a password, then:
/// 1. Extracts the salt, nonce, and ciphertext from the input
/// 2. Derives the encryption key from the password and salt
/// 3. Decrypts the string using AES-256-GCM
/// 4. Converts and returns the decrypted bytes as a UTF-8 string
///
/// # Arguments
/// * `encrypted` - The encrypted string bytes from `encrypt_string`
/// * `password` - The password used for encryption
///
/// # Returns
/// * `Result<String>` - The decrypted string or an error
///
/// # Errors
/// Returns an error if:
/// * The encrypted data format is invalid
/// * The password is incorrect
/// * Decryption fails
/// * The decrypted data is not valid UTF-8
///
/// # Examples
/// ```
/// # use rialo_cdk::wallet::encryption::{encrypt_string, decrypt_string};
/// # fn example() -> Result<(), Box<dyn std::error::Error>> {
/// # let original_text = "secret information";
/// # let encrypted = encrypt_string(original_text, "secure_password")?;
/// let decrypted_text = decrypt_string(&encrypted, "secure_password")?;
/// assert_eq!(decrypted_text, "secret information");
/// # Ok(())
/// # }
/// ```
#[cfg(feature = "encryption")]
pub fn decrypt_string(encrypted: &[u8], password: &str) -> Result<String> {
    // Check if encrypted data has enough bytes
    if encrypted.len() < SALT_LEN + NONCE_LEN + TAG_LEN {
        return Err(RialoError::Encryption("Invalid encrypted data".to_string()));
    }

    // Extract salt, nonce, and encrypted data
    let salt = &encrypted[0..SALT_LEN];
    let nonce = &encrypted[SALT_LEN..SALT_LEN + NONCE_LEN];
    let ciphertext = &encrypted[SALT_LEN + NONCE_LEN..];

    // Derive key from password and salt
    let key = derive_key(password, salt)?;

    // Create a fixed-size array for the nonce
    let mut nonce_array = [0u8; NONCE_LEN];
    nonce_array.copy_from_slice(nonce);

    let nonce_sequence = ring::aead::Nonce::assume_unique_for_key(nonce_array);
    let mut decrypted = ciphertext.to_vec();

    let decrypted_bytes = key
        .open_in_place(nonce_sequence, aead::Aad::empty(), &mut decrypted)
        .map_err(|_| RialoError::Password("Invalid password".to_string()))?;

    // Convert decrypted bytes back to string
    String::from_utf8(decrypted_bytes.to_vec())
        .map_err(|_| RialoError::Encryption("Failed to decode decrypted data as UTF-8".to_string()))
}