Feat/cs/cleanup

src/backup/crypto.rs

@@ -20,7 +20,7 @@ use sha2::Sha256;
 use crate::{
     backup::models::{
         file::WrappedKey,
-        keyring::{EncryptionKey, ProtectionClassKey},
+        keyring::{ClassKeyData, EncryptionKey, ProtectionClassKey},
         manifest::manifest_plist::ManifestData,
     },
     error::{BackupError, Result},
@@ -108,40 +108,43 @@ pub(crate) fn unlock_keys_from_manifest(
         .as_ref()
         .ok_or_else(|| BackupError::Crypto("BackupKeyBag not found in PlistInfo".to_string()))?;
 
-    for (class_id, class_key_data) in &key_ring.class_keys {
-        // Skip classes without WPKY
-        let Some(wpky) = &class_key_data.wpky else {
-            continue;
-        };
-
-        // Check wrap flags for passcode protection (bit 0x02)
-        let Some(wrap_bytes) = &class_key_data.wrap else {
-            continue;
-        };
-
-        // Parse wrap flag as big-endian u32
-        let wrap_val = u32::from_be_bytes(
-            wrap_bytes
-                .as_slice()
-                .try_into()
-                .map_err(|_| BackupError::KeyUnwrapFailed(*class_id))?,
-        );
-
-        if wrap_val & 0x02 == 0 {
-            continue; // Skip keys not protected by passcode
-        }
+    // Iterate over each class key in the key ring
+    for (&class_id, class_key_data) in &key_ring.class_keys {
+        match class_key_data {
+            //  Unwrapped key data with wrapped key
+            ClassKeyData {
+                wpky: Some(wpky),
+                wrap: Some(wrap_bytes),
+                ..
+            } => {
+                // Ensure wrap_bytes is exactly 4 bytes (u32)
+                let wrap_val = u32::from_be_bytes(
+                    wrap_bytes
+                        .as_slice()
+                        .try_into()
+                        .map_err(|_| BackupError::KeyUnwrapFailed(class_id))?,
+                );
+                // Check if the key is wrapped with AES Key Wrap
+                if wrap_val & 0x02 == 0 {
+                    // Not wrapped with AES Key Wrap, skip
+                    continue;
+                }
 
-        // Unwrap class key using AES key wrap (RFC 3394)
-        let unwrapped = aes_kw_unwrap(master_key, &WrappedKey::from(wpky.clone()))
-            .map_err(|_| BackupError::KeyUnwrapFailed(*class_id))?;
-
-        unlocked_keys.insert(
-            *class_id,
-            ProtectionClassKey {
-                class_id: *class_id,
-                key: unwrapped,
-            },
-        );
+                // Create the Key Encryption Key (KEK) from the master key
+                let unwrapped = aes_kw_unwrap(master_key, &WrappedKey::from(wpky.clone()))
+                    .map_err(|_| BackupError::KeyUnwrapFailed(class_id))?;
+
+                // Insert the unwrapped key into the map
+                unlocked_keys.insert(
+                    class_id,
+                    ProtectionClassKey {
+                        class_id,
+                        key: unwrapped,
+                    },
+                );
+            }
+            _ => continue,
+        }
     }
     Ok(unlocked_keys)
 }

src/backup/mod.rs

@@ -477,7 +477,7 @@ impl Backup {
     pub fn decrypt_entry_stream(
         &self,
         entry: &BackupFileEntry,
-    ) -> Result<crypto::AesCbcDecryptReader<BufReader<File>>> {
+    ) -> Result<AesCbcDecryptReader<BufReader<File>>> {
         if !self.is_encrypted() {
             return Err(BackupError::NotEncrypted);
         }

src/backup/models/manifest/manifest_plist.rs

@@ -68,35 +68,39 @@ impl Manifest {
     /// # Ok::<(), crabapple::error::BackupError>(())
     /// ```
     pub fn from_manifest_data(manifest_data: ManifestData, auth: &Authentication) -> Result<Self> {
-        let (main_decryption_key, unlocked_class_keys) = if manifest_data.is_encrypted {
-            let backup_key_ring = manifest_data.key_ring.as_ref().ok_or_else(|| {
-                BackupError::MissingPlistKey(
-                    "BackupKeyBag (required for encrypted backup)".to_string(),
-                )
-            })?;
-
-            let master_key = match auth {
-                Authentication::Password(password) => derive_key_from_password(
+        // Determine decryption strategy based on whether the backup is encrypted and the provided authentication.
+        let (main_decryption_key, unlocked_class_keys) = match (manifest_data.is_encrypted, auth) {
+            (true, Authentication::Password(password)) => {
+                let backup_key_ring = manifest_data.key_ring.as_ref().ok_or_else(|| {
+                    BackupError::MissingPlistKey(
+                        "BackupKeyBag (required for encrypted backup)".to_string(),
+                    )
+                })?;
+                let master_key = derive_key_from_password(
                     password.as_bytes(),
                     &backup_key_ring.dpsl,
                     backup_key_ring.dpic,
                     &backup_key_ring.salt,
                     backup_key_ring.iter,
-                )?,
-                Authentication::DerivedKey(key_hex) => hex_decode(key_hex)?.into(),
-                Authentication::None => return Err(BackupError::PasswordOrKeyIncorrect),
-            };
-
-            let unlocked_keys_map = unlock_keys_from_manifest(&master_key, &manifest_data)
-                .map_err(|_| BackupError::PasswordOrKeyIncorrect)?;
-
-            (Some(master_key), Some(unlocked_keys_map))
-        } else {
-            // Error if the backup is not encrypted but an authentication method is provided
-            if !matches!(auth, Authentication::None) {
-                return Err(BackupError::NotEncrypted);
+                )?;
+                let unlocked_keys_map = unlock_keys_from_manifest(&master_key, &manifest_data)
+                    .map_err(|_| BackupError::PasswordOrKeyIncorrect)?;
+                (Some(master_key), Some(unlocked_keys_map))
+            }
+            (true, Authentication::DerivedKey(key_hex)) => {
+                manifest_data.key_ring.as_ref().ok_or_else(|| {
+                    BackupError::MissingPlistKey(
+                        "BackupKeyBag (required for encrypted backup)".to_string(),
+                    )
+                })?;
+                let master_key = hex_decode(key_hex)?.into();
+                let unlocked_keys_map = unlock_keys_from_manifest(&master_key, &manifest_data)
+                    .map_err(|_| BackupError::PasswordOrKeyIncorrect)?;
+                (Some(master_key), Some(unlocked_keys_map))
             }
-            (None, None)
+            (true, Authentication::None) => return Err(BackupError::PasswordOrKeyRequired),
+            (false, Authentication::None) => (None, None),
+            (false, _) => return Err(BackupError::NotEncrypted),
         };
 
         Ok(Self {

src/error.rs

@@ -63,7 +63,12 @@ pub enum BackupError {
     KeyUnwrapFailed(u32),
 
     /// Cryptographic data had an unexpected length.
-    InvalidCryptoDataLength { expected: usize, actual: usize },
+    InvalidCryptoDataLength {
+        /// The expected length of the data in bytes.
+        expected: usize,
+        /// The actual length of the data in bytes.
+        actual: usize,
+    },
 
     /// Invalid TLV data encountered.
     InvalidTlvData(String),