//! Keys and other sensitive cryptographic material <span style="color:red">**HAZMAT**</span>.
//!
//! This module provides both originating keys ([`RootKey`]) and derived keys ([`DerivedKey`]) as well
//! as nonce types and other hazardous material.
//!
//! Originating keys ([`RootKey`]) are randomly generated, and have methods for their usage in
//! generating derived keys, as well methods for encrypting the key with a password for safe at-rest
//! storage.
//!
//! Derived keys ([`DerivedKey`]) are generated via a KDF over the originating key's entropy pool and a
//! context string incorporating both a random nonce and a namespace.
//!
//! # <span style="color:red">**DANGER**</span>
//!
//! This module deals in low level cryptographic details. It is advisable to not deal with this module
//! directly, and instead use a higher level API.

use argon2::Argon2;
use blake3::Hasher;
use chacha20::{
    cipher::{generic_array::GenericArray, NewCipher, StreamCipher},
    XChaCha20, XNonce,
};
use rand::Rng;
use serde::{Deserialize, Serialize};
use snafu::{ensure, ResultExt};
use zeroize::{Zeroize, Zeroizing};

use crate::{
    crypto::types::{CipherText, ClearText, ConstArray},
    error::{Argon2Failure, BadHMAC, Error},
};

/// Allows access to the subkeys of a key-like structure
pub trait Key {
    /// Provides the encryption key as a chacha [`Key`](chacha20::Key)
    fn encryption_key(&self) -> &chacha20::Key;
    /// Provides the hmac key
    fn hmac_key(&self) -> &[u8; 32];
}

/// Nonce/Salt value used in encryption
///
/// This is stored as a 24-byte array, for serialization, but is viewable as an `XChaCha` nonce
///
/// Nonces are intended to be always randomly generated, and there is intentionally no API for
/// reconstructing a nonce.
#[derive(Debug, Hash, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub struct Nonce(pub(crate) ConstArray<24>);

impl Nonce {
    /// Generates a new, random nonce
    pub fn random() -> Self {
        let mut data = [0_u8; 24];
        rand::thread_rng().fill(&mut data[..]);
        Self(data.into())
    }
    /// Gets a view of the nonce as an `XChaCha20` [`XNonce`]
    pub fn nonce(&self) -> &XNonce {
        GenericArray::from_slice(&self.0)
    }
}

/// Originating Key
///
/// This key consists of an independently, randomly generated encryption key, HMAC key, and an entropy
/// pool for generating derived keys.
///
/// While the components are generated independently from any user provided input, at rest encryption of
/// this key material is achieved by encrypting the `RootKey` using an argon2 derivation of a user
/// supplied password. This decoupling of key generation from user input allows several useful things,
/// such as allowing the password for a store to be changed without having to reencrypt all the data, as
/// well as making entire categories of attacks effectively impossible.
///
/// The included encryption and HMAC keys should only be used for encrypting top-level metadata, to
/// limit the amount of data encrypted with this key. Derived keys can be produced with a namespace
/// string via the `derive` method.
#[derive(Hash, Clone, Serialize, Deserialize, Zeroize)]
#[zeroize(drop)]
pub struct RootKey {
    /// The root cipher key, used for encrypting headers
    encryption: ConstArray<32>,
    /// The root HMAC key, used for validating headers
    hmac: ConstArray<32>,
    /// Random data used to derive encryption keys
    entropy: ConstArray<256>,
}

