/// This module provides the necessary encryption keys to securely encrypt
/// secrets in storage.
/// A [`StorageKey`] that can be used by the client to encrypt secrets and
/// a [`RemoteStorageKey`] that can be used by a server.
use super::{
    generic::{AssociatedData, EncryptionKey},
    Encrypted, MasterKey, SigningKeyPair,
};
use crate::{
    crypto::{data_blob::DataBlob, OpaqueSessionKey},
    types::database::account::UserId,
    LockKeeperError,
};
use rand::{CryptoRng, RngCore};
use std::path::Path;
use utilities::crypto::error::CryptoError;
use zeroize::{Zeroize, ZeroizeOnDrop};

/// The remote storage key is a default-length symmetric encryption key
/// for an AEAD scheme.
///
/// The remote storage key is used by the key server to securely encrypt
/// signing keys generated on (or plaintext imported to) the server.
#[derive(Debug, Clone, Zeroize, ZeroizeOnDrop)]
pub struct RemoteStorageKey(pub(super) EncryptionKey);

impl RemoteStorageKey {
    /// Generate a new 32-byte [`RemoteStorageKey`].
    pub fn generate(rng: &mut (impl CryptoRng + RngCore)) -> Self {
        Self(EncryptionKey::new(rng))
    }

    /// Returns the remote storage key found in the file at the given
    /// path
    pub fn read_from_file(path: impl AsRef<Path>) -> Result<Self, LockKeeperError> {
        let bytes = std::fs::read(&path)
            .map_err(|e| LockKeeperError::FileIo(e, path.as_ref().to_path_buf()))?;
        Self::from_bytes(&bytes)
    }

    /// Returns the remote storage key found in the given bytes
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, LockKeeperError> {
        let key = EncryptionKey::from_bytes(
            bytes
                .try_into()
                .map_err(|_| LockKeeperError::InvalidRemoteStorageKey)?,
            AssociatedData::new().with_str(EncryptionKey::domain_separator()),
        );

        Ok(Self(key))
    }

    /// Encrypt the given [`SigningKeyPair`] under the
    /// [`RemoteStorageKey`] using an AEAD scheme.
    pub fn encrypt_signing_key_pair(
        &self,
        rng: &mut (impl CryptoRng + RngCore),
        signing_key_pair: SigningKeyPair,
    ) -> Result<Encrypted<SigningKeyPair>, LockKeeperError> {
        Encrypted::encrypt(
            rng,
            &self.0,
            signing_key_pair.clone(),
            signing_key_pair.context(),
        )
        .map_err(LockKeeperError::Crypto)
    }

    /// Encrypt the given [`OpaqueSessionKey`] under the
    /// [`RemoteStorageKey`] using an AEAD scheme.
    pub fn encrypt_session_key(
        &self,
        rng: &mut (impl CryptoRng + RngCore),
        session_key: OpaqueSessionKey,
    ) -> Result<Encrypted<OpaqueSessionKey>, LockKeeperError> {
        Encrypted::encrypt(rng, &self.0, session_key.clone(), session_key.context())
            .map_err(LockKeeperError::Crypto)
    }

    /// Encrypt the given [`DataBlob`] under the [`RemoteStorageKey`] using an
    /// AEAD scheme.
    pub fn encrypt_data_blob(
        &self,
        rng: &mut (impl CryptoRng + RngCore),
        data_blob: DataBlob,
    ) -> Result<Encrypted<DataBlob>, LockKeeperError> {
        // Get the cloned context first. This is necessary as the `Encrypted::encrypt`
        // function will take our data_blob by value. The alternative would be
        // cloning the data_blob which would be expensive.
        let context = data_blob.context().clone();

        Encrypted::encrypt(rng, &self.0, data_blob, &context).map_err(LockKeeperError::Crypto)
    }
}

/// A storage key is a default-length symmetric encryption key for an
/// AEAD scheme. The storage key is used to encrypt stored secrets and signing
/// keys.
///
/// It is generated by the client and should never be revealed to the server or
/// the calling application.
/// It should not be stored or saved beyond the lifetime of a single
/// authentication session.
#[derive(Debug, Clone, PartialEq, Eq, ZeroizeOnDrop)]
pub struct StorageKey(pub(super) EncryptionKey);

impl StorageKey {
    pub(super) fn domain_separator() -> &'static str {
        "storage key"
    }

    /// Generate a new 32-byte [`StorageKey`].
    pub fn generate(rng: &mut (impl CryptoRng + RngCore)) -> Self {
        Self(EncryptionKey::new(rng))
    }
}

impl TryFrom<StorageKey> for Vec<u8> {
    type Error = CryptoError;

