From eaef2caca81497211f2364353494ee50e17a91f4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Garci=CC=81a?=
 <dani-garcia@users.noreply.github.com>
Date: Wed, 27 Mar 2024 18:58:08 +0100
Subject: [PATCH] Implement make_register_tde_keys

---
 crates/bitwarden-uniffi/src/auth/mod.rs  | 16 ++++++++
 crates/bitwarden/src/auth/client_auth.rs |  9 ++++
 crates/bitwarden/src/auth/mod.rs         |  4 ++
 crates/bitwarden/src/auth/tde.rs         | 52 ++++++++++++++++++++++++
 crates/bitwarden/src/client/client.rs    |  2 +-
 5 files changed, 82 insertions(+), 1 deletion(-)
 create mode 100644 crates/bitwarden/src/auth/tde.rs

diff --git a/crates/bitwarden-uniffi/src/auth/mod.rs b/crates/bitwarden-uniffi/src/auth/mod.rs
index 2a451ffdf..62b709045 100644
--- a/crates/bitwarden-uniffi/src/auth/mod.rs
+++ b/crates/bitwarden-uniffi/src/auth/mod.rs
@@ -2,6 +2,7 @@ use std::sync::Arc;
 
 use bitwarden::auth::{
     password::MasterPasswordPolicyOptions, AuthRequestResponse, RegisterKeyResponse,
+    RegisterTdeKeyResponse,
 };
 use bitwarden_crypto::{AsymmetricEncString, HashPurpose, Kdf, TrustDeviceResponse};
 
@@ -78,6 +79,21 @@ impl ClientAuth {
             .make_register_keys(email, password, kdf)?)
     }
 
+    /// Generate keys needed for TDE process
+    pub async fn make_register_tde_keys(
+        &self,
+        org_public_key: String,
+        remember_device: bool,
+    ) -> Result<RegisterTdeKeyResponse> {
+        Ok(self
+            .0
+             .0
+            .write()
+            .await
+            .auth()
+            .make_register_tde_keys(org_public_key, remember_device)?)
+    }
+
     /// Validate the user password
     ///
     /// To retrieve the user's password hash, use [`ClientAuth::hash_password`] with
diff --git a/crates/bitwarden/src/auth/client_auth.rs b/crates/bitwarden/src/auth/client_auth.rs
index 5f2002133..5029d389b 100644
--- a/crates/bitwarden/src/auth/client_auth.rs
+++ b/crates/bitwarden/src/auth/client_auth.rs
@@ -20,6 +20,7 @@ use crate::{
             MasterPasswordPolicyOptions,
         },
         register::{make_register_keys, register},
+        tde::{make_register_tde_keys, RegisterTdeKeyResponse},
         AuthRequestResponse, RegisterKeyResponse, RegisterRequest,
     },
     client::Kdf,
@@ -73,6 +74,14 @@ impl<'a> ClientAuth<'a> {
         make_register_keys(email, password, kdf)
     }
 
+    pub fn make_register_tde_keys(
+        &mut self,
+        org_public_key: String,
+        remember_device: bool,
+    ) -> Result<RegisterTdeKeyResponse> {
+        make_register_tde_keys(self.client, org_public_key, remember_device)
+    }
+
     pub async fn register(&mut self, input: &RegisterRequest) -> Result<()> {
         register(self.client, input).await
     }
diff --git a/crates/bitwarden/src/auth/mod.rs b/crates/bitwarden/src/auth/mod.rs
index 7918694e2..b1fe2dbda 100644
--- a/crates/bitwarden/src/auth/mod.rs
+++ b/crates/bitwarden/src/auth/mod.rs
@@ -20,6 +20,10 @@ mod auth_request;
 pub use auth_request::AuthRequestResponse;
 #[cfg(feature = "mobile")]
 pub(crate) use auth_request::{auth_request_decrypt_master_key, auth_request_decrypt_user_key};
+#[cfg(feature = "internal")]
+mod tde;
+#[cfg(feature = "internal")]
+pub use tde::RegisterTdeKeyResponse;
 
 #[cfg(feature = "internal")]
 use crate::{client::Kdf, error::Result};
diff --git a/crates/bitwarden/src/auth/tde.rs b/crates/bitwarden/src/auth/tde.rs
new file mode 100644
index 000000000..919ec25e1
--- /dev/null
+++ b/crates/bitwarden/src/auth/tde.rs
@@ -0,0 +1,52 @@
+use base64::{engine::general_purpose::STANDARD, Engine};
+use bitwarden_crypto::{
+    AsymmetricEncString, AsymmetricPublicCryptoKey, DeviceKey, EncString, SymmetricCryptoKey,
+    TrustDeviceResponse, UserKey,
+};
+
+use crate::{error::Result, Client};
+
+pub(super) fn make_register_tde_keys(
+    client: &mut Client,
+    org_public_key: String,
+    remember_device: bool,
+) -> Result<RegisterTdeKeyResponse> {
+    let public_key = AsymmetricPublicCryptoKey::from_der(&STANDARD.decode(org_public_key)?)?;
+
+    let mut rng = rand::thread_rng();
+
+    // Generate a new user key and key pair, and encrypt the user key with the org public key for
+    // admin password reset
+    let user_key = UserKey::new(SymmetricCryptoKey::generate(&mut rng));
+    let key_pair = user_key.make_key_pair()?;
+
+    let admin_reset =
+        AsymmetricEncString::encrypt_rsa2048_oaep_sha1(&user_key.0.to_vec(), &public_key)?;
+
+    let device_key = if remember_device {
+        Some(DeviceKey::trust_device(&user_key.0)?)
+    } else {
+        None
+    };
+
+    // Initialize the crypto with the generated user key, this way it doesn't need to leave the
+    // client
+    client.initialize_user_crypto_decrypted_key(user_key.0, key_pair.private.clone())?;
+
+    Ok(RegisterTdeKeyResponse {
+        private_key: key_pair.private,
+        public_key: key_pair.public,
+
+        admin_reset,
+        device_key,
+    })
+}
+
+#[cfg_attr(feature = "mobile", derive(uniffi::Record))]
+pub struct RegisterTdeKeyResponse {
+    pub private_key: EncString,
+    pub public_key: String,
+
+    pub admin_reset: AsymmetricEncString,
+    pub device_key: Option<TrustDeviceResponse>,
+}
diff --git a/crates/bitwarden/src/client/client.rs b/crates/bitwarden/src/client/client.rs
index 883720529..68442d5d3 100644
--- a/crates/bitwarden/src/client/client.rs
+++ b/crates/bitwarden/src/client/client.rs
@@ -265,7 +265,7 @@ impl Client {
         Ok(self.encryption_settings.as_ref().unwrap())
     }
 
-    #[cfg(feature = "mobile")]
+    #[cfg(feature = "internal")]
     pub(crate) fn initialize_user_crypto_decrypted_key(
         &mut self,
         user_key: SymmetricCryptoKey,