From 3dcb8d017b7805a7d7849a772b45ce9a6c6cc260 Mon Sep 17 00:00:00 2001
From: ant013 <ant013@mail.ru>
Date: Thu, 14 Sep 2023 14:03:27 +0600
Subject: [PATCH] Add enabled wallets to WalletBackup

---
 .../Core/Crypto/WalletBackup.swift            | 55 +++++++++++++++-
 .../Managers/CloudAccountBackupManager.swift  |  4 +-
 .../Core/Storage/LocalStorage.swift           | 63 ++++++++++++++++++-
 .../Backup/ICloud/BackupCloudModule.swift     |  2 +-
 .../BackupCloudPassphraseService.swift        |  7 ++-
 .../Backup/ICloud/WalletBackupConverter.swift | 62 +++++++++---------
 .../RestoreCloud/RestoreCloudModule.swift     |  1 +
 .../RestoreCloudPassphraseModule.swift        |  1 +
 .../RestoreCloudPassphraseService.swift       | 26 +++++++-
 .../RestoreCloudPassphraseViewModel.swift     | 10 ++-
 10 files changed, 185 insertions(+), 46 deletions(-)

diff --git a/UnstoppableWallet/UnstoppableWallet/Core/Crypto/WalletBackup.swift b/UnstoppableWallet/UnstoppableWallet/Core/Crypto/WalletBackup.swift
index 375f14e63f..63807ee6bf 100644
--- a/UnstoppableWallet/UnstoppableWallet/Core/Crypto/WalletBackup.swift
+++ b/UnstoppableWallet/UnstoppableWallet/Core/Crypto/WalletBackup.swift
@@ -7,9 +7,11 @@ class WalletBackup: Codable {
     let isManualBackedUp: Bool
     let version: Int
     let timestamp: TimeInterval?
+    let enabledWallets: [EnabledWallet]
 
     enum CodingKeys: String, CodingKey {
         case crypto
+        case enabledWallets = "enabled_wallets"
         case id
         case type
         case isManualBackedUp = "manual_backup"
@@ -17,8 +19,9 @@ class WalletBackup: Codable {
         case timestamp
     }
 
-    init(crypto: WalletBackupCrypto, id: String, type: AccountType.Abstract, isManualBackedUp: Bool, version: Int, timestamp: TimeInterval) {
+    init(crypto: WalletBackupCrypto, enabledWallets: [EnabledWallet], id: String, type: AccountType.Abstract, isManualBackedUp: Bool, version: Int, timestamp: TimeInterval) {
         self.crypto = crypto
+        self.enabledWallets = enabledWallets
         self.id = id
         self.type = type
         self.isManualBackedUp = isManualBackedUp
@@ -29,6 +32,7 @@ class WalletBackup: Codable {
     required init(from decoder: Decoder) throws {
         let container = try decoder.container(keyedBy: CodingKeys.self)
         crypto = try container.decode(WalletBackupCrypto.self, forKey: .crypto)
+        enabledWallets = (try? container.decode([EnabledWallet].self, forKey: .enabledWallets)) ?? []
         id = try container.decode(String.self, forKey: .id)
         type = try container.decode(AccountType.Abstract.self, forKey: .type)
         let isManualBackedUp = try? container.decode(Bool.self, forKey: .isManualBackedUp)
@@ -40,12 +44,59 @@ class WalletBackup: Codable {
     func encode(to encoder: Encoder) throws {
         var container = encoder.container(keyedBy: CodingKeys.self)
         try container.encode(crypto, forKey: .crypto)
+        try container.encode(enabledWallets, forKey: .enabledWallets)
         try container.encode(id, forKey: .id)
         try container.encode(type, forKey: .type)
         try container.encode(isManualBackedUp, forKey: .isManualBackedUp)
         try container.encode(version, forKey: .version)
         try container.encode(timestamp, forKey: .timestamp)
     }
-
 }
 
+extension WalletBackup {
+    struct EnabledWallet: Codable {
+        let tokenQueryId: String
+        let coinName: String?
+        let coinCode: String?
+        let tokenDecimals: Int?
+
+        enum CodingKeys: String, CodingKey {
+            case tokenQueryId = "token_query_id"
+            case coinName = "coin_name"
+            case coinCode = "coin_code"
+            case tokenDecimals = "decimals"
+        }
+
+        init(tokenQueryId: String, coinName: String?, coinCode: String?, tokenDecimals: Int?) {
+            self.tokenQueryId = tokenQueryId
+            self.coinName = coinName
+            self.coinCode = coinCode
+            self.tokenDecimals = tokenDecimals
+        }
+
+        init(_ wallet: Wallet) {
+            tokenQueryId = wallet.token.tokenQuery.id
+            coinName = wallet.coin.name
+            coinCode = wallet.coin.code
+            tokenDecimals = wallet.decimals
+        }
+
+        init(from decoder: Decoder) throws {
+            let container = try decoder.container(keyedBy: CodingKeys.self)
+            let tokenQueryId = try container.decode(String.self, forKey: .tokenQueryId)
+            let coinName = try? container.decode(String.self, forKey: .coinName)
+            let coinCode = try container.decode(String.self, forKey: .coinCode)
+            let tokenDecimals = try container.decode(Int.self, forKey: .tokenDecimals)
+
+            self.init(tokenQueryId: tokenQueryId, coinName: coinName, coinCode: coinCode, tokenDecimals: tokenDecimals)
+        }
+
+        public func encode(to encoder: Encoder) throws {
+            var container = encoder.container(keyedBy: CodingKeys.self)
+            try container.encode(tokenQueryId, forKey: .tokenQueryId)
+            try container.encode(coinName, forKey: .coinName)
+            try container.encode(coinCode, forKey: .coinCode)
+            try container.encode(tokenDecimals, forKey: .tokenDecimals)
+        }
+    }
+}
diff --git a/UnstoppableWallet/UnstoppableWallet/Core/Managers/CloudAccountBackupManager.swift b/UnstoppableWallet/UnstoppableWallet/Core/Managers/CloudAccountBackupManager.swift
index d88fcc347f..748dc386d6 100644
--- a/UnstoppableWallet/UnstoppableWallet/Core/Managers/CloudAccountBackupManager.swift
+++ b/UnstoppableWallet/UnstoppableWallet/Core/Managers/CloudAccountBackupManager.swift
@@ -134,14 +134,14 @@ extension CloudAccountBackupManager {
         iCloudUrl != nil
     }
 
-    func save(accountType: AccountType, isManualBackedUp: Bool, passphrase: String, name: String) throws {
+    func save(accountType: AccountType, wallets: [Wallet], isManualBackedUp: Bool, passphrase: String, name: String) throws {
         guard let iCloudUrl else {
             throw BackupError.urlNotAvailable
         }
 
         do {
             let name = name + Self.fileExtension
-            let encoded = try WalletBackupConverter.encode(accountType: accountType, isManualBackedUp: isManualBackedUp, passphrase: passphrase)
+            let encoded = try WalletBackupConverter.encode(accountType: accountType, wallets: wallets.map { WalletBackup.EnabledWallet($0) }, isManualBackedUp: isManualBackedUp, passphrase: passphrase)
 
             try fileStorage.write(directoryUrl: iCloudUrl, filename: name, data: encoded)
             logger?.log(level: .debug, message: "CloudAccountManager.downloadItems, save \(name)")
diff --git a/UnstoppableWallet/UnstoppableWallet/Core/Storage/LocalStorage.swift b/UnstoppableWallet/UnstoppableWallet/Core/Storage/LocalStorage.swift
index 36dafbf00d..e7549ad6cc 100644
--- a/UnstoppableWallet/UnstoppableWallet/Core/Storage/LocalStorage.swift
+++ b/UnstoppableWallet/UnstoppableWallet/Core/Storage/LocalStorage.swift
@@ -1,6 +1,6 @@
 import Foundation
-import StorageKit
 import MarketKit
+import StorageKit
 
 class LocalStorage {
     private let agreementAcceptedKey = "i_understand_key"
@@ -25,11 +25,9 @@ class LocalStorage {
     init(storage: StorageKit.ILocalStorage) {
         self.storage = storage
     }
-
 }
 
 extension LocalStorage {
-
     var debugLog: String? {
         get { storage.value(for: debugLogKey) }
         set { storage.set(value: newValue, for: debugLogKey) }
@@ -100,5 +98,64 @@ extension LocalStorage {
         get { storage.value(for: keyTelegramSupportRequested) ?? false }
         set { storage.set(value: newValue, for: keyTelegramSupportRequested) }
     }
+}
 
+extension LocalStorage {
+    var backup: [String: Any] {
+        var fields: [String: Any] = [
+            "lock_time_enabled": lockTimeEnabled.description,
+            "contacts_sync": remoteContactsSync.description,
+            "show_indicators": indicatorsShown.description,
+        ]
+
+        if let chartIndicators {
+            fields["chart_indicators"] = chartIndicators.hs.hexString
+        }
+
+        let providers: [String: String] = BlockchainType
+            .supported
+            .filter { !$0.allowedProviders.isEmpty }
+            .map { ($0.uid, defaultProvider(blockchainType: $0).rawValue) }
+            .reduce(into: [:]) { $0[$1.0] = $1.1 }
+
+        fields["swap_providers"] = providers
+
+        return fields
+    }
+
+    private func bool(value: Any?) -> Bool? {
+        if let string = value as? String,
+           let enabled = Bool(string) {
+            return enabled
+        }
+        return nil
+    }
+
+    func restore(backup _: [String: Any]) {
+        if let lockTime = bool(value: backup["lock_time_enabled"]) {
+            lockTimeEnabled = lockTime
+        }
+        if let contactsSync = bool(value: backup["contacts_sync"]) {
+            remoteContactsSync = contactsSync
+        }
+        if let showIndicators = bool(value: backup["show_indicators"]) {
+            indicatorsShown = showIndicators
+        }
+
+        if let chartIndicators = backup["chart_indicators"] as? String,
+           let data = chartIndicators.hs.hexData {
+            self.chartIndicators = data
+        }
+
+        if let providers = backup["swap_providers"] as? [String: String] {
+            providers.forEach { key, value in
+                let type = BlockchainType(uid: key)
+                if let provider = SwapModule.Dex.Provider(rawValue: value),
+                   !type.allowedProviders.isEmpty {
+
+                    setDefaultProvider(blockchainType: type, provider: provider)
+                }
+            }
+        }
+    }
 }
diff --git a/UnstoppableWallet/UnstoppableWallet/Modules/Backup/ICloud/BackupCloudModule.swift b/UnstoppableWallet/UnstoppableWallet/Modules/Backup/ICloud/BackupCloudModule.swift
index a82d186a2b..dabc1a6829 100644
--- a/UnstoppableWallet/UnstoppableWallet/Modules/Backup/ICloud/BackupCloudModule.swift
+++ b/UnstoppableWallet/UnstoppableWallet/Modules/Backup/ICloud/BackupCloudModule.swift
@@ -22,7 +22,7 @@ class BackupCloudModule {
     }
 
     static func backupPassword(account: Account, name: String) -> UIViewController {
-        let service = BackupCloudPassphraseService(iCloudManager: App.shared.cloudAccountBackupManager, account: account, name: name)
+        let service = BackupCloudPassphraseService(iCloudManager: App.shared.cloudAccountBackupManager, walletManager: App.shared.walletManager, account: account, name: name)
         let viewModel = BackupCloudPassphraseViewModel(service: service)
         let controller = BackupCloudPassphraseViewController(viewModel: viewModel)
 
diff --git a/UnstoppableWallet/UnstoppableWallet/Modules/Backup/ICloud/Passphrase/BackupCloudPassphraseService.swift b/UnstoppableWallet/UnstoppableWallet/Modules/Backup/ICloud/Passphrase/BackupCloudPassphraseService.swift
index ba6e86e757..2c522f8eee 100644
--- a/UnstoppableWallet/UnstoppableWallet/Modules/Backup/ICloud/Passphrase/BackupCloudPassphraseService.swift
+++ b/UnstoppableWallet/UnstoppableWallet/Modules/Backup/ICloud/Passphrase/BackupCloudPassphraseService.swift
@@ -2,14 +2,16 @@ import Foundation
 
 class BackupCloudPassphraseService {
     private let iCloudManager: CloudAccountBackupManager
+    private let walletManager: WalletManager
     private let account: Account
     private let name: String
 
     var passphrase: String = ""
     var passphraseConfirmation: String = ""
 
-    init(iCloudManager: CloudAccountBackupManager, account: Account, name: String) {
+    init(iCloudManager: CloudAccountBackupManager, walletManager: WalletManager, account: Account, name: String) {
         self.iCloudManager = iCloudManager
+        self.walletManager = walletManager
         self.account = account
         self.name = name
     }
@@ -41,7 +43,8 @@ extension BackupCloudPassphraseService {
         }
 
         do {
-            try iCloudManager.save(accountType: account.type, isManualBackedUp: account.backedUp, passphrase: passphrase, name: name)
+            let wallets = App.shared.walletManager.wallets(account: account)
+            try iCloudManager.save(accountType: account.type, wallets: wallets, isManualBackedUp: account.backedUp, passphrase: passphrase, name: name)
         } catch {
             if case .urlNotAvailable = error as? CloudAccountBackupManager.BackupError {
                 throw CreateError.urlNotAvailable
diff --git a/UnstoppableWallet/UnstoppableWallet/Modules/Backup/ICloud/WalletBackupConverter.swift b/UnstoppableWallet/UnstoppableWallet/Modules/Backup/ICloud/WalletBackupConverter.swift
index 806b88432b..19917aaeb7 100644
--- a/UnstoppableWallet/UnstoppableWallet/Modules/Backup/ICloud/WalletBackupConverter.swift
+++ b/UnstoppableWallet/UnstoppableWallet/Modules/Backup/ICloud/WalletBackupConverter.swift
@@ -1,43 +1,45 @@
 import Foundation
+import MarketKit
 
 class WalletBackupConverter {
-    private static let version = 1
+    private static let version = 2
 
-    static func encode(accountType: AccountType, isManualBackedUp: Bool, passphrase: String) throws -> Data {
+    static func encode(accountType: AccountType, wallets: [WalletBackup.EnabledWallet], isManualBackedUp: Bool, passphrase: String) throws -> Data {
         let message = accountType.uniqueId(hashed: false)
         let iv = BackupCryptoHelper.generateInitialVector().hs.hex
 
         let cipherText = try BackupCryptoHelper.AES128(
-                operation: .encrypt,
-                ivHex: iv,
-                pass: passphrase,
-                message: message,
-                kdf: .defaultBackup
+            operation: .encrypt,
+            ivHex: iv,
+            pass: passphrase,
+            message: message,
+            kdf: .defaultBackup
         )
         let encodedCipherText = cipherText.base64EncodedString()
         let mac = try BackupCryptoHelper.mac(
-                pass: passphrase,
-                message: encodedCipherText.hs.data,
-                kdf: .defaultBackup
+            pass: passphrase,
+            message: encodedCipherText.hs.data,
+            kdf: .defaultBackup
         )
 
         let crypto = WalletBackupCrypto(
-                cipher: BackupCryptoHelper.defaultCypher,
-                cipherParams: CipherParams(iv: iv),
-                cipherText: encodedCipherText,
-                kdf: BackupCryptoHelper.defaultKdf,
-                kdfParams: .defaultBackup,
-                mac: mac.hs.hex)
+            cipher: BackupCryptoHelper.defaultCypher,
+            cipherParams: CipherParams(iv: iv),
+            cipherText: encodedCipherText,
+            kdf: BackupCryptoHelper.defaultKdf,
+            kdfParams: .defaultBackup,
+            mac: mac.hs.hex
+        )
         let backup = WalletBackup(
-                crypto: crypto,
-                id: accountType.uniqueId().hs.hex,
-                type: AccountType.Abstract(accountType),
-                isManualBackedUp: isManualBackedUp,
-                version: Self.version,
-                timestamp: Date().timeIntervalSince1970
+            crypto: crypto,
+            enabledWallets: wallets,
+            id: accountType.uniqueId().hs.hex,
+            type: AccountType.Abstract(accountType),
+            isManualBackedUp: isManualBackedUp,
+            version: Self.version,
+            timestamp: Date().timeIntervalSince1970
         )
         return try JSONEncoder().encode(backup)
-
     }
 
     static func decode(data: Data, passphrase: String) throws -> AccountType {
@@ -48,11 +50,12 @@ class WalletBackupConverter {
         }
 
         let decryptData = try BackupCryptoHelper.AES128(
-                operation: .decrypt,
-                ivHex: backup.crypto.cipherParams.iv,
-                pass: passphrase,
-                message: message,
-                kdf: .defaultBackup)
+            operation: .decrypt,
+            ivHex: backup.crypto.cipherParams.iv,
+            pass: passphrase,
+            message: message,
+            kdf: .defaultBackup
+        )
 
         guard let accountType = AccountType.decode(uniqueId: decryptData, type: backup.type) else {
             throw CodingError.cantDecodeAccountType
@@ -60,14 +63,11 @@ class WalletBackupConverter {
 
         return accountType
     }
-
 }
 
 extension WalletBackupConverter {
-
     enum CodingError: Error {
         case cantDecodeCipherText
         case cantDecodeAccountType
     }
-
 }
diff --git a/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudModule.swift b/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudModule.swift
index 3555a45c19..b6fe4d0854 100644
--- a/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudModule.swift
+++ b/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudModule.swift
@@ -22,6 +22,7 @@ struct RestoreCloudModule {
         let name: String
         let accountType: AccountType
         let isManualBackedUp: Bool
+        let showSelectCoins: Bool
     }
 
 }
diff --git a/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudPassphrase/RestoreCloudPassphraseModule.swift b/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudPassphrase/RestoreCloudPassphraseModule.swift
index d6c399489f..a215dd5fc3 100644
--- a/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudPassphrase/RestoreCloudPassphraseModule.swift
+++ b/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudPassphrase/RestoreCloudPassphraseModule.swift
@@ -8,6 +8,7 @@ class RestoreCloudPassphraseModule {
                 iCloudManager: App.shared.cloudAccountBackupManager,
                 accountFactory: App.shared.accountFactory,
                 accountManager: App.shared.accountManager,
+                walletManager: App.shared.walletManager,
                 item: item
         )
         let viewModel = RestoreCloudPassphraseViewModel(service: service)
diff --git a/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudPassphrase/RestoreCloudPassphraseService.swift b/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudPassphrase/RestoreCloudPassphraseService.swift
index 238e4c3594..875feebdc0 100644
--- a/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudPassphrase/RestoreCloudPassphraseService.swift
+++ b/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudPassphrase/RestoreCloudPassphraseService.swift
@@ -4,16 +4,34 @@ class RestoreCloudPassphraseService {
     private let iCloudManager: CloudAccountBackupManager
     private let accountFactory: AccountFactory
     private let accountManager: AccountManager
+    private let walletManager: WalletManager
 
     private let restoredBackup: RestoreCloudModule.RestoredBackup
 
     var passphrase: String = ""
 
-    init(iCloudManager: CloudAccountBackupManager, accountFactory: AccountFactory, accountManager: AccountManager, item: RestoreCloudModule.RestoredBackup) {
+    init(iCloudManager: CloudAccountBackupManager, accountFactory: AccountFactory, accountManager: AccountManager, walletManager: WalletManager, item: RestoreCloudModule.RestoredBackup) {
         self.iCloudManager = iCloudManager
         self.accountFactory = accountFactory
         self.accountManager = accountManager
-        self.restoredBackup = item
+        self.walletManager = walletManager
+        restoredBackup = item
+    }
+
+    private func createAccount(accountType: AccountType) {
+        let account = accountFactory.account(type: accountType, origin: .restored, backedUp: restoredBackup.walletBackup.isManualBackedUp, name: restoredBackup.name)
+        accountManager.save(account: account)
+
+        let wallets = restoredBackup.walletBackup.enabledWallets.map {
+            EnabledWallet(
+                    tokenQueryId: $0.tokenQueryId,
+                    accountId: account.id,
+                    coinName: $0.coinName,
+                    coinCode: $0.coinCode,
+                    tokenDecimals: $0.tokenDecimals
+            )
+        }
+        walletManager.save(enabledWallets: wallets)
     }
 
 }
@@ -77,10 +95,12 @@ extension RestoreCloudPassphraseService {
                 accountManager.save(account: account)
                 return .success
             default:
+                createAccount(accountType: accountType)
                 return .restoredAccount(RestoreCloudModule.RestoredAccount(
                         name: restoredBackup.name,
                         accountType: accountType,
-                        isManualBackedUp: restoredBackup.walletBackup.isManualBackedUp
+                        isManualBackedUp: restoredBackup.walletBackup.isManualBackedUp,
+                        showSelectCoins: restoredBackup.walletBackup.enabledWallets.isEmpty
                 ))
             }
         } catch {
diff --git a/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudPassphrase/RestoreCloudPassphraseViewModel.swift b/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudPassphrase/RestoreCloudPassphraseViewModel.swift
index 5aa92a7d20..0f9d9da4e9 100644
--- a/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudPassphrase/RestoreCloudPassphraseViewModel.swift
+++ b/UnstoppableWallet/UnstoppableWallet/Modules/RestoreAccount/RestoreCloud/RestoreCloudPassphrase/RestoreCloudPassphraseViewModel.swift
@@ -67,8 +67,14 @@ extension RestoreCloudPassphraseViewModel {
                 self?.processing = false
 
                 switch result {
-                case .success: self?.successSubject.send()
-                case .restoredAccount(let account): self?.openSelectCoinsSubject.send(account)
+                case .success:
+                    self?.successSubject.send()
+                case .restoredAccount(let account):
+                    if account.showSelectCoins {
+                        self?.openSelectCoinsSubject.send(account)
+                    } else {
+                        self?.successSubject.send()
+                    }
                 }
             } catch {
                 switch (error as? RestoreCloudPassphraseService.RestoreError) {