olas-operate-middleware 0.10.20__py3-none-any.whl → 0.11.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
operate/wallet/master.py CHANGED
@@ -20,16 +20,13 @@
20
20
  """Master key implementation"""
21
21
 
22
22
  import json
23
- import logging
24
23
  import os
25
24
  import typing as t
26
- from copy import deepcopy
27
25
  from dataclasses import dataclass, field
28
26
  from pathlib import Path
29
27
 
30
28
  from aea.crypto.base import Crypto, LedgerApi
31
- from aea.crypto.registries import make_ledger_api
32
- from aea_ledger_ethereum import DEFAULT_GAS_PRICE_STRATEGIES, EIP1559, GWEI, to_wei
29
+ from aea.helpers.logging import setup_logger
33
30
  from aea_ledger_ethereum.ethereum import EthereumApi, EthereumCrypto
34
31
  from autonomy.chain.base import registry_contracts
35
32
  from autonomy.chain.config import ChainType as ChainProfile
@@ -43,18 +40,19 @@ from operate.constants import (
43
40
  ZERO_ADDRESS,
44
41
  )
45
42
  from operate.ledger import (
46
- get_default_rpc,
43
+ get_default_ledger_api,
44
+ make_chain_ledger_api,
47
45
  update_tx_with_gas_estimate,
48
46
  update_tx_with_gas_pricing,
49
47
  )
50
- from operate.ledger.profiles import ERC20_TOKENS, OLAS, USDC
51
- from operate.operate_types import Chain, LedgerType
48
+ from operate.ledger.profiles import DUST, ERC20_TOKENS, format_asset_amount
49
+ from operate.operate_types import Chain, EncryptedData, LedgerType
52
50
  from operate.resource import LocalResource
53
51
  from operate.utils import create_backup
54
52
  from operate.utils.gnosis import add_owner
55
53
  from operate.utils.gnosis import create_safe as create_gnosis_safe
