bitwarden-crypto 3.0.0

use std::pin::Pin;

use bitwarden_api_key_connector::models::user_key_response_model::UserKeyResponseModel;
use bitwarden_encoding::B64;
use hybrid_array::Array;
use rand::RngExt;
use tracing::instrument;
use typenum::U32;

use crate::{
    BitwardenLegacyKeyBytes, CryptoError, EncString, KeyDecryptable, KeySlotIds, KeyStoreContext,
    SymmetricCryptoKey, keys::utils::stretch_key,
};

/// Key connector key, used to protect the user key.
#[derive(Clone)]
pub struct KeyConnectorKey(pub(super) Pin<Box<Array<u8, U32>>>);

impl KeyConnectorKey {
    /// Make a new random key for KeyConnector.
    pub fn make() -> Self {
        let mut rng = rand::rng();
        let mut key = Box::pin(Array::<u8, U32>::default());

        rng.fill(key.as_mut_slice());
        KeyConnectorKey(key)
    }

    /// Wraps (encrypts) a user key from the key store using this key connector key.
    ///
    /// The user key identified by `user_key_id` is read from the context and encrypted.
    #[cfg_attr(feature = "dangerous-crypto-debug", instrument(skip(ctx), err))]
    #[cfg_attr(not(feature = "dangerous-crypto-debug"), instrument(skip_all, err))]
    pub fn wrap_user_key<Ids: KeySlotIds>(
        &self,
        user_key_id: Ids::Symmetric,
        ctx: &KeyStoreContext<Ids>,
    ) -> crate::error::Result<EncString> {
        #[allow(deprecated)]
        let user_key = ctx.dangerous_get_symmetric_key(user_key_id)?;
        self.encrypt_user_key(user_key)
    }

    /// Unwraps (decrypts) a user key and stores it in the key store context.
    ///
    /// Returns the local key identifier for the unwrapped user key.
    #[cfg_attr(feature = "dangerous-crypto-debug", instrument(skip(ctx), err))]
    #[cfg_attr(not(feature = "dangerous-crypto-debug"), instrument(skip_all, err))]
    pub fn unwrap_user_key<Ids: KeySlotIds>(
        &self,
        wrapped_user_key: EncString,
        ctx: &mut KeyStoreContext<Ids>,
    ) -> crate::error::Result<Ids::Symmetric> {
        let user_key = self.decrypt_user_key(wrapped_user_key)?;
        Ok(ctx.add_local_symmetric_key(user_key))
    }

    /// Wraps the user key with this key connector key.
    #[cfg_attr(feature = "dangerous-crypto-debug", instrument(err))]
    #[cfg_attr(not(feature = "dangerous-crypto-debug"), instrument(skip_all, err))]
    pub fn encrypt_user_key(
        &self,
        user_key: &SymmetricCryptoKey,
    ) -> crate::error::Result<EncString> {
        let stretched_key = stretch_key(&self.0);
        let user_key_bytes = user_key.to_encoded();
        EncString::encrypt_aes256_hmac(user_key_bytes.as_ref(), &stretched_key)
    }

    /// Unwraps the user key with this key connector key.
    #[cfg_attr(feature = "dangerous-crypto-debug", instrument(err))]
    #[cfg_attr(not(feature = "dangerous-crypto-debug"), instrument(skip_all, err))]
    pub fn decrypt_user_key(
        &self,
        user_key: EncString,
    ) -> crate::error::Result<SymmetricCryptoKey> {
        let dec: Vec<u8> = match user_key {
            // Legacy. user_keys were encrypted using `Aes256Cbc_B64` a long time ago. We've since
            // moved to using `Aes256Cbc_HmacSha256_B64`. However, we still need to support
            // decrypting these old keys.
            EncString::Aes256Cbc_B64 { iv, ref data } => {
                let legacy_key = self.0.clone();
                crate::aes::decrypt_aes256(&iv, data.clone(), &legacy_key)
                    .map_err(|_| CryptoError::Decrypt)?
            }
            EncString::Aes256Cbc_HmacSha256_B64 { .. } => {
                let stretched_key = SymmetricCryptoKey::Aes256CbcHmacKey(stretch_key(&self.0));
                user_key.decrypt_with_key(&stretched_key)?
            }
            _ => {
                return Err(CryptoError::OperationNotSupported(
                    crate::error::UnsupportedOperationError::EncryptionNotImplementedForKey,
                ));
            }
        };

        SymmetricCryptoKey::try_from(&BitwardenLegacyKeyBytes::from(dec))
    }
}

