From de6d0f308498a2fcaca88c8dddba4920c682045e Mon Sep 17 00:00:00 2001
From: Joel-David <git@joeldavid.dev>
Date: Wed, 19 Jun 2024 00:13:47 +0800
Subject: [PATCH] cleanup vaultservice

---
 Chronos/Services/CryptoService.swift |  4 ++--
 Chronos/Services/VaultService.swift  | 29 +++++++++++++++++++++++++---
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/Chronos/Services/CryptoService.swift b/Chronos/Services/CryptoService.swift
index 67eb296..3487f83 100644
--- a/Chronos/Services/CryptoService.swift
+++ b/Chronos/Services/CryptoService.swift
@@ -41,8 +41,8 @@ public class CryptoService {
         }
     }
 
-    func unwrapMasterKeyWithUserPassword(password: [UInt8], isRestore: Bool = false) async -> Bool {
-        guard let vault = vaultService.getVault(context: nil, isRestore: isRestore) else {
+    func unwrapMasterKeyWithUserPassword(password: [UInt8], isRestore _: Bool = false) async -> Bool {
+        guard let vault = vaultService.getVaultFromCloudContainer() else {
             return false
         }
 
diff --git a/Chronos/Services/VaultService.swift b/Chronos/Services/VaultService.swift
index 561deb5..6a784f7 100644
--- a/Chronos/Services/VaultService.swift
+++ b/Chronos/Services/VaultService.swift
@@ -27,17 +27,17 @@ public class VaultService {
         }
     }
 
-    func getVault(context: ModelContext? = nil, isRestore: Bool = false) -> Vault? {
+    func getVaultFromCloudContainer() -> Vault? {
         guard let vaultId: UUID = stateService.getVaultId() else {
             logger.error("vaultId not found in AppStorage")
             return nil
         }
 
         let predicate = #Predicate<Vault> { $0.vaultId == vaultId }
-        let _context = context ?? ModelContext(swiftDataService.getModelContainer(isRestore: isRestore))
+        let context = ModelContext(swiftDataService.getCloudModelContainer())
 
         do {
-            let vaultArr = try _context.fetch(FetchDescriptor<Vault>(predicate: predicate))
+            let vaultArr = try context.fetch(FetchDescriptor<Vault>(predicate: predicate))
             if let vault = vaultArr.first {
                 logger.info("Returning vault \(vault.name)")
                 return vault
@@ -53,6 +53,29 @@ public class VaultService {
 }
 
 extension VaultService {
+    private func getVault(context: ModelContext) -> Vault? {
+        guard let vaultId: UUID = stateService.getVaultId() else {
+            logger.error("vaultId not found in AppStorage")
+            return nil
+        }
+
+        let predicate = #Predicate<Vault> { $0.vaultId == vaultId }
+
+        do {
+            let vaultArr = try context.fetch(FetchDescriptor<Vault>(predicate: predicate))
+            if let vault = vaultArr.first {
+                logger.info("Returning vault \(vault.name)")
+                return vault
+            } else {
+                logger.error("No vaults found with the given vaultId")
+                return nil
+            }
+        } catch {
+            logger.error("Failed to fetch vaults: \(error.localizedDescription)")
+            return nil
+        }
+    }
+
     func insertEncryptedToken(_ encryptedToken: EncryptedToken) {
         let context = ModelContext(swiftDataService.getModelContainer())