impl RootKey {
    /// Generates a new `RootKey`
    ///
    /// This method uses a cryptograpically secure random number generator to fill the encryption key, HMAC
    /// key, and entropy pools with random data.
    ///
    /// `RootKey`s should always be randomly generated, and there is intentionally no API for recreating a
    /// specific `RootKey`
    pub fn random() -> Self {
        let mut rand = rand::thread_rng();
        // Construct ourself first, and mutate in place, to limit the possibility of key material
        // getting leaked in a place that `zeroize` can't reach
        let mut ret = Self::null();
        // Fill keys
        rand.fill(&mut ret.encryption[..]);
        rand.fill(&mut ret.hmac[..]);
        rand.fill(&mut ret.entropy[..]);

        ret
    }
    /// Creates an all zero `RootKey`, also known as 'The null key'
    ///
    /// # <span style="color:red">**DANGER**</span>
    ///
    /// This method exists because this library does not support operating on plaintext at rest data, so the
    /// all-zero key is used as a known-ahead-of-time key for passwordless use.
    ///
    /// This is, hopefully obviously, incredibly insecure, and should only ever be called when storing data
    /// in plaintext would be appropriate.
    pub fn null() -> Self {
        RootKey {
            encryption: [0_u8; 32].into(),
            hmac: [0_u8; 32].into(),
            entropy: [0_u8; 256].into(),
        }
    }
    /// Encrypts this key, producing a [`EncryptedRootKey`], with the provided password.
    ///
    /// This uses a byte slice rather than a string to provide more flexibility.
    ///
    /// Argon2 and a random salt are used to generate 64 bytes of key material from the user supplied
    /// password, the first 32 bytes of which are used as the encryption key, and the last 32 bytes are used
    /// as an HMAC key. This allows us to use the password to provide both encryption and authentication.
    ///
    /// # Errors
    ///
    /// Will return:
    ///   * `Error::Argon2Failure` if the argon2 key derivation fails
    pub fn encrypt(&self, password: &[u8]) -> Result<EncryptedRootKey, Error> {
        // Generate a random salt
        let mut salt = [0_u8; 32];
        rand::thread_rng().fill(&mut salt[..]);
        // Prepare the argon2 instance
        let argon = Argon2::default();
        // Prepare the output buffer
        let mut argon_output = Zeroizing::new([0_u8; 64]);
        argon
            .hash_password_into(password, &salt, &mut argon_output[..])
            .context(Argon2Failure)?;
        // Use the first half of the argon output as the encryption key
        let encryption_key: &chacha20::Key = GenericArray::from_slice(&argon_output[0..32]);
        // Use the second half as the hmac key
        let mut hmac_key = Zeroizing::new([0_u8; 32]);
        hmac_key.copy_from_slice(&argon_output[32..]);

        // Get a random nonce
        let nonce = Nonce::random();
        // Serialize ourself
        let mut serial = serde_cbor::to_vec(self).expect("Infallible");
        // Encrypt ourself
        let mut chacha = XChaCha20::new(encryption_key, nonce.nonce());
        chacha.apply_keystream(&mut serial[..]);
        // Calculate the HMAC
        let hmac: [u8; 32] = blake3::keyed_hash(&hmac_key, &serial).into();
        Ok(EncryptedRootKey {
            nonce,
            hmac: hmac.into(),
            salt: salt.into(),
            payload: serial,
        })
    }
    /// Creates a [`DerivedKey`] from this [`RootKey`] using the provided namespace as part of the
    /// context string
    ///
    /// This method will generate, using a CSPRNG a random, 32 character, alphanumeric nonce (~192 bits of
    /// entropy) to include in the context string, which will then be combined with the provided namespace
    /// and fed into Blake3's key derivation mode. Blake3 is used to derive 64 bytes of key material, the
    /// first 32 bytes of which are used as an encryption key, and the last 32 bytes of which are used as an
    /// HMAC key.
    ///
    /// While [`DerivedKey`] can be used quite safely for a large number of encryptions, it is wise to limit
    /// the usage of an individual [`DerivedKey`] as much as possible, to limit the fallout of any
    /// accidental/unintentional nonce reuses.
    pub fn derive(&self, namespace: &str) -> DerivedKey {
        // Get 32 alphanumerics for the nonce
        let mut nonce = Zeroizing::new(String::new());
        let mut rand = rand::thread_rng();
        while nonce.len() < 32 {
            let ch: char = rand.gen();
            if ch.is_ascii_alphanumeric() {
                nonce.push(ch);
            }
        }
        // Setup the context string
        let context_string = format!("snapper-box nonce: {} namespace: {}", &*nonce, namespace);
        self.derive_with_context(context_string)
    }
    /// Creates a [`DerivedKey`] from this [`RootKey`] with a specified context string
    ///
    /// # <span style="color:red">**DANGER**</span>
    ///
    /// The `derive` method intentionally includes a random component to facilitate key rotation. Using this
    /// method for any other purpose then to rederive a lost key is dangerous, as it can lead to unintended
    /// key reuse.
    pub fn derive_with_context(&self, context_string: String) -> DerivedKey {
        // Setup the hasher
        let mut hasher = Hasher::new_derive_key(&context_string);
        // Load in the entropy
        hasher.update(&self.entropy[..]);
        // Make the final derived key
        let mut ret = DerivedKey {
            encryption: [0_u8; 32].into(),
            hmac: [0_u8; 32].into(),
            context_string,
        };
        // Load in the keys, and return
        let mut output = hasher.finalize_xof();
        output.fill(&mut ret.encryption[..]);
        output.fill(&mut ret.hmac[..]);
        ret
    }
}