impl std::fmt::Debug for KeyConnectorKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut debug_struct = f.debug_struct("KeyConnectorKey");
        #[cfg(feature = "dangerous-crypto-debug")]
        debug_struct.field("key", &self.0.as_slice());
        debug_struct.finish()
    }
}

impl From<KeyConnectorKey> for B64 {
    fn from(key: KeyConnectorKey) -> Self {
        B64::from(key.0.as_slice())
    }
}

impl TryFrom<UserKeyResponseModel> for KeyConnectorKey {
    type Error = CryptoError;

    fn try_from(s: UserKeyResponseModel) -> Result<Self, Self::Error> {
        let bytes = B64::try_from(s.key).map_err(|_| CryptoError::InvalidKey)?;

        Ok(KeyConnectorKey(Box::pin(
            Array::<u8, U32>::try_from(bytes.as_bytes()).map_err(|_| CryptoError::InvalidKeyLen)?,
        )))
    }
}

#[cfg(test)]
mod tests {
    use bitwarden_encoding::B64;
    use coset::iana::KeyOperation;
    use rand_chacha::rand_core::SeedableRng;

    use super::KeyConnectorKey;
    use crate::{
        BitwardenLegacyKeyBytes, EncString, SymmetricCryptoKey, UserKey,
        store::KeyStore,
        traits::tests::{TestIds, TestSymmKey},
    };

    const KEY_CONNECTOR_KEY_BYTES: [u8; 32] = [
        31, 79, 104, 226, 150, 71, 177, 90, 194, 80, 172, 209, 17, 129, 132, 81, 138, 167, 69, 167,
        254, 149, 2, 27, 39, 197, 64, 42, 22, 195, 86, 75,
    ];

    #[test]
    fn test_make_two_different_keys() {
        let key1 = KeyConnectorKey::make();
        let key2 = KeyConnectorKey::make();
        assert_ne!(key1.0.as_slice(), key2.0.as_slice());
    }

    #[test]
    fn test_into_base64() {
        let key: B64 = KeyConnectorKey(Box::pin(KEY_CONNECTOR_KEY_BYTES.into())).into();

        assert_eq!(
            "H09o4pZHsVrCUKzREYGEUYqnRaf+lQIbJ8VAKhbDVks=",
            key.to_string()
        );
    }

    #[test]
    fn test_decrypt_user_key_aes256_cbc() {
        let key_connector_key = "hvBMMb1t79YssFZkpetYsM3deyVuQv4r88Uj9gvYe08=".to_string();
        let key_connector_key = SymmetricCryptoKey::try_from(key_connector_key).unwrap();
        let SymmetricCryptoKey::Aes256CbcKey(key_connector_key) = &key_connector_key else {
            panic!("Key Connector key is not an Aes256CbcKey");
        };

        let key_connector_key = KeyConnectorKey(key_connector_key.enc_key.clone());

        let user_key: EncString = "0.tn/heK4HLbbEe+yEkC+kvw==|8QM94f7aVTtjm/bmvRdVxOxiLiiZtHYYO7+oBdjFCkilncesx0iVrXPl+tMKqW+Jo7+FtZdPNsTrL6RdoG7i5QbCRVwK+9010+xm7MTQY8s=".parse().unwrap();

        let decrypted_user_key = key_connector_key.decrypt_user_key(user_key).unwrap();
        let SymmetricCryptoKey::Aes256CbcHmacKey(user_key_unwrapped) = &decrypted_user_key else {
            panic!("User key is not an Aes256CbcHmacKey");
        };

        assert_eq!(
            user_key_unwrapped.enc_key.as_slice(),
            [
                116, 170, 187, 43, 80, 212, 193, 202, 234, 181, 57, 66, 151, 249, 59, 47, 70, 16,
                57, 4, 170, 78, 85, 241, 152, 232, 91, 57, 9, 87, 209, 245,
            ]
        );
        assert_eq!(
            user_key_unwrapped.mac_key.as_slice(),
            [
                40, 245, 106, 140, 2, 225, 138, 213, 98, 223, 92, 168, 135, 208, 22, 194, 31, 21,
                178, 252, 203, 198, 35, 174, 53, 218, 254, 151, 235, 57, 7, 98,
            ]
        );
    }