56
54
  from operate.utils.gnosis import (
57
- drain_eoa,
55
+ estimate_transfer_tx_fee,
58
56
  get_asset_balance,
59
57
  get_owners,
60
58
  remove_owner,
@@ -64,6 +62,9 @@ from operate.utils.gnosis import transfer as transfer_from_safe
64
62
  from operate.utils.gnosis import transfer_erc20_from_safe
65
63
 
66
64
 
65
+ logger = setup_logger(name="master_wallet")
66
+
67
+
67
68
  # TODO Organize exceptions definition
68
69
  class InsufficientFundsException(Exception):
69
70
  """Insufficient funds exception."""
@@ -81,6 +82,7 @@ class MasterWallet(LocalResource):
81
82
  safe_nonce: t.Optional[int] = None
82
83
 
83
84
  _key: str
85
+ _mnemonic: str
84
86
  _crypto: t.Optional[Crypto] = None
85
87
  _password: t.Optional[str] = None
86
88
  _crypto_cls: t.Type[Crypto]
@@ -109,59 +111,42 @@ class MasterWallet(LocalResource):
109
111
  """Key path."""
110
112
  return self.path / self._key
111
113
 
114
+ @property
115
+ def mnemonic_path(self) -> Path:
116
+ """Mnemonic path."""
117
+ return self.path / self._mnemonic
118
+
119
+ @staticmethod
112
120
  def ledger_api(
113
- self,
114
121
  chain: Chain,
115
122
  rpc: t.Optional[str] = None,
116
123
  ) -> LedgerApi:
117
124
  """Get ledger api object."""
118
- gas_price_strategies = deepcopy(DEFAULT_GAS_PRICE_STRATEGIES)
119
- if chain in (Chain.BASE, Chain.MODE, Chain.OPTIMISM):
120
- gas_price_strategies[EIP1559]["fallback_estimate"]["maxFeePerGas"] = to_wei(
121
- 5, GWEI
122
- )
123
-
124
- return make_ledger_api(
125
- self.ledger_type.name.lower(),
126
- address=(rpc or get_default_rpc(chain=chain)),
127
- chain_id=chain.id,
128
- gas_price_strategies=gas_price_strategies,
129
- )
125
+ if not rpc:
126
+ return get_default_ledger_api(chain=chain)
127
+ return make_chain_ledger_api(chain=chain, rpc=rpc)
130
128
 
131
- def transfer(
129
+ def transfer( # pylint: disable=too-many-arguments
132
130
  self,
133
131
  to: str,
134
132
  amount: int,
135
133
  chain: Chain,
134
+ asset: str = ZERO_ADDRESS,
136
135
  from_safe: bool = True,
137
136
  rpc: t.Optional[str] = None,
138
137
  ) -> t.Optional[str]:
139
138
  """Transfer funds to the given account."""
140
139
  raise NotImplementedError()
141
140
 
142
- # pylint: disable=too-many-arguments
143
- def transfer_erc20(
144
- self,
145
- token: str,
146
- to: str,
147
- amount: int,
148
- chain: Chain,
149
- from_safe: bool = True,
150
- rpc: t.Optional[str] = None,
151
- ) -> t.Optional[str]:
152
- """Transfer funds to the given account."""
153
- raise NotImplementedError()
154
-
155
- def transfer_asset(
141
+ def transfer_from_safe_then_eoa(
156
142
  self,
157
143
  to: str,
158
144
  amount: int,
159
145
  chain: Chain,
160
146
  asset: str = ZERO_ADDRESS,
161
- from_safe: bool = True,
162
147
  rpc: t.Optional[str] = None,
163
- ) -> t.Optional[str]:
164
- """Transfer erc20/native assets to the given account."""
148
+ ) -> t.List[str]:
149
+ """Transfer assets to the given account using Safe balance first, and EOA balance for leftover."""
165
150
  raise NotImplementedError()
166
151
 
167
152
  def drain(
@@ -179,6 +164,10 @@ class MasterWallet(LocalResource):
179
164
  """Create a new master wallet."""
180
165
  raise NotImplementedError()
181
166
 
167
+ def decrypt_mnemonic(self, password: str) -> t.Optional[t.List[str]]:
168
+ """Retrieve the mnemonic"""
169
+ raise NotImplementedError()
170
+
182
171
  def create_safe(
183
172
  self,
184
173
  chain: Chain,
@@ -217,6 +206,24 @@ class MasterWallet(LocalResource):
217
206
  """Updates password using the mnemonic."""
218
207
  raise NotImplementedError()
219
208
 
209
+ def get_balance(
210
+ self, chain: Chain, asset: str = ZERO_ADDRESS, from_safe: bool = True
211
+ ) -> int:
212
+ """Get wallet balance on a given chain."""
213
+ if from_safe:
214
+ if chain not in self.safes:
215
+ raise ValueError(f"Wallet does not have a Safe on chain {chain}.")
216
+
217
+ address = self.safes[chain]
218
+ else:
219
+ address = self.address
220
+
221
+ return get_asset_balance(
222
+ ledger_api=get_default_ledger_api(chain),
223
+ asset_address=asset,
224
+ address=address,
225
+ )
226
+
220
227
  # TODO move to resource.py if used in more resources similarly
221
228
  @property
222
229
  def extended_json(self) -> t.Dict:
@@ -243,12 +250,61 @@ class EthereumMasterWallet(MasterWallet):
243
250
 
244
251
  _file = ledger_type.config_file
245
252
  _key = ledger_type.key_file
253
+ _mnemonic = ledger_type.mnemonic_file
246
254
  _crypto_cls = EthereumCrypto
247
255
 
256
+ def _pre_transfer_checks(
257
+ self,
258
+ to: str,
259
+ amount: int,
260
+ chain: Chain,
261
+ from_safe: bool,
262
+ asset: str = ZERO_ADDRESS,
263
+ ) -> str:
264
+ """Checks conditions before transfer. Returns the to address checksummed."""
265
+ if amount <= 0:
266
+ raise ValueError(
267
+ "Transfer amount must be greater than zero, not transferring."
268
+ )
269
+
270
+ to = Web3().to_checksum_address(to)
271
+ if from_safe and chain not in self.safes:
272
+ raise ValueError(f"Wallet does not have a Safe on chain {chain}.")
273
+
274
+ balance = self.get_balance(chain=chain, asset=asset, from_safe=from_safe)
275
+ if balance < amount:
276
+ source = "Master Safe" if from_safe else " Master EOA"
277
+ source_address = self.safes[chain] if from_safe else self.address
278
+ raise InsufficientFundsException(
279
+ f"Cannot transfer {format_asset_amount(chain, asset, amount)} from {source} {source_address} to {to} on chain {chain.name}. "
280
+ f"Balance: {format_asset_amount(chain, asset, balance)}. Missing: {format_asset_amount(chain, asset, amount - balance)}."
281
+ )
282
+
283
+ return to
284
+
248
285
  def _transfer_from_eoa(
249
286
  self, to: str, amount: int, chain: Chain, rpc: t.Optional[str] = None
250
287
  ) -> t.Optional[str]:
251
288
  """Transfer funds from EOA wallet."""
289
+ balance = self.get_balance(chain=chain, from_safe=False)
290
+ tx_fee = estimate_transfer_tx_fee(
291
+ chain=chain, sender_address=self.address, to=to
292
+ )
293
+ if balance - tx_fee < amount <= balance:
294
+ # we assume that the user wants to drain the EOA
295
+ # we also account for dust here because withdraw call use some EOA balance to drain the safes first
296
+ amount = balance - tx_fee
297
+ if amount <= 0:
298
+ logger.warning(
299
+ f"Not enough balance to cover gas fees for transfer of {amount} on chain {chain} from EOA {self.address}. "
300
+ f"Balance is {balance}, estimated fee is {tx_fee}. Not transferring."
301
+ )
302
+ return None
303
+
304
+ to = self._pre_transfer_checks(
305
+ to=to, amount=amount, chain=chain, from_safe=False
306
+ )
307
+
252
308
  ledger_api = t.cast(EthereumApi, self.ledger_api(chain=chain, rpc=rpc))
253
309
  tx_helper = TxSettler(
254
310
  ledger_api=ledger_api,
@@ -292,13 +348,14 @@ class EthereumMasterWallet(MasterWallet):
292
348
  self, to: str, amount: int, chain: Chain, rpc: t.Optional[str] = None
293
349
  ) -> t.Optional[str]:
294
350
  """Transfer funds from safe wallet."""
295
- if self.safes is None:
296
- raise ValueError("Safes not initialized")
351
+ to = self._pre_transfer_checks(
352
+ to=to, amount=amount, chain=chain, from_safe=True
353
+ )
297
354
 
298
355
  return transfer_from_safe(
299
356
  ledger_api=self.ledger_api(chain=chain, rpc=rpc),
300
357
  crypto=self.crypto,
301
- safe=t.cast(str, self.safes[chain]),
358
+ safe=self.safes[chain],
302
359
  to=to,
303
360
  amount=amount,
304
361
  )
@@ -312,14 +369,15 @@ class EthereumMasterWallet(MasterWallet):
312
369
  rpc: t.Optional[str] = None,
313
370
  ) -> t.Optional[str]:
314
371
  """Transfer erc20 from safe wallet."""
315
- if self.safes is None:
316
- raise ValueError("Safes not initialized")
372
+ to = self._pre_transfer_checks(
373
+ to=to, amount=amount, chain=chain, from_safe=True, asset=token
374
+ )
317
375
 
318
376
  return transfer_erc20_from_safe(
319
377
  ledger_api=self.ledger_api(chain=chain, rpc=rpc),
320
378
  crypto=self.crypto,
321
379
  token=token,
322
- safe=t.cast(str, self.safes[chain]), # type: ignore
380
+ safe=self.safes[chain],
323
381
  to=to,
324
382
  amount=amount,
325
383
  )
@@ -333,6 +391,10 @@ class EthereumMasterWallet(MasterWallet):
333
391
  rpc: t.Optional[str] = None,
334
392
  ) -> t.Optional[str]:
335
393
  """Transfer erc20 from EOA wallet."""
394
+ to = self._pre_transfer_checks(
395
+ to=to, amount=amount, chain=chain, from_safe=False, asset=token
396
+ )
397
+
336
398
  wallet_address = self.address
337
399
  ledger_api = t.cast(EthereumApi, self.ledger_api(chain=chain, rpc=rpc))
338
400
  tx_settler = TxSettler(
@@ -375,136 +437,109 @@ class EthereumMasterWallet(MasterWallet):
375
437
  tx_hash = tx_receipt.get("transactionHash", "").hex()
376
438
  return tx_hash
377
439
 
378
- def transfer(
440
+ def transfer( # pylint: disable=too-many-arguments
379
441
  self,
380
442
  to: str,
381
443
  amount: int,
382
444
  chain: Chain,
445
+ asset: str = ZERO_ADDRESS,
383
446
  from_safe: bool = True,
384
447
  rpc: t.Optional[str] = None,
385
448
  ) -> t.Optional[str]:
386
449
  """Transfer funds to the given account."""
387
- if amount <= 0:
388
- return None
389
-
390
450
  if from_safe:
391
- sender = t.cast(str, self.safes[chain])
392
- sender_str = f"Safe {sender}"
393
- else:
394
- sender = self.crypto.address
395
- sender_str = f"EOA {sender}"
396
-
397
- ledger_api = self.ledger_api(chain=chain, rpc=rpc)
398
- to = ledger_api.api.to_checksum_address(to)
399
- balance = ledger_api.get_balance(address=sender)
400
-
401
- if balance < amount:
402
- raise InsufficientFundsException(
403
- f"Cannot transfer {amount} native units from {sender_str} to {to} on chain {chain.value.capitalize()}. "
404
- f"Balance of {sender_str} is {balance} native units on chain {chain.value.capitalize()}."
405
- )
451
+ if asset == ZERO_ADDRESS:
452
+ return self._transfer_from_safe(
453
+ to=to,
454
+ amount=amount,
455
+ chain=chain,
456
+ rpc=rpc,
457
+ )
406
458
 
407
- if from_safe:
408
- return self._transfer_from_safe(
459
+ return self._transfer_erc20_from_safe(
460
+ token=asset,
409
461
  to=to,
410
462
  amount=amount,
411
463
  chain=chain,
412
464
  rpc=rpc,
413
465
  )
414
- return self._transfer_from_eoa(
415
- to=to,
416
- amount=amount,
417
- chain=chain,
418
- rpc=rpc,
419
- )
420
-
421
- # pylint: disable=too-many-arguments
422
- def transfer_erc20(
423
- self,
424
- token: str,
425
- to: str,
426
- amount: int,
427
- chain: Chain,
428
- from_safe: bool = True,
429
- rpc: t.Optional[str] = None,
430
- ) -> t.Optional[str]:
431
- """Transfer funds to the given account."""
432
- if amount <= 0:
433
- return None
434
-
435
- if from_safe:
436
- sender = t.cast(str, self.safes[chain])
437
- sender_str = f"Safe {sender}"
438
- else:
439
- sender = self.crypto.address
440
- sender_str = f"EOA {sender}"
441
-
442
- ledger_api = self.ledger_api(chain=chain, rpc=rpc)
443
- to = ledger_api.api.to_checksum_address(to)
444
- balance = (
445
- registry_contracts.erc20.get_instance(
446
- ledger_api=ledger_api,
447
- contract_address=token,
448
- )
449
- .functions.balanceOf(sender)
450
- .call()
451
- )
452
-
453
- tokens = {OLAS[chain]: "OLAS", USDC[chain]: "USDC"}
454
- token_name = tokens.get(token, token)
455
-
456
- if balance < amount:
457
- raise InsufficientFundsException(
458
- f"Cannot transfer {amount} {token_name} from {sender_str} to {to} on chain {chain.value.capitalize()}. "
459
- f"Balance of {sender_str} is {balance} {token_name} on chain {chain.value.capitalize()}."
460
- )
461
466
 
462
- if from_safe:
463
- return self._transfer_erc20_from_safe(
464
- token=token,
467
+ if asset == ZERO_ADDRESS:
468
+ return self._transfer_from_eoa(
465
469
  to=to,
466
470
  amount=amount,
467
471
  chain=chain,
468
472
  rpc=rpc,
469
473
  )
474
+
470
475
  return self._transfer_erc20_from_eoa(
471
- token=token,
476
+ token=asset,
472
477
  to=to,
473
478
  amount=amount,
474
479
  chain=chain,
475
480
  rpc=rpc,
476
481
  )
477
482
 
478
- def transfer_asset(
483
+ def transfer_from_safe_then_eoa(
479
484
  self,
480
485
  to: str,
481
486
  amount: int,
482
487
  chain: Chain,
483
488
  asset: str = ZERO_ADDRESS,
484
- from_safe: bool = True,
485
489
  rpc: t.Optional[str] = None,
486
- ) -> t.Optional[str]:
490
+ ) -> t.List[str]:
487
491
  """
488
- Transfer assets to the given account.
492
+ Transfer assets to the given account using Safe balance first, and EOA balance for leftover.
489
493
 
490
494
  If asset is a zero address, transfer native currency.
491
495
  """
496
+ safe_balance = self.get_balance(chain=chain, asset=asset, from_safe=True)
497
+ eoa_balance = self.get_balance(chain=chain, asset=asset, from_safe=False)
498
+ balance = safe_balance + eoa_balance
492
499
  if asset == ZERO_ADDRESS:
493
- return self.transfer(
500
+ # to account for gas fees burned in previous txs
501
+ # in this case we will set the amount = eoa_balance below
502
+ balance += DUST[chain]
503
+
504
+ if balance < amount:
505
+ raise InsufficientFundsException(
506
+ f"Cannot transfer {format_asset_amount(chain, asset, amount)} to {to} on chain {chain.name}. "
507
+ f"Balance of Master Safe {self.safes[chain]}: {format_asset_amount(chain, asset, safe_balance)}. "
508
+ f"Balance of Master EOA {self.address}: {format_asset_amount(chain, asset, eoa_balance)}. "
509
+ f"Missing: {format_asset_amount(chain, asset, amount - balance)}."
510
+ )
511
+
512
+ tx_hashes = []
513
+ from_safe_amount = min(safe_balance, amount)
514
+ if from_safe_amount > 0:
515
+ tx_hash = self.transfer(
494
516
  to=to,
495
- amount=amount,
517
+ amount=from_safe_amount,
496
518
  chain=chain,
497
- from_safe=from_safe,
519
+ asset=asset,
520
+ from_safe=True,
498
521
  rpc=rpc,
499
522
  )
500
- return self.transfer_erc20(
501
- token=asset,
502
- to=to,
503
- amount=amount,
504
- chain=chain,
505
- from_safe=from_safe,
506
- rpc=rpc,
507
- )
523
+ if tx_hash:
524
+ tx_hashes.append(tx_hash)
525
+ amount -= from_safe_amount
526
+
527
+ if amount > 0:
528
+ eoa_balance = self.get_balance(chain=chain, asset=asset, from_safe=False)
529
+ if (
530
+ asset == ZERO_ADDRESS
531
+ and eoa_balance <= amount <= eoa_balance + DUST[chain]
532
+ ):
533
+ # to make the internal function drain the EOA
534
+ amount = eoa_balance
535
+
536
+ tx_hash = self.transfer(
537
+ to=to, amount=amount, chain=chain, asset=asset, from_safe=False, rpc=rpc
538
+ )
539
+ if tx_hash:
540
+ tx_hashes.append(tx_hash)
541
+
542
+ return tx_hashes
508
543
 
509
544
  def drain(
510
545
  self,
@@ -514,23 +549,13 @@ class EthereumMasterWallet(MasterWallet):
514
549
  rpc: t.Optional[str] = None,
515
550
  ) -> None:
516
551
  """Drain all erc20/native assets to the given account."""
517
-
518
- ledger_api = self.ledger_api(chain=chain, rpc=rpc)
519
- assets = {token[chain] for token in ERC20_TOKENS}
520
-
521
- if from_safe:
522
- assets.add(ZERO_ADDRESS)
523
-
552
+ assets = [token[chain] for token in ERC20_TOKENS.values()] + [ZERO_ADDRESS]
524
553
  for asset in assets:
525
- balance = get_asset_balance(
526
- ledger_api=ledger_api,
527
- asset_address=asset,
528
- address=self.safes[chain] if from_safe else self.crypto.address,
529
- )
554
+ balance = self.get_balance(chain=chain, asset=asset, from_safe=from_safe)
530
555
  if balance <= 0:
531
556
  continue
532
557
 
533
- self.transfer_asset(
558
+ self.transfer(
534
559
  to=withdrawal_address,
535
560
  amount=balance,
536
561
  chain=chain,
@@ -539,14 +564,6 @@ class EthereumMasterWallet(MasterWallet):
539
564
  rpc=rpc,
540
565
  )
541
566
 
542
- if not from_safe:
543
- drain_eoa(
544
- ledger_api=ledger_api,
545
- crypto=self.crypto,
546
- withdrawal_address=withdrawal_address,
547
- chain_id=chain.id,
548
- )
549
-
550
567
  @classmethod
551
568
  def new(
552
569
  cls, password: str, path: Path
@@ -555,12 +572,26 @@ class EthereumMasterWallet(MasterWallet):
555
572
  # Backport support on aea
556
573
 
557
574
  eoa_wallet_path = path / cls._key
575
+ eoa_mnemonic_path = path / cls._mnemonic
576
+
558
577
  if eoa_wallet_path.exists():
559
578
  raise FileExistsError(f"Wallet file already exists at {eoa_wallet_path}.")
560
579
 
580
+ if eoa_mnemonic_path.exists():
581
+ raise FileExistsError(
582
+ f"Mnemonic file already exists at {eoa_mnemonic_path}."
583
+ )
584
+
585
+ eoa_wallet_path.parent.mkdir(parents=True, exist_ok=True)
586
+
587
+ # Store private key (Ethereum V3 keystore JSON) and encrypted mnemonic
561
588
  account = Account()
562
589
  account.enable_unaudited_hdwallet_features()
563
590
  crypto, mnemonic = account.create_with_mnemonic()
591
+ encrypted_mnemonic = EncryptedData.new(
592
+ path=eoa_mnemonic_path, password=password, plaintext_bytes=mnemonic.encode()
593
+ )
594
+ encrypted_mnemonic.store()
564
595
  eoa_wallet_path.write_text(
565
596
  data=json.dumps(
566
597
  Account.encrypt(
@@ -578,6 +609,17 @@ class EthereumMasterWallet(MasterWallet):
578
609
  wallet.password = password
579
610
  return wallet, mnemonic.split()
580
611
 
612
+ def decrypt_mnemonic(self, password: str) -> t.Optional[t.List[str]]:
613
+ """Retrieve the mnemonic"""
614
+ eoa_mnemonic_path = self.path / self.ledger_type.mnemonic_file
615
+
616
+ if not eoa_mnemonic_path.exists():
617
+ return None
618
+
619
+ encrypted_mnemonic = EncryptedData.load(eoa_mnemonic_path)
620
+ mnemonic = encrypted_mnemonic.decrypt(password).decode("utf-8")
621
+ return mnemonic.split()
622
+
581
623
  def update_password(self, new_password: str) -> None:
582
624
  """Updates password."""
583
625
  create_backup(self.path / self._key)
@@ -638,8 +680,9 @@ class EthereumMasterWallet(MasterWallet):
638
680
  rpc: t.Optional[str] = None,
639
681
  ) -> t.Optional[str]:
640
682
  """Create safe."""
641
- if chain in self.safe_chains:
642
- return None
683
+ if chain in self.safes:
684
+ raise ValueError(f"Wallet already has a Safe on chain {chain}.")
685
+
643
686
  safe, self.safe_nonce, tx_hash = create_gnosis_safe(
644
687
  ledger_api=self.ledger_api(chain=chain, rpc=rpc),
645
688
  crypto=self.crypto,
@@ -661,9 +704,9 @@ class EthereumMasterWallet(MasterWallet):
661
704
  ) -> bool:
662
705
  """Adds a backup owner if not present, or updates it by the provided backup owner. Setting a None backup owner will remove the current one, if any."""
663
706
  ledger_api = self.ledger_api(chain=chain, rpc=rpc)
664
- if chain not in self.safes: # type: ignore
665
- raise ValueError(f"Safes not created for chain {chain}!")
666
- safe = t.cast(str, self.safes[chain]) # type: ignore
707
+ if chain not in self.safes:
708
+ raise ValueError(f"Wallet does not have a Safe on chain {chain}.")
709
+ safe = t.cast(str, self.safes[chain])
667
710
  owners = get_owners(ledger_api=ledger_api, safe=safe)
668
711
 
669
712
  if len(owners) > 2:
@@ -718,39 +761,35 @@ class EthereumMasterWallet(MasterWallet):
718
761
  def extended_json(self) -> t.Dict:
719
762
  """Get JSON representation with extended information (e.g., safe owners)."""
720
763
  rpc = None
721
- tokens = (OLAS, USDC)
722
764
  wallet_json = self.json
723
765
 
724
- if not self.safes:
725
- return wallet_json
726
-
766
+ balances: t.Dict[str, t.Dict[str, t.Dict[str, int]]] = {}
727
767
  owner_sets = set()
728
768
  for chain, safe in self.safes.items():
769
+ chain_str = chain.value
729
770
  ledger_api = self.ledger_api(chain=chain, rpc=rpc)
730
771
  owners = get_owners(ledger_api=ledger_api, safe=safe)
731
772
  owners.remove(self.address)
732
773
 
733
- balances: t.Dict[str, int] = {}
734
- balances[ZERO_ADDRESS] = ledger_api.get_balance(safe) or 0
735
- for token in tokens:
736
- balance = (
737
- registry_contracts.erc20.get_instance(
738
- ledger_api=ledger_api,
739
- contract_address=token[chain],
740
- )
741
- .functions.balanceOf(safe)
742
- .call()
743
- )
744
- balances[token[chain]] = balance
774
+ balances[chain_str] = {self.address: {}, safe: {}}
745
775
 
776
+ assets = [token[chain] for token in ERC20_TOKENS.values()] + [ZERO_ADDRESS]
777
+ for asset in assets:
778
+ balances[chain_str][self.address][asset] = self.get_balance(
779
+ chain=chain, asset=asset, from_safe=False
780
+ )
781
+ balances[chain_str][safe][asset] = self.get_balance(
782
+ chain=chain, asset=asset, from_safe=True
783
+ )
746
784
  wallet_json["safes"][chain.value] = {
747
785
  wallet_json["safes"][chain.value]: {
748
786
  "backup_owners": owners,
749
- "balances": balances,
787
+ "balances": balances[chain_str][safe],
750
788
  }
751
789
  }
752
790
  owner_sets.add(frozenset(owners))
753
791
 
792
+ wallet_json["balances"] = balances
754
793
  wallet_json["extended_json"] = True
755
794
  wallet_json["consistent_safe_address"] = len(set(self.safes.values())) == 1
756
795
  wallet_json["consistent_backup_owner"] = len(owner_sets) == 1
@@ -850,12 +889,10 @@ class MasterWalletManager:
850
889
  def __init__(
851
890
  self,
852
891
  path: Path,
853
- logger: logging.Logger,
854
892
  password: t.Optional[str] = None,
855
893
  ) -> None:
856
894
  """Initialize master wallet manager."""
857
895
  self.path = path
858
- self.logger = logger
859
896
  self._password = password
860
897
 
861
898
  @property
@@ -26,7 +26,7 @@ from logging import Logger
26
26
  from pathlib import Path
27
27
 
28
28
  from operate.account.user import UserAccount
29
- from operate.constants import USER_JSON, WALLETS_DIR
29
+ from operate.constants import MSG_INVALID_PASSWORD, USER_JSON, WALLETS_DIR
30
30
  from operate.utils.gnosis import get_owners
31
31
  from operate.wallet.master import MasterWalletManager
32
32
 
@@ -49,7 +49,7 @@ class WalletRecoveryManager:
49
49
  logger: Logger,
50
50
  wallet_manager: MasterWalletManager,
51
51
  ) -> None:
52
- """Initialize master wallet manager."""
52
+ """Initialize wallet recovery manager."""
53
53
  self.path = path
54
54
  self.logger = logger
55
55
  self.wallet_manager = wallet_manager
@@ -77,7 +77,7 @@ class WalletRecoveryManager:
77
77
 
78
78
  new_wallets_path = new_root / WALLETS_DIR
79
79
  new_wallet_manager = MasterWalletManager(
80
- path=new_wallets_path, logger=self.logger, password=new_password
80
+ path=new_wallets_path, password=new_password
81
81
  )
82
82
  new_wallet_manager.setup()
83
83
 
@@ -145,10 +145,10 @@ class WalletRecoveryManager:
145
145
 
146
146
  new_user_account = UserAccount.load(new_root / USER_JSON)
147
147
  if not new_user_account.is_valid(password=password):
148
- raise ValueError("Password is not valid.")
148
+ raise ValueError(MSG_INVALID_PASSWORD)
149
149
 
150
150
  new_wallet_manager = MasterWalletManager(
151
- path=new_wallets_path, logger=self.logger, password=password
151
+ path=new_wallets_path, password=password
152
152
  )
153
153
 
154
154
  ledger_types = {item.ledger_type for item in self.wallet_manager}