From 833d7594aef46ded786ef14dd5434d7ee4685d48 Mon Sep 17 00:00:00 2001
From: Hinton <hinton@users.noreply.github.com>
Date: Fri, 12 Jan 2024 10:49:40 +0100
Subject: [PATCH] Initial work on TDE

---
 crates/bitwarden/src/crypto/device_key.rs | 109 ++++++++++++++++++++++
 crates/bitwarden/src/crypto/mod.rs        |   3 +-
 crates/bitwarden/src/crypto/rsa.rs        |  12 ++-
 3 files changed, 120 insertions(+), 4 deletions(-)
 create mode 100644 crates/bitwarden/src/crypto/device_key.rs

diff --git a/crates/bitwarden/src/crypto/device_key.rs b/crates/bitwarden/src/crypto/device_key.rs
new file mode 100644
index 000000000..46d22c08d
--- /dev/null
+++ b/crates/bitwarden/src/crypto/device_key.rs
@@ -0,0 +1,109 @@
+use rsa::{
+    pkcs8::{DecodePrivateKey, EncodePrivateKey, EncodePublicKey},
+    Oaep, RsaPrivateKey,
+};
+
+use sha1::Sha1;
+
+use super::{
+    rsa::generate_rsa, AsymmEncString, EncString, KeyDecryptable, KeyEncryptable,
+    SymmetricCryptoKey, UserKey,
+};
+use crate::error::Result;
+
+/// Device Key
+///
+/// Encrypts the DevicePrivateKey
+/// Allows the device to decrypt the UserKey, via the DevicePrivateKey.
+#[derive(Debug)]
+pub(crate) struct DeviceKey(SymmetricCryptoKey);
+
+#[derive(Debug)]
+struct CreateDeviceKey {
+    device_key: DeviceKey,
+    /// UserKey encrypted with DevicePublicKey
+    protected_user_key: AsymmEncString,
+    /// DevicePrivateKey encrypted with [DeviceKey]
+    protected_device_private_key: EncString,
+    /// DevicePublicKey encrypted with [UserKey](super::UserKey)
+    protected_device_public_key: EncString,
+}
+
+// We need to support the following scenarios:
+//
+// - Creating device key and encrypting users key
+// - Decrypting users key using device key
+
+impl DeviceKey {
+    /// Generate a new device key
+    fn trust_device(user_key: UserKey) -> Result<CreateDeviceKey> {
+        let device_key = DeviceKey(SymmetricCryptoKey::generate(rand::thread_rng()));
+
+        let device_private_key = generate_rsa();
+
+        let mut rng = rand::thread_rng();
+        let padding = Oaep::new::<Sha1>();
+
+        let protected_user_key = AsymmEncString::Rsa2048_OaepSha1_B64 {
+            data: device_private_key
+                .to_public_key()
+                .encrypt(&mut rng, padding, &user_key.0.key)
+                .map_err(|e| e.to_string())?,
+        };
+
+        let spki = device_private_key
+            .to_public_key()
+            .to_public_key_der()
+            .map_err(|_| "Invalid key")?;
+
+        let protected_device_public_key = spki.as_bytes().encrypt_with_key(&user_key.0)?;
+
+        let pkcs8 = device_private_key
+            .to_pkcs8_der()
+            .map_err(|_| "Invalid key")?;
+
+        let protected_device_private_key = pkcs8.as_bytes().encrypt_with_key(&device_key.0)?;
+
+        Ok(CreateDeviceKey {
+            device_key,
+            protected_user_key,
+            protected_device_private_key,
+            protected_device_public_key,
+        })
+    }
+
+    /// Decrypt the user key using the device key
+    async fn decrypt_user_key(
+        &self,
+        protected_device_private_key: EncString,
+        protected_user_key: AsymmEncString,
+    ) -> Result<UserKey> {
+        let device_private_key: Vec<u8> = protected_device_private_key.decrypt_with_key(&self.0)?;
+        let device_private_key = RsaPrivateKey::from_pkcs8_der(device_private_key.as_slice())
+            .map_err(|e| e.to_string())?;
+
+        let user_key: SymmetricCryptoKey = protected_user_key
+            .decrypt(&device_private_key)?
+            .as_slice()
+            .try_into()?;
+
+        Ok(UserKey(user_key))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::crypto::symmetric_crypto_key::derive_symmetric_key;
+
+    use super::*;
+
+    #[test]
+    fn test_trust_device() {
+        let user_key = UserKey(derive_symmetric_key("test"));
+
+        // Call trust_device function
+        let result = DeviceKey::trust_device(user_key).unwrap();
+
+        println!("result: {:?}", result);
+    }
+}
diff --git a/crates/bitwarden/src/crypto/mod.rs b/crates/bitwarden/src/crypto/mod.rs
index 5355a337a..5f8a161de 100644
--- a/crates/bitwarden/src/crypto/mod.rs
+++ b/crates/bitwarden/src/crypto/mod.rs
@@ -44,6 +44,8 @@ pub use symmetric_crypto_key::SymmetricCryptoKey;
 mod shareable_key;
 pub(crate) use shareable_key::derive_shareable_key;
 
+#[cfg(feature = "internal")]
+mod device_key;
 #[cfg(feature = "internal")]
 mod master_key;
 #[cfg(feature = "internal")]
@@ -58,7 +60,6 @@ pub(crate) use user_key::UserKey;
 mod rsa;
 #[cfg(feature = "internal")]
 pub use self::rsa::RsaKeyPair;
-
 #[cfg(feature = "internal")]
 mod fingerprint;
 #[cfg(feature = "internal")]
diff --git a/crates/bitwarden/src/crypto/rsa.rs b/crates/bitwarden/src/crypto/rsa.rs
index 9641237fc..693d32a10 100644
--- a/crates/bitwarden/src/crypto/rsa.rs
+++ b/crates/bitwarden/src/crypto/rsa.rs
@@ -18,9 +18,7 @@ pub struct RsaKeyPair {
 }
 
 pub(super) fn make_key_pair(key: &SymmetricCryptoKey) -> Result<RsaKeyPair> {
-    let mut rng = rand::thread_rng();
-    let bits = 2048;
-    let priv_key = RsaPrivateKey::new(&mut rng, bits).expect("failed to generate a key");
+    let priv_key = generate_rsa();
     let pub_key = RsaPublicKey::from(&priv_key);
 
     let spki = pub_key
@@ -39,3 +37,11 @@ pub(super) fn make_key_pair(key: &SymmetricCryptoKey) -> Result<RsaKeyPair> {
         private: protected,
     })
 }
+
+// TODO: Move this to AsymmCryptoKey
+/// Generate a new random AsymmetricCryptoKey (RSA-2048)
+pub(crate) fn generate_rsa() -> RsaPrivateKey {
+    let mut rng = rand::thread_rng();
+    let bits = 2048;
+    RsaPrivateKey::new(&mut rng, bits).expect("failed to generate a key")
+}