    fn try_from(key: StorageKey) -> Result<Self, Self::Error> {
        let key = Vec::<u8>::try_from(key.0.to_owned())?;
        let bytes = StorageKey::domain_separator()
            .as_bytes()
            .iter()
            .copied()
            .chain(key)
            .collect();
        Ok(bytes)
    }
}

impl TryFrom<Vec<u8>> for StorageKey {
    type Error = CryptoError;
    fn try_from(bytes: Vec<u8>) -> Result<Self, Self::Error> {
        let expected_domain_sep = StorageKey::domain_separator().as_bytes();
        let domain_separator = bytes
            .get(0..expected_domain_sep.len())
            .ok_or(CryptoError::ConversionError)?;
        if expected_domain_sep != domain_separator {
            return Err(CryptoError::ConversionError);
        }

        // Take off the domain separator
        let key_bytes = bytes
            .get(expected_domain_sep.len()..)
            .ok_or(CryptoError::ConversionError)?
            .to_vec();

        Ok(Self(key_bytes.try_into()?))
    }
}

impl Encrypted<StorageKey> {
    /// Decrypt a storage key. This should be run as part of the subprotocol to
    /// retrieve a storage key from the server.
    ///
    /// This must be run by the client. It takes the following steps:
    /// 1. Derive the decryption key from the master key using the associated
    ///    data
    /// 2. Decrypt the encrypted storage key using the decryption key
    /// 3. Return the decrypted [`StorageKey`]
    pub fn decrypt_storage_key(
        self,
        master_key: MasterKey,
        user_id: &UserId,
    ) -> Result<StorageKey, LockKeeperError> {
        // Check that the associated data is correct.
        let expected_aad = AssociatedData::new()
            .with_bytes(user_id.clone())
            .with_str(StorageKey::domain_separator());
        if self.associated_data != expected_aad {
            return Err(CryptoError::DecryptionFailed.into());
        }

        let decryption_key = master_key.derive_key(self.associated_data.clone())?;
        let decrypted = self.decrypt_inner(&decryption_key)?;
        Ok(decrypted)
    }
}

#[cfg(test)]
pub(super) mod test {
    use super::*;
    use crate::crypto::{test::create_test_session_key, KeyId};
    use rand::{rngs::StdRng, Rng, SeedableRng};
    use std::collections::HashSet;

    // In practice, an export key will be a pseudorandom output from OPAQUE.
    // We'll use random bytes for the test key.
    pub fn create_test_export_key(rng: &mut (impl CryptoRng + RngCore)) -> [u8; 64] {
        let mut key = [0_u8; 64];
        rng.try_fill(&mut key)
            .expect("Failed to generate random key");

        key
    }

    #[test]
    fn storage_key_to_vec_u8_conversion_works() -> Result<(), CryptoError> {
        let mut rng = rand::thread_rng();

        for _ in 0..1000 {
            let storage_key = StorageKey::generate(&mut rng);

            let vec: Vec<u8> = storage_key.clone().try_into()?;
            let output_storage_key = vec.try_into()?;

            assert_eq!(storage_key, output_storage_key);
        }
        Ok(())
    }

    #[test]
    fn storage_key_generation_produces_unique_storage_keys() {
        let mut rng = rand::thread_rng();
        let mut uniq = HashSet::new();

        assert!((0..1000)
            .map(|_| StorageKey::generate(&mut rng))
            .all(|storage_key| uniq.insert(storage_key.0.to_owned())));
    }

    #[test]
    fn storage_keys_are_32_bytes() {
        let mut rng = rand::thread_rng();
        let storage_key = StorageKey::generate(&mut rng);
        assert_eq!(32, storage_key.0.len())
    }

    #[test]
    fn storage_key_encryption_works() -> Result<(), LockKeeperError> {
        let mut rng = rand::thread_rng();
        let export_key = create_test_export_key(&mut rng);
        let user_id = UserId::new(&mut rng)?;

        // Set up matching RNGs to check behavior of the utility function.
        let seed = b"not-random seed for convenience!";
        let mut rng = StdRng::from_seed(*seed);
        let mut rng_copy = StdRng::from_seed(*seed);

        // Encrypt the storage key
        let master_key = MasterKey::derive_master_key(export_key.into())?;
        let storage_key = StorageKey::generate(&mut rng);

        let encrypted_key =
            master_key
                .clone()
                .encrypt_storage_key(&mut rng, storage_key.clone(), &user_id)?;

        // Make sure the utility function gives the same result when encrypting the
        // storage key
        let utility_encrypted_key = master_key
            .clone()
            .create_and_encrypt_storage_key(&mut rng_copy, &user_id)?;
        assert_eq!(utility_encrypted_key, encrypted_key);

        // Decrypt the storage key and check that it worked
        let decrypted_key = encrypted_key.decrypt_storage_key(master_key, &user_id)?;
        assert_eq!(storage_key, decrypted_key);

        Ok(())
    }

