olas-operate-middleware 0.13.1__py3-none-any.whl → 0.13.3__py3-none-any.whl

operate/wallet/wallet_recovery_manager.py

@@ -19,6 +19,7 @@
 
 """Wallet recovery manager"""
 
+import enum
 import shutil
 import typing as t
 import uuid
@@ -49,6 +50,19 @@ RECOVERY_NEW_OBJECTS_DIR = "new"
 RECOVERY_OLD_OBJECTS_DIR = "old"
 
 
+class WalletRecoveryStatus(str, enum.Enum):
+    """ProviderRequestStatus"""
+
+    NOT_PREPARED = "NOT_PREPARED"
+    PREPARED = "PREPARED"
+    IN_PROGRESS = "IN_PROGRESS"
+    COMPLETED = "COMPLETED"
+
+    def __str__(self) -> str:
+        """__str__"""
+        return self.value
+
+
 class WalletRecoveryError(Exception):
     """WalletRecoveryError"""
 
@@ -125,20 +139,16 @@ class WalletRecoveryManager:
                     )
 
         last_prepared_bundle_id = self.data.last_prepared_bundle_id
-        if last_prepared_bundle_id is not None:
-            (
-                _,
-                num_safes_with_new_wallet,
-                _,
-                _,
-            ) = self._get_swap_status(last_prepared_bundle_id)
-            if num_safes_with_new_wallet > 0:
-                self.logger.info(
-                    f"[WALLET RECOVERY MANAGER] Uncompleted bundle {last_prepared_bundle_id} has Safes with new wallet."
-                )
-                return self._load_bundle(
-                    bundle_id=last_prepared_bundle_id, new_password=new_password
-                )
+        if (
+            last_prepared_bundle_id is not None
+            and self.status()["num_safes_with_new_wallet"] > 0
+        ):
+            self.logger.info(
+                f"[WALLET RECOVERY MANAGER] Uncompleted bundle {last_prepared_bundle_id} has Safes with new wallet."
+            )
+            return self._load_bundle(
+                bundle_id=last_prepared_bundle_id, new_password=new_password
+            )
 
         # Create new recovery bundle
         bundle_id = f"{RECOVERY_BUNDLE_PREFIX}{str(uuid.uuid4())}"
@@ -178,23 +188,45 @@ class WalletRecoveryManager:
         )
         return self._load_bundle(bundle_id=bundle_id, new_password=new_password)
 
-    def _get_swap_status(self, bundle_id: str) -> t.Tuple[int, int, int, int]:
+    def _load_bundle(  # pylint: disable=too-many-locals
+        self, bundle_id: str, new_password: t.Optional[str] = None
+    ) -> t.Dict:
         new_root = self.path / bundle_id / RECOVERY_NEW_OBJECTS_DIR
+
+        new_user_account = UserAccount.load(new_root / USER_JSON)
+        if new_password is not None and not new_user_account.is_valid(
+            password=new_password
+        ):
+            raise ValueError(MSG_INVALID_PASSWORD)
+
         new_wallets_path = new_root / WALLETS_DIR
-        new_wallet_manager = MasterWalletManager(path=new_wallets_path, password=None)
+        new_wallet_manager = MasterWalletManager(
+            path=new_wallets_path, password=new_password
+        )
 
         num_safes = 0
         num_safes_with_new_wallet = 0
         num_safes_with_old_wallet = 0
         num_safes_with_both_wallets = 0
+        backup_owner_sets = set()
 
+        wallets = []
         for wallet in self.wallet_manager:
             new_wallet = next(
                 (w for w in new_wallet_manager if w.ledger_type == wallet.ledger_type)
             )
+            new_mnemonic = None
+            if new_password:
+                new_mnemonic = new_wallet.decrypt_mnemonic(password=new_password)
+
+            wallet_json = wallet.json
+
             for chain, safe in wallet.safes.items():
+                chain_str = chain.value
                 ledger_api = get_default_ledger_api(chain)
                 owners = get_owners(ledger_api=ledger_api, safe=safe)
+                backup_owners = list(set(owners) - {wallet.address, new_wallet.address})
+                backup_owner_sets.add(frozenset(backup_owners))
 
                 num_safes += 1
                 if new_wallet.address in owners and wallet.address in owners:
@@ -204,42 +236,54 @@ class WalletRecoveryManager:
                 if wallet.address in owners:
                     num_safes_with_old_wallet += 1
 