impl Key for RootKey {
    /// Provides the encryption key as a chacha [`Key`](chacha20::Key)
    fn encryption_key(&self) -> &chacha20::Key {
        GenericArray::from_slice(&self.encryption)
    }
    /// Provides the hmac key as a reference to the underlying array
    fn hmac_key(&self) -> &[u8; 32] {
        self.hmac.as_ref()
    }
}

/// A [`RootKey`] that has been encrypted with an argon2 derivation of a users password
///
/// See the [`RootKey`] docs for a description of the method used for encryption
#[derive(Debug, Hash, Clone, Serialize, Deserialize)]
pub struct EncryptedRootKey {
    /// The nonce used for the encryption
    nonce: Nonce,
    /// The HMAC tag
    hmac: ConstArray<32>,
    /// The salt used for argon2
    salt: ConstArray<32>,
    /// The encrypted payload
    #[serde(with = "serde_bytes")]
    payload: Vec<u8>,
}

impl EncryptedRootKey {
    /// Attempts to decrypt the key with the provided password, and provide it as a [`RootKey`]
    ///
    /// # Errors
    ///
    /// Will return:
    ///   * `Error::Argon2Failure` if the argon2 key derivation fails
    ///   * `Error::BadHmac` if the hmac verification failed, either due to incorrect password or corruption
    ///   * `Error::KeyDeserialization` if the key fails to deserialize, which really shouldn't happen
    pub fn decrypt(&self, password: &[u8]) -> Result<RootKey, Error> {
        // Prepare the argon2 instance
        let argon = Argon2::default();
        // Prepare the output buffer
        let mut argon_output = Zeroizing::new([0_u8; 64]);
        argon
            .hash_password_into(password, &self.salt, &mut argon_output[..])
            .context(Argon2Failure)?;
        // Use the first half of the argon output as the encryption key
        let encryption_key: &chacha20::Key = GenericArray::from_slice(&argon_output[0..32]);
        // Use the second half as the hmac key
        let mut hmac_key = Zeroizing::new([0_u8; 32]);
        hmac_key.copy_from_slice(&argon_output[32..]);
        // Verify the HMAC
        let hmac = blake3::keyed_hash(&hmac_key, &self.payload);
        ensure!(hmac.eq(&*self.hmac), BadHMAC);
        // Decrypt the data
        let mut data = Zeroizing::new(self.payload.clone());
        let mut chacha = XChaCha20::new(encryption_key, self.nonce.nonce());
        chacha.apply_keystream(&mut data[..]);
        // Deserialize the data
        match serde_cbor::from_slice(&data) {
            Ok(x) => Ok(x),
            Err(_) => {
                // Do not preserve the serde error, this may leak secrets into logs
                Err(Error::KeyDeserialization)
            }
        }
    }
}