    #[test]
    fn storage_key_retrieval_requires_correct_aad() -> Result<(), LockKeeperError> {
        let mut rng = rand::thread_rng();
        let export_key = create_test_export_key(&mut rng);
        let user_id = UserId::new(&mut rng)?;

        let master_key = MasterKey::derive_master_key(export_key.into())?;
        let storage_key = StorageKey::generate(&mut rng);

        let bad_aad = AssociatedData::new()
            .with_bytes(user_id.clone())
            .with_str(StorageKey::domain_separator())
            .with_str("some other content that makes this incorrect");

        let encrypted_storage_key =
            Encrypted::encrypt(&mut rng, &master_key.0, storage_key, &bad_aad)?;
        assert!(encrypted_storage_key
            .decrypt_storage_key(master_key, &user_id)
            .is_err());

        Ok(())
    }

    #[test]
    fn storage_key_requires_correct_encrypted_blob() -> Result<(), LockKeeperError> {
        let mut rng = rand::thread_rng();

        // Set up correct parameters to encrypt a storage key.
        let user_id = UserId::new(&mut rng)?;
        let export_key = create_test_export_key(&mut rng);
        let aad = AssociatedData::new()
            .with_bytes(user_id.clone())
            .with_str(StorageKey::domain_separator());

        // Encrypt any old key (and make sure it's decryptable in general)
        let fake_key = EncryptionKey::new(&mut rng);
        let encrypted_fake_key = Encrypted::encrypt(
            &mut rng,
            &MasterKey::derive_master_key(export_key.into())?.0,
            fake_key,
            &aad,
        )?;
        assert!(encrypted_fake_key
            .clone()
            .decrypt_inner(&MasterKey::derive_master_key(export_key.into())?.0)
            .is_ok());

        // Serialize the fake key, and pretend it's a storage key when you deserialize.
        let fake_storage_key: Encrypted<StorageKey> =
            serde_json::from_str(&serde_json::to_string(&encrypted_fake_key).unwrap()).unwrap();

        // Decryption must fail.
        assert!(fake_storage_key
            .decrypt_storage_key(MasterKey::derive_master_key(export_key.into())?, &user_id)
            .is_err());

        Ok(())
    }

    #[test]
    fn remote_storage_key_from_bytes() -> Result<(), LockKeeperError> {
        let mut rng = rand::thread_rng();
        let encryption_key = RemoteStorageKey::generate(&mut rng);
        let encryption_key_bytes = encryption_key.0.clone().into_bytes();
        let new_encryption_key = RemoteStorageKey::from_bytes(&encryption_key_bytes[..])?;
        assert_eq!(encryption_key.0, new_encryption_key.0);
        Ok(())
    }

    #[test]
    fn remote_storage_key_from_bytes_fails_wrong_length() -> Result<(), LockKeeperError> {
        let mut rng = rand::thread_rng();
        let encryption_key = RemoteStorageKey::generate(&mut rng);
        let encryption_key_bytes = encryption_key.0.clone().into_bytes();
        let new_encryption_key = RemoteStorageKey::from_bytes(&encryption_key_bytes[..31]);
        assert_eq!(
            new_encryption_key.unwrap_err().to_string(),
            LockKeeperError::InvalidRemoteStorageKey.to_string()
        );
        Ok(())
    }

    #[test]
    fn remote_storage_key_works() -> Result<(), LockKeeperError> {
        let mut rng = rand::thread_rng();
        let encryption_key = RemoteStorageKey::generate(&mut rng);
        let user_id = UserId::new(&mut rng)?;
        let key_id = KeyId::generate(&mut rng, &user_id)?;
        let signing_key = SigningKeyPair::remote_generate(&mut rng, &user_id, &key_id);
        let encrypted = encryption_key.encrypt_signing_key_pair(&mut rng, signing_key.clone())?;

        let result = encrypted.decrypt_signing_key_by_server(&encryption_key, user_id, key_id)?;
        assert_eq!(result, signing_key);
        Ok(())
    }

    #[test]
    fn session_key_encryption_works() -> Result<(), LockKeeperError> {
        let mut rng = rand::thread_rng();
        let encryption_key = RemoteStorageKey::generate(&mut rng);
        let session_key = create_test_session_key(&mut rng);
        let encrypted = encryption_key.encrypt_session_key(&mut rng, session_key.clone())?;

        let result = encrypted.decrypt_session_key(&encryption_key)?;
        assert_eq!(result, session_key);
        Ok(())
    }
}