-        return (
-            num_safes,
-            num_safes_with_new_wallet,
-            num_safes_with_old_wallet,
-            num_safes_with_both_wallets,
-        )
-
-    def _load_bundle(self, bundle_id: str, new_password: str) -> t.Dict:
-        new_root = self.path / bundle_id / RECOVERY_NEW_OBJECTS_DIR
-
-        new_user_account = UserAccount.load(new_root / USER_JSON)
-        if not new_user_account.is_valid(password=new_password):
-            raise ValueError(MSG_INVALID_PASSWORD)
-
-        new_wallets_path = new_root / WALLETS_DIR
-        new_wallet_manager = MasterWalletManager(
-            path=new_wallets_path, password=new_password
-        )
+                wallet_json["safes"][chain_str] = {
+                    safe: {
+                        "owners": owners,
+                        "backup_owners": backup_owners,
+                        "owner_to_remove": (
+                            wallet.address if wallet.address in owners else None
+                        ),
+                        "owner_to_add": (
+                            new_wallet.address
+                            if new_wallet.address not in owners
+                            else None
+                        ),
+                    }
+                }
 
-        wallets = []
-        for wallet in self.wallet_manager:
-            ledger_type = wallet.ledger_type
-            new_wallet = new_wallet_manager.load(ledger_type=ledger_type)
-            new_mnemonic = None
-            if new_password:
-                new_mnemonic = new_wallet.decrypt_mnemonic(password=new_password)
             wallets.append(
                 {
-                    "current_wallet": wallet.json,
+                    "current_wallet": wallet_json,
                     "new_wallet": new_wallet.json,
                     "new_mnemonic": new_mnemonic,
                 }
             )
+
+        if num_safes_with_new_wallet == 0:
+            status = WalletRecoveryStatus.PREPARED
+        elif num_safes_with_new_wallet < num_safes:
+            status = WalletRecoveryStatus.IN_PROGRESS
+        else:
+            status = WalletRecoveryStatus.COMPLETED
+
         return {
             "id": bundle_id,
             "wallets": wallets,
+            "status": status,
+            "all_safes_have_backup_owner": all(
+                len(owners) >= 1 for owners in backup_owner_sets
+            ),
+            "consistent_backup_owner": len(backup_owner_sets) == 1,
+            "consistent_backup_owner_count": all(
+                len(owners) == 1 for owners in backup_owner_sets
+            ),
+            "prepared": bundle_id is not None,
+            "has_swaps": num_safes_with_new_wallet > 0,
+            "has_pending_swaps": num_safes_with_new_wallet < num_safes,
+            "num_safes": num_safes,
+            "num_safes_with_new_wallet": num_safes_with_new_wallet,
+            "num_safes_with_old_wallet": num_safes_with_old_wallet,
+            "num_safes_with_both_wallets": num_safes_with_both_wallets,
         }
 
     def recovery_requirements(  # pylint: disable=too-many-locals
@@ -317,27 +361,37 @@ class WalletRecoveryManager:
     def status(self) -> t.Dict[str, t.Any]:
         """Get recovery status."""
         bundle_id = self.data.last_prepared_bundle_id
-        if not bundle_id:
+        if bundle_id is None:
+            backup_owner_sets = set()
+            for wallet in self.wallet_manager:
+                for chain, safe in wallet.safes.items():
+                    ledger_api = get_default_ledger_api(chain)
+                    owners = get_owners(ledger_api=ledger_api, safe=safe)
+                    backup_owners = list(set(owners) - {wallet.address})
+                    backup_owner_sets.add(frozenset(backup_owners))
+
             return {
+                "id": None,
+                "wallets": [],
+                "status": WalletRecoveryStatus.NOT_PREPARED,
+                "all_safes_have_backup_owner": all(
+                    len(owners) >= 1 for owners in backup_owner_sets
+                ),
+                "consistent_backup_owner": len(backup_owner_sets) == 1,
+                "consistent_backup_owner_count": all(
+                    len(owners) == 1 for owners in backup_owner_sets
+                ),
                 "prepared": False,
                 "bundle_id": bundle_id,
                 "has_swaps": False,
                 "has_pending_swaps": False,
+                "num_safes": 0,
+                "num_safes_with_new_wallet": 0,
+                "num_safes_with_old_wallet": 0,
+                "num_safes_with_both_wallets": 0,
             }
 
-        (
-            num_safes,
-            num_safes_with_new_wallet,
-            _,
-            _,
-        ) = self._get_swap_status(bundle_id)
-
-        return {
-            "prepared": bundle_id is not None,
-            "bundle_id": bundle_id,
-            "has_swaps": num_safes_with_new_wallet > 0,
-            "has_pending_swaps": num_safes_with_new_wallet < num_safes,
-        }
+        return self._load_bundle(bundle_id=bundle_id)
 
     def complete_recovery(  # pylint: disable=too-many-locals,too-many-statements
         self, raise_if_inconsistent_owners: bool = True