/// A key that has been derived from a [`RootKey`]
///
/// This will have an encryption key and an hmac key derived from the [`RootKey`]'s entropy pool, using
/// Blake3 in key-derivation mode, keyed with the [`RootKey`]'s hmac key.
///
/// This struct also contains the context string that was used to derive it, as a way to validate key
/// provenance, or possibly reconstruct a key after a corruption occurs.
#[derive(Hash, Clone, Serialize, Deserialize, Zeroize)]
#[zeroize(drop)]
pub struct DerivedKey {
    /// The encryption key
    encryption: ConstArray<32>,
    /// The root HMAC key, used for validating headers
    hmac: ConstArray<32>,
    /// The context string used to produce this key
    context_string: String,
}

impl Key for DerivedKey {
    /// Provides the encryption key as a chacha [`Key`](chacha20::Key)
    fn encryption_key(&self) -> &chacha20::Key {
        GenericArray::from_slice(&self.encryption)
    }
    /// Provides the hmac key
    fn hmac_key(&self) -> &[u8; 32] {
        self.hmac.as_ref()
    }
}

impl DerivedKey {
    /// Encrypts the given [`DerivedKey`] into an [`EncryptedDerivedKey`] using the given [`RootKey`].
    ///
    /// This method packs the derived key into a [`ClearText`] box, and then encrypts it to a [`CipherText`]
    /// box, without enabling compression.
    ///
    /// See the documentation for [`CipherText`] for a description of the encryption method.
    ///
    /// # Errors
    ///
    /// Will error if the encryption or serialization fail
    pub fn encrypt(&self, root_key: &RootKey) -> Result<EncryptedDerivedKey<'static>, Error> {
        // First get a cleartext containing ourself
        let cleartext = ClearText::new(self)?;
        // Then encrypt it with the root key, without compression
        let ciphertext = cleartext.encrypt(root_key, None)?;
        Ok(EncryptedDerivedKey {
            encrypted_key: ciphertext,
        })
    }
}

/// A [`DerivedKey`] that has been encrypted with a [`RootKey`]
///
/// See the documentation for [`CipherText`] for a description of the encryption method.
#[derive(Hash, Clone, Serialize, Deserialize)]
pub struct EncryptedDerivedKey<'a> {
    /// The encrypted key value
    encrypted_key: CipherText<'a>,
}

impl EncryptedDerivedKey<'_> {
    /// Decrypts the given [`EncryptedDerivedKey`] into a [`DerivedKey`], using the provided [`RootKey`].
    ///
    /// # Errors
    ///
    /// Will error if the decryption or deserialization fails
    pub fn decrypt(&self, root_key: &RootKey) -> Result<DerivedKey, Error> {
        // Attempt to get the cleartext
        let cleartext = self.encrypted_key.decrypt(root_key)?;
        // Attempt to deserialize it
        cleartext.deserialize()
    }
}

