[PM-23085] Add encrypt_with_key method to cipher_client (#350)

## 🎟️ Tracking

[PM-23085](https://bitwarden.atlassian.net/browse/PM-23085)

## 📔 Objective

Implement `encrypt_with_key` for the `cipher_client` so that the TS
clients can re-encrypt a cipher with a new user key during key rotation.
Once key rotation is implemented fully in the SDK
([PM-23084](https://bitwarden.atlassian.net/browse/PM-23084)) this
method can either be removed or updated to accept a symmetric key
identifier instead of the raw base64 encoded key.

## ⏰ Reminders before review

- Contributor guidelines followed
- All formatters and local linters executed and passed
- Written new unit and / or integration tests where applicable
- Protected functional changes with optionality (feature flags)
- Used internationalization (i18n) for all UI strings
- CI builds passed
- Communicated to DevOps any deployment requirements
- Updated any necessary documentation (Confluence, contributing docs) or
informed the documentation
  team

## 🦮 Reviewer guidelines

<!-- Suggested interactions but feel free to use (or not) as you desire!
-->

- 👍 (`:+1:`) or similar for great changes
- 📝 (`:memo:`) or ℹ️ (`:information_source:`) for notes or general info
- ❓ (`:question:`) for questions
- 🤔 (`:thinking:`) or 💭 (`:thought_balloon:`) for more open inquiry
that's not quite a confirmed
  issue and could potentially benefit from discussion
- 🎨 (`:art:`) for suggestions / improvements
- ❌ (`:x:`) or ⚠️ (`:warning:`) for more significant problems or
concerns needing attention
- 🌱 (`:seedling:`) or ♻️ (`:recycle:`) for future improvements or
indications of technical debt
- ⛏ (`:pick:`) for minor or nitpick changes


[PM-23085]:
https://bitwarden.atlassian.net/browse/PM-23085?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ
[PM-23084]:
https://bitwarden.atlassian.net/browse/PM-23084?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

Author: shane-melton

crates/bitwarden-vault/src/cipher/cipher.rs

@@ -28,7 +28,7 @@ use super::{
     secure_note, ssh_key,
 };
 use crate::{
-    password_history, Fido2CredentialFullView, Fido2CredentialView, Login, LoginView,
+    password_history, EncryptError, Fido2CredentialFullView, Fido2CredentialView, Login, LoginView,
     VaultParseError,
 };
 
@@ -42,6 +42,8 @@ pub enum CipherError {
     VaultLocked(#[from] VaultLockedError),
     #[error(transparent)]
     CryptoError(#[from] CryptoError),
+    #[error(transparent)]
+    EncryptError(#[from] EncryptError),
     #[error("This cipher contains attachments without keys. Those attachments will need to be reuploaded to complete the operation")]
     AttachmentsWithoutKeys,
 }
@@ -526,32 +528,55 @@ impl CipherView {
         Ok(())
     }
 
-    #[allow(missing_docs)]
+    /// Moves the cipher to an organization by re-encrypting the cipher keys with the organization
+    /// key and assigning the organization ID to the cipher.
+    ///
+    /// # Arguments
+    /// * `ctx` - The key store context where the cipher keys will be re-encrypted
+    /// * `organization_id` - The ID of the organization to move the cipher to
     pub fn move_to_organization(
         &mut self,
         ctx: &mut KeyStoreContext<KeyIds>,
         organization_id: Uuid,
     ) -> Result<(), CipherError> {
-        let old_key = self.key_identifier();
         let new_key = SymmetricKeyId::Organization(organization_id);
 
+        self.reencrypt_cipher_keys(ctx, new_key)?;
+        self.organization_id = Some(organization_id);
+
+        Ok(())
+    }
+
+    /// Re-encrypt the cipher key(s) using a new wrapping key.
+    ///
+    /// If the cipher has a cipher key, it will be re-encrypted with the new wrapping key.
+    /// Otherwise, the cipher will re-encrypt all attachment keys and FIDO2 credential keys
+    pub fn reencrypt_cipher_keys(
+        &mut self,
+        ctx: &mut KeyStoreContext<KeyIds>,
+        new_wrapping_key: SymmetricKeyId,
+    ) -> Result<(), CipherError> {
+        let old_key = self.key_identifier();
+
         // If any attachment is missing a key we can't reencrypt the attachment keys
         if self.attachments.iter().flatten().any(|a| a.key.is_none()) {
             return Err(CipherError::AttachmentsWithoutKeys);
         }
 
-        // If the cipher has a key, we need to re-encrypt it with the new organization key
-        if let Some(cipher_key) = &mut self.key {
-            let tmp_cipher_key_id = SymmetricKeyId::Local("cipher_key");
-            ctx.unwrap_symmetric_key(old_key, tmp_cipher_key_id, cipher_key)?;
-            *cipher_key = ctx.wrap_symmetric_key(new_key, tmp_cipher_key_id)?;
+        // If the cipher has a key, reencrypt it with the new wrapping key
+        if self.key.is_some() {
+            // Decrypt the current cipher key using the existing wrapping key
+            let cipher_key = Cipher::decrypt_cipher_key(ctx, old_key, &self.key)?;
+
+            // Wrap the cipher key with the new wrapping key
+            self.key = Some(ctx.wrap_symmetric_key(new_wrapping_key, cipher_key)?);
         } else {
-            // If the cipher does not have a key, we need to reencrypt all attachment keys
-            self.reencrypt_attachment_keys(ctx, old_key, new_key)?;
-            self.reencrypt_fido2_credentials(ctx, old_key, new_key)?;
+            // The cipher does not have a key, we must reencrypt all attachment keys and FIDO2
+            // credentials individually
+            self.reencrypt_attachment_keys(ctx, old_key, new_wrapping_key)?;
+            self.reencrypt_fido2_credentials(ctx, old_key, new_wrapping_key)?;
         }
 
-        self.organization_id = Some(organization_id);
         Ok(())
     }
 
@@ -999,6 +1024,48 @@ mod tests {
         assert!(cipher.attachments.unwrap()[0].key.is_none());
     }
 
+    #[test]
+    fn test_reencrypt_cipher_key() {
+        let old_key = SymmetricCryptoKey::make_aes256_cbc_hmac_key();
+        let new_key = SymmetricCryptoKey::make_aes256_cbc_hmac_key();
+        let key_store = create_test_crypto_with_user_key(old_key);
+        let mut ctx = key_store.context_mut();
+
+        let mut cipher = generate_cipher();
+        cipher
+            .generate_cipher_key(&mut ctx, cipher.key_identifier())
+            .unwrap();
+
+        // Re-encrypt the cipher key with a new wrapping key
+        let new_key_id: SymmetricKeyId = SymmetricKeyId::Local("new_cipher_key");
+        #[allow(deprecated)]
+        ctx.set_symmetric_key(new_key_id, new_key).unwrap();
+
+        cipher.reencrypt_cipher_keys(&mut ctx, new_key_id).unwrap();
+
+        // Check that the cipher key can be unwrapped with the new key
+        assert!(cipher.key.is_some());
+        assert!(ctx
+            .unwrap_symmetric_key(new_key_id, new_key_id, &cipher.key.unwrap())
+            .is_ok());
+    }
+
+    #[test]
+    fn test_reencrypt_cipher_key_ignores_missing_key() {
+        let key = SymmetricCryptoKey::make_aes256_cbc_hmac_key();
+        let key_store = create_test_crypto_with_user_key(key);
+        let mut ctx = key_store.context_mut();
+        let mut cipher = generate_cipher();
+
+        // The cipher does not have a key, so re-encryption should not add one
+        cipher
+            .reencrypt_cipher_keys(&mut ctx, SymmetricKeyId::Local("new_cipher_key"))
+            .unwrap();
+
+        // Check that the cipher key is still None
+        assert!(cipher.key.is_none());
+    }
+
     #[test]
     fn test_move_user_cipher_to_org() {
         let org = uuid::Uuid::new_v4();

crates/bitwarden-vault/src/cipher/cipher_client.rs

@@ -1,5 +1,5 @@
-use bitwarden_core::{Client, OrganizationId};
-use bitwarden_crypto::IdentifyKey;
+use bitwarden_core::{key_management::SymmetricKeyId, Client, OrganizationId};
+use bitwarden_crypto::{CompositeEncryptable, IdentifyKey, SymmetricCryptoKey};
 #[cfg(feature = "wasm")]
 use wasm_bindgen::prelude::*;
 
@@ -46,6 +46,57 @@ impl CiphersClient {
         })
     }
 
+    /// Encrypt a cipher with the provided key. This should only be used when rotating encryption
+    /// keys in the Web client.
+    ///
+    /// Until key rotation is fully implemented in the SDK, this method must be provided the new
+    /// symmetric key in base64 format. See PM-23084
+    ///
+    /// If the cipher has a CipherKey, it will be re-encrypted with the new key.
+    /// If the cipher does not have a CipherKey and CipherKeyEncryption is enabled, one will be
+    /// generated using the new key. Otherwise, the cipher's data will be encrypted with the new
+    /// key directly.
+    #[cfg(feature = "wasm")]
+    pub fn encrypt_cipher_for_rotation(
+        &self,
+        mut cipher_view: CipherView,
+        new_key_b64: String,
+    ) -> Result<EncryptionContext, CipherError> {
+        let new_key = SymmetricCryptoKey::try_from(new_key_b64)?;
+
+        let user_id = self
+            .client
+            .internal
+            .get_user_id()
+            .ok_or(EncryptError::MissingUserId)?;
+        let key_store = self.client.internal.get_key_store();
+        let mut ctx = key_store.context();
+
+        // Set the new key in the key store context
+        const NEW_KEY_ID: SymmetricKeyId = SymmetricKeyId::Local("new_cipher_key");
+        #[allow(deprecated)]
+        ctx.set_symmetric_key(NEW_KEY_ID, new_key)?;
+
+        if cipher_view.key.is_none()
+            && self
+                .client
+                .internal
+                .get_flags()
+                .enable_cipher_key_encryption
+        {
+            cipher_view.generate_cipher_key(&mut ctx, NEW_KEY_ID)?;
+        } else {
+            cipher_view.reencrypt_cipher_keys(&mut ctx, NEW_KEY_ID)?;
+        }
+
+        let cipher = cipher_view.encrypt_composite(&mut ctx, NEW_KEY_ID)?;
+
+        Ok(EncryptionContext {
+            cipher,
+            encrypted_for: user_id,
+        })
+    }
+
     #[allow(missing_docs)]
     pub fn decrypt(&self, cipher: Cipher) -> Result<CipherView, DecryptError> {
         let key_store = self.client.internal.get_key_store();
@@ -127,9 +178,10 @@ impl CiphersClient {
 mod tests {
 
     use bitwarden_core::client::test_accounts::test_bitwarden_com_account;
+    use bitwarden_crypto::CryptoError;
 
     use super::*;
-    use crate::{Attachment, CipherRepromptType, CipherType, Login, VaultClientExt};
+    use crate::{Attachment, CipherRepromptType, CipherType, Login, LoginView, VaultClientExt};
 
     fn test_cipher() -> Cipher {
         Cipher {
@@ -170,6 +222,46 @@ mod tests {
         }
     }
 
+    fn test_cipher_view() -> CipherView {
+        let test_id: uuid::Uuid = "fd411a1a-fec8-4070-985d-0e6560860e69".parse().unwrap();
+        CipherView {
+            r#type: CipherType::Login,
+            login: Some(LoginView {
+                username: Some("test_username".to_string()),
+                password: Some("test_password".to_string()),
+                password_revision_date: None,
+                uris: None,
+                totp: None,
+                autofill_on_page_load: None,
+                fido2_credentials: None,
+            }),
+            id: Some(test_id),
+            organization_id: None,
+            folder_id: None,
+            collection_ids: vec![],
+            key: None,
+            name: "My test login".to_string(),
+            notes: None,
+            identity: None,
+            card: None,
+            secure_note: None,
+            ssh_key: None,
+            favorite: false,
+            reprompt: CipherRepromptType::None,
+            organization_use_totp: true,
+            edit: true,
+            permissions: None,
+            view_password: true,
+            local_data: None,
+            attachments: None,
+            fields: None,
+            password_history: None,
+            creation_date: "2024-01-30T17:55:36.150Z".parse().unwrap(),
+            deleted_date: None,
+            revision_date: "2024-01-30T17:55:36.150Z".parse().unwrap(),
+        }
+    }
+
     fn test_attachment_legacy() -> Attachment {
         Attachment {
             id: Some("uf7bkexzag04d3cw04jsbqqkbpbwhxs0".to_string()),
@@ -417,4 +509,28 @@ mod tests {
 
         assert_eq!(content, b"Hello");
     }
+
+    #[tokio::test]
+    async fn test_encrypt_cipher_for_rotation() {
+        let client = Client::init_test_account(test_bitwarden_com_account()).await;
+
+        let new_key = SymmetricCryptoKey::make_aes256_cbc_hmac_key();
+
+        let cipher_view = test_cipher_view();
+        let new_key_b64 = new_key.to_base64();
+
+        let ctx = client
+            .vault()
+            .ciphers()
+            .encrypt_cipher_for_rotation(cipher_view, new_key_b64)
+            .unwrap();
+
+        assert!(ctx.cipher.key.is_some());
+
+        // Decrypting the cipher "normally" will fail because it was encrypted with a new key
+        assert!(matches!(
+            client.vault().ciphers().decrypt(ctx.cipher).err(),
+            Some(DecryptError::Crypto(CryptoError::InvalidMac))
+        ));
+    }
 }