    #[test]
    fn test_encrypt_decrypt_user_key_aes256_cbc_hmac() {
        let rng = rand_chacha::ChaCha8Rng::from_seed([0u8; 32]);

        let key_connector_key = KeyConnectorKey(Box::pin(KEY_CONNECTOR_KEY_BYTES.into()));

        let user_key = SymmetricCryptoKey::make_aes256_cbc_hmac_key_internal(rng);
        let wrapped_user_key = key_connector_key.encrypt_user_key(&user_key).unwrap();
        let user_key = UserKey::new(user_key);

        let decrypted_user_key = key_connector_key
            .decrypt_user_key(wrapped_user_key)
            .unwrap();

        let SymmetricCryptoKey::Aes256CbcHmacKey(user_key_unwrapped) = &decrypted_user_key else {
            panic!("User key is not an Aes256CbcHmacKey");
        };

        assert_eq!(
            user_key_unwrapped.enc_key.as_slice(),
            [
                62, 0, 239, 47, 137, 95, 64, 214, 127, 91, 184, 232, 31, 9, 165, 161, 44, 132, 14,
                195, 206, 154, 127, 59, 24, 27, 225, 136, 239, 113, 26, 30
            ]
        );
        assert_eq!(
            user_key_unwrapped.mac_key.as_slice(),
            [
                152, 76, 225, 114, 185, 33, 111, 65, 159, 68, 83, 103, 69, 109, 86, 25, 49, 74, 66,
                163, 218, 134, 176, 1, 56, 123, 253, 184, 14, 12, 254, 66
            ]
        );

        assert_eq!(
            decrypted_user_key, user_key.0,
            "Decrypted key doesn't match user key"
        );
    }

    #[test]
    fn test_encrypt_decrypt_user_key_xchacha20_poly1305() {
        let key_connector_key = KeyConnectorKey(Box::pin(KEY_CONNECTOR_KEY_BYTES.into()));

        let user_key_b64: B64 = "pQEEAlDib+JxbqMBlcd3KTUesbufAzoAARFvBIQDBAUGIFggt79surJXmqhPhYuuqi9ZyPfieebmtw2OsmN5SDrb4yUB".parse()
            .unwrap();
        let user_key =
            SymmetricCryptoKey::try_from(&BitwardenLegacyKeyBytes::from(&user_key_b64)).unwrap();
        let wrapped_user_key = key_connector_key.encrypt_user_key(&user_key).unwrap();
        let user_key = UserKey::new(user_key);

        let decrypted_user_key = key_connector_key
            .decrypt_user_key(wrapped_user_key)
            .unwrap();

        let SymmetricCryptoKey::XChaCha20Poly1305Key(user_key_unwrapped) = &decrypted_user_key
        else {
            panic!("User key is not an XChaCha20Poly1305Key");
        };

        assert_eq!(
            user_key_unwrapped.enc_key.as_slice(),
            [
                183, 191, 108, 186, 178, 87, 154, 168, 79, 133, 139, 174, 170, 47, 89, 200, 247,
                226, 121, 230, 230, 183, 13, 142, 178, 99, 121, 72, 58, 219, 227, 37
            ]
        );
        assert_eq!(
            user_key_unwrapped.key_id.as_slice(),
            [
                226, 111, 226, 113, 110, 163, 1, 149, 199, 119, 41, 53, 30, 177, 187, 159
            ]
        );
        assert_eq!(
            user_key_unwrapped.supported_operations,
            [
                KeyOperation::Encrypt,
                KeyOperation::Decrypt,
                KeyOperation::WrapKey,
                KeyOperation::UnwrapKey
            ]
        );

        assert_eq!(
            decrypted_user_key, user_key.0,
            "Decrypted key doesn't match user key"
        );
    }

    #[test]
    fn test_wrap_unwrap_user_key_aes256_cbc_hmac() {
        let rng = rand_chacha::ChaCha8Rng::from_seed([0u8; 32]);
        let key_connector_key = KeyConnectorKey(Box::pin(KEY_CONNECTOR_KEY_BYTES.into()));

        let store: KeyStore<TestIds> = KeyStore::default();
        let mut ctx = store.context_mut();

        let user_key = SymmetricCryptoKey::make_aes256_cbc_hmac_key_internal(rng);
        #[allow(deprecated)]
        ctx.set_symmetric_key(TestSymmKey::A(0), user_key.clone())
            .expect("set_symmetric_key should succeed");

        let wrapped = key_connector_key
            .wrap_user_key(TestSymmKey::A(0), &ctx)
            .expect("wrap_user_key should succeed");

        let unwrapped_id = key_connector_key
            .unwrap_user_key(wrapped, &mut ctx)
            .expect("unwrap_user_key should succeed");

        #[allow(deprecated)]
        let unwrapped = ctx
            .dangerous_get_symmetric_key(unwrapped_id)
            .expect("unwrapped key should be in context");

        assert_eq!(&user_key, unwrapped);
    }
}