/// Unit tests
#[cfg(test)]
mod tests {
    use super::*;
    /// Unit tests for [`RootKey`]
    mod root_key {
        use super::*;
        /// Make sure the null key is all zeros
        #[test]
        fn null_is_zeros() {
            let key = RootKey::null();
            assert_eq!(key.encryption, [0_u8; 32].into());
            assert_eq!(key.hmac, [0_u8; 32].into());
            assert_eq!(key.entropy, [0_u8; 256].into());
        }
        /// Make sure randomly generated key has no zero segments
        #[test]
        fn random_is_not_zeros() {
            let key = RootKey::random();
            assert_ne!(key.encryption, [0_u8; 32].into());
            assert_ne!(key.hmac, [0_u8; 32].into());
            assert_ne!(key.entropy, [0_u8; 256].into());
        }
    }
    /// Unit tests for [`Nonce`]
    mod nonce {
        use super::*;
        /// Make sure the nonce is non-zero
        #[test]
        fn non_zero() {
            let nonce = Nonce::random();
            assert_ne!(nonce.0, [0_u8; 24].into());
        }
    }
    /// Unit tests for [`EncryptedRootKey`]
    mod encrypted_root_key {
        use super::*;
        /// Test round trip encryption/decryption of a [`RootKey`]
        #[test]
        fn round_trip() {
            let key = RootKey::random();
            let password = "password".as_bytes();
            let encrypted = key.encrypt(password).expect("Failed to encrypt key");
            let decrypted = encrypted.decrypt(password).expect("Failed to decrypt key");
            assert_eq!(decrypted.encryption, key.encryption);
            assert_eq!(decrypted.hmac, key.hmac);
            assert_eq!(decrypted.entropy, key.entropy);
        }
        /// Make sure the wrong password can't be used to decrypt a key
        #[test]
        fn bad_password_failure() {
            let key = RootKey::random();
            let password = "password".as_bytes();
            let wrong_password = "wrong password".as_bytes();
            let encrypted = key.encrypt(password).expect("Failed to encrypt key");
            let decrypted = encrypted.decrypt(wrong_password);
            assert!(decrypted.is_err());
        }
        /// Make sure corruption is detected
        #[test]
        fn corruption_failure() {
            let key = RootKey::random();
            let password = "password".as_bytes();
            let mut encrypted = key.encrypt(password).expect("Failed to encrypt key");
            // Corrupt the first byte of the payload
            encrypted.payload[0] = encrypted.payload[0].wrapping_add(1_u8);
            let decrypted = encrypted.decrypt(password);
            match decrypted {
                Ok(_) => panic!("Somehow decrypted corrupted data"),
                Err(e) => assert!(matches!(e, Error::BadHMAC)),
            }
        }
    }
    /// Unit tests for [`DerivedKey`]
    mod derived_key {
        use super::*;
        /// Ensure derived keys aren't zeros, and that the encryption key and hmac key aren't the
        /// same
        #[test]
        fn not_zero() {
            let root_key = RootKey::random();
            let derived_key = root_key.derive("namespace");
            assert_ne!(derived_key.encryption, [0_u8; 32].into());
            assert_ne!(derived_key.hmac, [0_u8; 32].into());
            assert_ne!(derived_key.encryption, derived_key.hmac);
        }
        /// Ensure that repeated calls to derive key are different
        #[test]
        fn non_repeatable() {
            let root_key = RootKey::random();
            let derived_key_1 = root_key.derive("namespace");
            let derived_key_2 = root_key.derive("namespace");

            assert_ne!(derived_key_1.encryption, derived_key_2.encryption);
            assert_ne!(derived_key_1.hmac, derived_key_2.hmac);
            assert_ne!(derived_key_1.context_string, derived_key_2.context_string);
        }
        /// Ensure that using the same context string twice gives the same key
        #[test]
        fn repeatable() {
            let root_key = RootKey::random();
            let derived_key_1 = root_key.derive_with_context("Some context goes here".to_string());
            let derived_key_2 = root_key.derive_with_context("Some context goes here".to_string());

            assert_eq!(derived_key_1.encryption, derived_key_2.encryption);
            assert_eq!(derived_key_1.hmac, derived_key_2.hmac);
            assert_eq!(derived_key_1.context_string, derived_key_2.context_string);
        }
    }
    /// Unit tests for [`EncryptedDerivedKey`]
    mod enc_derived_key {
        use super::*;
        /// Test round trip encryption/decryption of a [`DerivedKey`]
        #[test]
        fn round_trip() {
            let root_key = RootKey::random();
            let derived_key_orig = root_key.derive("testing");
            let enc_derived_key = derived_key_orig
                .encrypt(&root_key)
                .expect("Failed to encrypt key");
            let derived_key_deser = enc_derived_key
                .decrypt(&root_key)
                .expect("Failed to decrypt key");
            assert_eq!(derived_key_deser.encryption, derived_key_orig.encryption);
            assert_eq!(derived_key_deser.hmac, derived_key_orig.hmac);
            assert_eq!(
                derived_key_deser.context_string,
                derived_key_orig.context_string
            );
        }
    }
}