ncrypt-me 0.3.0

use super::*;

use chacha20poly1305::{
   AeadCore, KeyInit, XChaCha20Poly1305,
   aead::{Aead, OsRng, Payload, generic_array::GenericArray, rand_core::RngCore},
};
use secure_types::SecureBytes;

/*
██████████████████████████████████████████████████████████████████████████████
█                                                                            █
█                           nCrypt File Format                               █
█                                                                            █
█    ┌───────────┬──────────────────┬──────────────────┬───────────────┐     █
█    │   Header  │ EncryptedInfo Len│  EncryptedInfo   │ Encrypted Data│     █
█    │  8 bytes  │     4 bytes      │  Dyn Size        │ Dyn Size      │     █
█    └───────────┴──────────────────┴──────────────────┴───────────────┘     █
█                                                                            █
█                                                                            █
██████████████████████████████████████████████████████████████████████████████
*/

// File Headers
pub const HEADER: &[u8; 8] = b"nCrypt1\0";

pub const HEADER_02: &[u8; 8] = b"nCrypt2\0";

/// Encrypts the given data
///
/// ### Arguments
///
/// - `argon2` - The Argon2 instance to use for the password hashing
/// - `data` - The data to encrypt
/// - `credentials` - The credentials to use for encryption
pub fn encrypt_data(
   argon2: Argon2,
   data: SecureBytes,
   credentials: Credentials,
) -> Result<Vec<u8>, Error> {
   let (encrypted_data, info) = encrypt(argon2, credentials, data)?;

   let encoded_info = info.encode();

   // Construct the file format
   let mut result = Vec::new();

   // Append the header
   result.extend_from_slice(HEADER_02);

   // Append the EncryptedInfo Length
   let info_length = encoded_info.len() as u32;
   result.extend_from_slice(&info_length.to_le_bytes());

   // Append the EncryptedInfo
   result.extend_from_slice(&encoded_info);

   // Append the encrypted Data
   result.extend_from_slice(&encrypted_data);

   Ok(result)
}

fn encrypt(
   argon2: Argon2,
   credentials: Credentials,
   data: SecureBytes,
) -> Result<(Vec<u8>, EncryptedInfo), Error> {
   credentials.is_valid()?;

   if argon2.hash_length < 32 {
      return Err(Error::HashLength);
   }

   let mut password_salt = vec![0u8; RECOMMENDED_SALT_LEN];
   let mut username_salt = vec![0u8; RECOMMENDED_SALT_LEN];

   OsRng
      .try_fill_bytes(&mut password_salt)
      .map_err(|e| Error::Custom(e.to_string()))?;
   OsRng
      .try_fill_bytes(&mut username_salt)
      .map_err(|e| Error::Custom(e.to_string()))?;

   let cipher_nonce = XChaCha20Poly1305::generate_nonce(&mut OsRng);

   let mut aad = credentials
      .username
      .unlock_str(|username_str| argon2.hash_password(&username_str, username_salt.clone()))
      .map_err(|e| Error::Custom(e.to_string()))?;

   let password_hash = credentials
      .password
      .unlock_str(|password_str| argon2.hash_password(&password_str, password_salt.clone()))
      .map_err(|e| Error::Custom(e.to_string()))?;

   data.unlock_slice(|data| {
      let payload = Payload {
         msg: data,
         aad: &aad,
      };

      let cipher = xchacha20_poly_1305(password_hash);

      let encrypted_data_res = cipher.encrypt(&cipher_nonce, payload);
      aad.zeroize();

      let encrypted_data = match encrypted_data_res {
         Ok(data) => data,
         Err(e) => {
            return Err(Error::EncryptionFailed(e.to_string()));
         }
      };

      let info = EncryptedInfo::new(
         password_salt,
         username_salt,
         cipher_nonce.to_vec(),
         argon2,
      );

      Ok((encrypted_data, info))
   })
}

pub(crate) fn xchacha20_poly_1305(mut hash_output: Vec<u8>) -> XChaCha20Poly1305 {
   let mut key = GenericArray::clone_from_slice(&hash_output[..32]);
   hash_output.zeroize();

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