//! An implementation of symmetric keys for ChaCha20Poly1305 encryption.
//!
//! The ciphertext is authenticated.
//! The 12-byte nonce is appended to the generated ciphertext.
//! Keys are 32 bytes in size.

use chacha20poly1305::{ChaCha20Poly1305, KeyInit, Nonce, aead::Aead};

use crate::error::CryptoError;

#[derive(Debug, PartialEq, Eq)]
pub struct SymmetricKey([u8; 32]);

/// Size of a nonce in ChaCha20Poly1305.
const NONCE_SIZE: usize = 12;

impl SymmetricKey {
    /// Generate a new random SymmetricKey.
    pub fn new() -> Self {
        let mut key = [0u8; 32];
        rand::fill(&mut key);
        Self(key)
    }

    /// Create a new key from existing bytes.
    pub(crate) fn from_bytes(bytes: &[u8]) -> Result<SymmetricKey, CryptoError> {
        if bytes.len() == 32 {
            let mut key = [0u8; 32];
            key.copy_from_slice(bytes);
            Ok(SymmetricKey(key))
        } else {
            Err(CryptoError::InvalidKeySize)
        }
    }

    /// View the raw bytes of this key
    pub(crate) fn as_raw_bytes(&self) -> &[u8; 32] {
        &self.0
    }

    /// Encrypt a plaintext with the key. A nonce is generated and appended to the end of the
    /// message.
    pub fn encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>, CryptoError> {
        // Create cipher
        let cipher = ChaCha20Poly1305::new_from_slice(&self.0)
            .expect("Key is a fixed 32 byte array so size is always ok");

        // Generate random nonce
        let mut nonce_bytes = [0u8; NONCE_SIZE];
        rand::fill(&mut nonce_bytes);
        let nonce = Nonce::from_slice(&nonce_bytes);

        // Encrypt message
        let mut ciphertext = cipher
            .encrypt(nonce, plaintext)
            .map_err(|_| CryptoError::EncryptionFailed)?;

        // Append nonce to ciphertext
        ciphertext.extend_from_slice(&nonce_bytes);

        Ok(ciphertext)
    }

    /// Decrypts a ciphertext with appended nonce.
    pub fn decrypt(&self, ciphertext: &[u8]) -> Result<Vec<u8>, CryptoError> {
        // Check if ciphertext is long enough to contain a nonce
        if ciphertext.len() <= NONCE_SIZE {
            return Err(CryptoError::DecryptionFailed);
        }

        // Extract nonce from the end of ciphertext
        let ciphertext_len = ciphertext.len() - NONCE_SIZE;
        let nonce_bytes = &ciphertext[ciphertext_len..];
        let ciphertext = &ciphertext[0..ciphertext_len];

        // Create cipher
        let cipher = ChaCha20Poly1305::new_from_slice(&self.0)
            .expect("Key is a fixed 32 byte array so size is always ok");

        let nonce = Nonce::from_slice(nonce_bytes);

        // Decrypt message
        cipher
            .decrypt(nonce, ciphertext)
            .map_err(|_| CryptoError::DecryptionFailed)
    }

    /// Derives a new symmetric key from a password.
    ///
    /// Derivation is done using pbkdf2 with Sha256 hashing.
    pub fn derive_from_password(password: &str) -> Self {
        /// Salt to use for PBKDF2. This needs to be consistent accross runs to generate the same
        /// key. Additionally, it does not really matter what this is, as long as its unique.
        const SALT: &[u8; 10] = b"vault_salt";
        /// Amount of rounds to use for key generation. More rounds => more cpu time. Changing this
        /// also chagnes the generated keys.
        const ROUNDS: u32 = 100_000;

        let mut key = [0; 32];

        pbkdf2::pbkdf2_hmac::<sha2::Sha256>(password.as_bytes(), SALT, ROUNDS, &mut key);

        Self(key)
    }
}

#[cfg(test)]
mod tests {

    /// Using the same password derives the same key
    #[test]
    fn same_password_derives_same_key() {
        const EXPECTED_KEY: [u8; 32] = [
            4, 179, 233, 202, 225, 70, 211, 200, 7, 73, 115, 1, 85, 149, 90, 42, 160, 68, 16, 106,
            136, 19, 197, 195, 153, 145, 179, 21, 37, 13, 37, 90,
        ];
        const PASSWORD: &str = "test123";

        let key = super::SymmetricKey::derive_from_password(PASSWORD);

        assert_eq!(key.0, EXPECTED_KEY);
    }

    /// Make sure an encrypted value with some key can be decrypted with the same key
    #[test]
    fn can_decrypt() {
        let key = super::SymmetricKey::new();

        let message = b"this is a message to decrypt";

        let enc = key.encrypt(message).expect("Can encrypt message");
        let dec = key.decrypt(&enc).expect("Can decrypt message");

        assert_eq!(message.as_slice(), dec.as_slice());
    }

    /// Make sure a value encrypted with one key can't be decrypted with a different key. Since we
    /// use AEAD encryption we will notice this when trying to decrypt
    #[test]
    fn different_key_cant_decrypt() {
        let key1 = super::SymmetricKey::new();
        let key2 = super::SymmetricKey::new();

        let message = b"this is a message to decrypt";

        let enc = key1.encrypt(message).expect("Can encrypt message");
        let dec = key2.decrypt(&enc);

        assert!(dec.is_err());
    }
}