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

operate/wallet/master.py

@@ -20,16 +20,13 @@
 """Master key implementation"""
 
 import json
-import logging
 import os
 import typing as t
-from copy import deepcopy
 from dataclasses import dataclass, field
 from pathlib import Path
 
 from aea.crypto.base import Crypto, LedgerApi
-from aea.crypto.registries import make_ledger_api
-from aea_ledger_ethereum import DEFAULT_GAS_PRICE_STRATEGIES, EIP1559, GWEI, to_wei
+from aea.helpers.logging import setup_logger
 from aea_ledger_ethereum.ethereum import EthereumApi, EthereumCrypto
 from autonomy.chain.base import registry_contracts
 from autonomy.chain.config import ChainType as ChainProfile
@@ -42,15 +39,20 @@ from operate.constants import (
     ON_CHAIN_INTERACT_TIMEOUT,
     ZERO_ADDRESS,
 )
-from operate.ledger import get_default_rpc
-from operate.ledger.profiles import ERC20_TOKENS, OLAS, USDC
-from operate.operate_types import Chain, LedgerType
+from operate.ledger import (
+    get_default_ledger_api,
+    make_chain_ledger_api,
+    update_tx_with_gas_estimate,
+    update_tx_with_gas_pricing,
+)
+from operate.ledger.profiles import DUST, ERC20_TOKENS, format_asset_amount
+from operate.operate_types import Chain, EncryptedData, LedgerType
 from operate.resource import LocalResource
 from operate.utils import create_backup
 from operate.utils.gnosis import add_owner
 from operate.utils.gnosis import create_safe as create_gnosis_safe
 from operate.utils.gnosis import (
-    drain_eoa,
+    estimate_transfer_tx_fee,
     get_asset_balance,
     get_owners,
     remove_owner,
@@ -60,6 +62,9 @@ from operate.utils.gnosis import transfer as transfer_from_safe
 from operate.utils.gnosis import transfer_erc20_from_safe
 
 
+logger = setup_logger(name="master_wallet")
+
+
 # TODO Organize exceptions definition
 class InsufficientFundsException(Exception):
     """Insufficient funds exception."""
@@ -77,6 +82,7 @@ class MasterWallet(LocalResource):
     safe_nonce: t.Optional[int] = None
 
     _key: str
+    _mnemonic: str
     _crypto: t.Optional[Crypto] = None
     _password: t.Optional[str] = None
     _crypto_cls: t.Type[Crypto]
@@ -105,59 +111,42 @@ class MasterWallet(LocalResource):
         """Key path."""
         return self.path / self._key
 
+    @property
+    def mnemonic_path(self) -> Path:
+        """Mnemonic path."""
+        return self.path / self._mnemonic
+
+    @staticmethod
     def ledger_api(
-        self,
         chain: Chain,
         rpc: t.Optional[str] = None,
     ) -> LedgerApi:
         """Get ledger api object."""
-        gas_price_strategies = deepcopy(DEFAULT_GAS_PRICE_STRATEGIES)
-        if chain in (Chain.BASE, Chain.MODE, Chain.OPTIMISM):
-            gas_price_strategies[EIP1559]["fallback_estimate"]["maxFeePerGas"] = to_wei(
-                5, GWEI
-            )
-
-        return make_ledger_api(
-            self.ledger_type.name.lower(),
-            address=(rpc or get_default_rpc(chain=chain)),
-            chain_id=chain.id,
-            gas_price_strategies=gas_price_strategies,
-        )
-
-    def transfer(
-        self,
-        to: str,
-        amount: int,
-        chain: Chain,
-        from_safe: bool = True,
-        rpc: t.Optional[str] = None,
-    ) -> t.Optional[str]:
-        """Transfer funds to the given account."""
-        raise NotImplementedError()
+        if not rpc:
+            return get_default_ledger_api(chain=chain)
+        return make_chain_ledger_api(chain=chain, rpc=rpc)
 
-    # pylint: disable=too-many-arguments
-    def transfer_erc20(
+    def transfer(  # pylint: disable=too-many-arguments
         self,
-        token: str,
         to: str,
         amount: int,
         chain: Chain,
+        asset: str = ZERO_ADDRESS,
         from_safe: bool = True,
         rpc: t.Optional[str] = None,
     ) -> t.Optional[str]:
         """Transfer funds to the given account."""
         raise NotImplementedError()
 
-    def transfer_asset(
+    def transfer_from_safe_then_eoa(
         self,
         to: str,
         amount: int,
         chain: Chain,
         asset: str = ZERO_ADDRESS,
-        from_safe: bool = True,
         rpc: t.Optional[str] = None,
-    ) -> t.Optional[str]:
-        """Transfer erc20/native assets to the given account."""
+    ) -> t.List[str]:
+        """Transfer assets to the given account using Safe balance first, and EOA balance for leftover."""
         raise NotImplementedError()
 
     def drain(
@@ -175,6 +164,10 @@ class MasterWallet(LocalResource):
         """Create a new master wallet."""
         raise NotImplementedError()
 
+    def decrypt_mnemonic(self, password: str) -> t.Optional[t.List[str]]:
+        """Retrieve the mnemonic"""
+        raise NotImplementedError()
+
     def create_safe(
         self,
         chain: Chain,
@@ -213,6 +206,24 @@ class MasterWallet(LocalResource):
         """Updates password using the mnemonic."""
         raise NotImplementedError()
 
+    def get_balance(
+        self, chain: Chain, asset: str = ZERO_ADDRESS, from_safe: bool = True
+    ) -> int:
+        """Get wallet balance on a given chain."""
+        if from_safe:
+            if chain not in self.safes:
+                raise ValueError(f"Wallet does not have a Safe on chain {chain}.")
+
+            address = self.safes[chain]
+        else:
+            address = self.address
+
+        return get_asset_balance(
+            ledger_api=get_default_ledger_api(chain),
+            asset_address=asset,
+            address=address,
+        )
+
     # TODO move to resource.py if used in more resources similarly
     @property
     def extended_json(self) -> t.Dict:
@@ -239,12 +250,61 @@ class EthereumMasterWallet(MasterWallet):
 
     _file = ledger_type.config_file
     _key = ledger_type.key_file
+    _mnemonic = ledger_type.mnemonic_file
     _crypto_cls = EthereumCrypto
 
+    def _pre_transfer_checks(
+        self,
+        to: str,
+        amount: int,
+        chain: Chain,
+        from_safe: bool,
+        asset: str = ZERO_ADDRESS,
+    ) -> str:
+        """Checks conditions before transfer. Returns the to address checksummed."""
+        if amount <= 0:
+            raise ValueError(
+                "Transfer amount must be greater than zero, not transferring."
+            )
+
+        to = Web3().to_checksum_address(to)
+        if from_safe and chain not in self.safes:
+            raise ValueError(f"Wallet does not have a Safe on chain {chain}.")
+
+        balance = self.get_balance(chain=chain, asset=asset, from_safe=from_safe)
+        if balance < amount:
+            source = "Master Safe" if from_safe else " Master EOA"
+            source_address = self.safes[chain] if from_safe else self.address
+            raise InsufficientFundsException(
+                f"Cannot transfer {format_asset_amount(chain, asset, amount)} from {source} {source_address} to {to} on chain {chain.name}. "
+                f"Balance: {format_asset_amount(chain, asset, balance)}. Missing: {format_asset_amount(chain, asset, amount - balance)}."
+            )
+
+        return to
+
     def _transfer_from_eoa(
         self, to: str, amount: int, chain: Chain, rpc: t.Optional[str] = None
     ) -> t.Optional[str]:
         """Transfer funds from EOA wallet."""
+        balance = self.get_balance(chain=chain, from_safe=False)
+        tx_fee = estimate_transfer_tx_fee(
+            chain=chain, sender_address=self.address, to=to
+        )
+        if balance - tx_fee < amount <= balance:
+            # we assume that the user wants to drain the EOA
+            # we also account for dust here because withdraw call use some EOA balance to drain the safes first
+            amount = balance - tx_fee
+            if amount <= 0:
+                logger.warning(
+                    f"Not enough balance to cover gas fees for transfer of {amount} on chain {chain} from EOA {self.address}. "
+                    f"Balance is {balance}, estimated fee is {tx_fee}. Not transferring."
+                )
+                return None
+
+        to = self._pre_transfer_checks(
+            to=to, amount=amount, chain=chain, from_safe=False
+        )
+
         ledger_api = t.cast(EthereumApi, self.ledger_api(chain=chain, rpc=rpc))
         tx_helper = TxSettler(
             ledger_api=ledger_api,
@@ -288,13 +348,14 @@ class EthereumMasterWallet(MasterWallet):
         self, to: str, amount: int, chain: Chain, rpc: t.Optional[str] = None
     ) -> t.Optional[str]:
         """Transfer funds from safe wallet."""
-        if self.safes is None:
-            raise ValueError("Safes not initialized")
+        to = self._pre_transfer_checks(
+            to=to, amount=amount, chain=chain, from_safe=True
+        )
 
         return transfer_from_safe(
             ledger_api=self.ledger_api(chain=chain, rpc=rpc),
             crypto=self.crypto,
-            safe=t.cast(str, self.safes[chain]),
+            safe=self.safes[chain],
             to=to,
             amount=amount,
         )
@@ -308,14 +369,15 @@ class EthereumMasterWallet(MasterWallet):
         rpc: t.Optional[str] = None,
     ) -> t.Optional[str]:
         """Transfer erc20 from safe wallet."""
-        if self.safes is None:
-            raise ValueError("Safes not initialized")
+        to = self._pre_transfer_checks(
+            to=to, amount=amount, chain=chain, from_safe=True, asset=token
+        )
 
         return transfer_erc20_from_safe(
             ledger_api=self.ledger_api(chain=chain, rpc=rpc),
             crypto=self.crypto,
             token=token,
-            safe=t.cast(str, self.safes[chain]),  # type: ignore
+            safe=self.safes[chain],
             to=to,
             amount=amount,
         )
@@ -329,6 +391,10 @@ class EthereumMasterWallet(MasterWallet):
         rpc: t.Optional[str] = None,
     ) -> t.Optional[str]:
         """Transfer erc20 from EOA wallet."""
+        to = self._pre_transfer_checks(
+            to=to, amount=amount, chain=chain, from_safe=False, asset=token
+        )
+
         wallet_address = self.address
         ledger_api = t.cast(EthereumApi, self.ledger_api(chain=chain, rpc=rpc))
         tx_settler = TxSettler(
@@ -357,10 +423,9 @@ class EthereumMasterWallet(MasterWallet):
                     "nonce": ledger_api.api.eth.get_transaction_count(wallet_address),
                 }
             )
-            return ledger_api.update_with_gas_estimate(
-                transaction=tx,
-                raise_on_try=False,
-            )
+            update_tx_with_gas_pricing(tx, ledger_api)
+            update_tx_with_gas_estimate(tx, ledger_api)
+            return tx
 
         setattr(tx_settler, "build", _build_transfer_tx)  # noqa: B010
         tx_receipt = tx_settler.transact(
@@ -372,136 +437,109 @@ class EthereumMasterWallet(MasterWallet):
         tx_hash = tx_receipt.get("transactionHash", "").hex()
         return tx_hash
 
-    def transfer(
+    def transfer(  # pylint: disable=too-many-arguments
         self,
         to: str,
         amount: int,
         chain: Chain,
+        asset: str = ZERO_ADDRESS,
         from_safe: bool = True,
         rpc: t.Optional[str] = None,
     ) -> t.Optional[str]:
         """Transfer funds to the given account."""
-        if amount <= 0:
-            return None
-
         if from_safe:
-            sender = t.cast(str, self.safes[chain])
-            sender_str = f"Safe {sender}"
-        else:
-            sender = self.crypto.address
-            sender_str = f"EOA {sender}"
-
-        ledger_api = self.ledger_api(chain=chain, rpc=rpc)
-        to = ledger_api.api.to_checksum_address(to)
-        balance = ledger_api.get_balance(address=sender)
-
-        if balance < amount:
-            raise InsufficientFundsException(
-                f"Cannot transfer {amount} native units from {sender_str} to {to} on chain {chain.value.capitalize()}. "
-                f"Balance of {sender_str} is {balance} native units on chain {chain.value.capitalize()}."
-            )
+            if asset == ZERO_ADDRESS:
+                return self._transfer_from_safe(
+                    to=to,
+                    amount=amount,
+                    chain=chain,
+                    rpc=rpc,
+                )
 
-        if from_safe:
-            return self._transfer_from_safe(
+            return self._transfer_erc20_from_safe(
+                token=asset,
                 to=to,
                 amount=amount,
                 chain=chain,
                 rpc=rpc,
             )
-        return self._transfer_from_eoa(
-            to=to,
-            amount=amount,
-            chain=chain,
-            rpc=rpc,
-        )
-
-    # pylint: disable=too-many-arguments
-    def transfer_erc20(
-        self,
-        token: str,
-        to: str,
-        amount: int,
-        chain: Chain,
-        from_safe: bool = True,
-        rpc: t.Optional[str] = None,
-    ) -> t.Optional[str]:
-        """Transfer funds to the given account."""
-        if amount <= 0:
-            return None
-
-        if from_safe:
-            sender = t.cast(str, self.safes[chain])
-            sender_str = f"Safe {sender}"
-        else:
-            sender = self.crypto.address
-            sender_str = f"EOA {sender}"
-
-        ledger_api = self.ledger_api(chain=chain, rpc=rpc)
-        to = ledger_api.api.to_checksum_address(to)
-        balance = (
-            registry_contracts.erc20.get_instance(
-                ledger_api=ledger_api,
-                contract_address=token,
-            )
-            .functions.balanceOf(sender)
-            .call()
-        )
-
-        tokens = {OLAS[chain]: "OLAS", USDC[chain]: "USDC"}
-        token_name = tokens.get(token, token)
-
-        if balance < amount:
-            raise InsufficientFundsException(
-                f"Cannot transfer {amount} {token_name} from {sender_str} to {to} on chain {chain.value.capitalize()}. "
-                f"Balance of {sender_str} is {balance} {token_name} on chain {chain.value.capitalize()}."
-            )
 
-        if from_safe:
-            return self._transfer_erc20_from_safe(
-                token=token,
+        if asset == ZERO_ADDRESS:
+            return self._transfer_from_eoa(
                 to=to,
                 amount=amount,
                 chain=chain,
                 rpc=rpc,
             )
+
         return self._transfer_erc20_from_eoa(
-            token=token,
+            token=asset,
             to=to,
             amount=amount,
             chain=chain,
             rpc=rpc,
         )
 
-    def transfer_asset(
+    def transfer_from_safe_then_eoa(
         self,
         to: str,
         amount: int,
         chain: Chain,
         asset: str = ZERO_ADDRESS,
-        from_safe: bool = True,
         rpc: t.Optional[str] = None,
-    ) -> t.Optional[str]:
+    ) -> t.List[str]:
         """
-        Transfer assets to the given account.
+        Transfer assets to the given account using Safe balance first, and EOA balance for leftover.
 
         If asset is a zero address, transfer native currency.
         """
+        safe_balance = self.get_balance(chain=chain, asset=asset, from_safe=True)
+        eoa_balance = self.get_balance(chain=chain, asset=asset, from_safe=False)
+        balance = safe_balance + eoa_balance
         if asset == ZERO_ADDRESS:
-            return self.transfer(
+            # to account for gas fees burned in previous txs
+            # in this case we will set the amount = eoa_balance below
+            balance += DUST[chain]
+
+        if balance < amount:
+            raise InsufficientFundsException(
+                f"Cannot transfer {format_asset_amount(chain, asset, amount)} to {to} on chain {chain.name}. "
+                f"Balance of Master Safe {self.safes[chain]}: {format_asset_amount(chain, asset, safe_balance)}. "
+                f"Balance of Master EOA {self.address}: {format_asset_amount(chain, asset, eoa_balance)}. "
+                f"Missing: {format_asset_amount(chain, asset, amount - balance)}."
+            )
+
+        tx_hashes = []
+        from_safe_amount = min(safe_balance, amount)
+        if from_safe_amount > 0:
+            tx_hash = self.transfer(
                 to=to,
-                amount=amount,
+                amount=from_safe_amount,
                 chain=chain,
-                from_safe=from_safe,
+                asset=asset,
+                from_safe=True,
                 rpc=rpc,
             )
-        return self.transfer_erc20(
-            token=asset,
-            to=to,
-            amount=amount,
-            chain=chain,
-            from_safe=from_safe,
-            rpc=rpc,
-        )
+            if tx_hash:
+                tx_hashes.append(tx_hash)
+        amount -= from_safe_amount
+
+        if amount > 0:
+            eoa_balance = self.get_balance(chain=chain, asset=asset, from_safe=False)
+            if (
+                asset == ZERO_ADDRESS
+                and eoa_balance <= amount <= eoa_balance + DUST[chain]
+            ):
+                # to make the internal function drain the EOA
+                amount = eoa_balance
+
+            tx_hash = self.transfer(
+                to=to, amount=amount, chain=chain, asset=asset, from_safe=False, rpc=rpc
+            )
+            if tx_hash:
+                tx_hashes.append(tx_hash)
+
+        return tx_hashes
 
     def drain(
         self,
@@ -511,23 +549,13 @@ class EthereumMasterWallet(MasterWallet):
         rpc: t.Optional[str] = None,
     ) -> None:
         """Drain all erc20/native assets to the given account."""
-
-        ledger_api = self.ledger_api(chain=chain, rpc=rpc)
-        assets = {token[chain] for token in ERC20_TOKENS}
-
-        if from_safe:
-            assets.add(ZERO_ADDRESS)
-
+        assets = [token[chain] for token in ERC20_TOKENS.values()] + [ZERO_ADDRESS]
         for asset in assets:
-            balance = get_asset_balance(
-                ledger_api=ledger_api,
-                asset_address=asset,
-                address=self.safes[chain] if from_safe else self.crypto.address,
-            )
+            balance = self.get_balance(chain=chain, asset=asset, from_safe=from_safe)
             if balance <= 0:
                 continue
 
-            self.transfer_asset(
+            self.transfer(
                 to=withdrawal_address,
                 amount=balance,
                 chain=chain,
@@ -536,14 +564,6 @@ class EthereumMasterWallet(MasterWallet):
                 rpc=rpc,
             )
 
-        if not from_safe:
-            drain_eoa(
-                ledger_api=ledger_api,
-                crypto=self.crypto,
-                withdrawal_address=withdrawal_address,
-                chain_id=chain.id,
-            )
-
     @classmethod
     def new(
         cls, password: str, path: Path
@@ -552,12 +572,26 @@ class EthereumMasterWallet(MasterWallet):
         # Backport support on aea
 
         eoa_wallet_path = path / cls._key
+        eoa_mnemonic_path = path / cls._mnemonic
+
         if eoa_wallet_path.exists():
             raise FileExistsError(f"Wallet file already exists at {eoa_wallet_path}.")
 
+        if eoa_mnemonic_path.exists():
+            raise FileExistsError(
+                f"Mnemonic file already exists at {eoa_mnemonic_path}."
+            )
+
+        eoa_wallet_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Store private key (Ethereum V3 keystore JSON) and encrypted mnemonic
         account = Account()
         account.enable_unaudited_hdwallet_features()
         crypto, mnemonic = account.create_with_mnemonic()
+        encrypted_mnemonic = EncryptedData.new(
+            path=eoa_mnemonic_path, password=password, plaintext_bytes=mnemonic.encode()
+        )
+        encrypted_mnemonic.store()
         eoa_wallet_path.write_text(
             data=json.dumps(
                 Account.encrypt(
@@ -575,6 +609,17 @@ class EthereumMasterWallet(MasterWallet):
         wallet.password = password
         return wallet, mnemonic.split()
 
+    def decrypt_mnemonic(self, password: str) -> t.Optional[t.List[str]]:
+        """Retrieve the mnemonic"""
+        eoa_mnemonic_path = self.path / self.ledger_type.mnemonic_file
+
+        if not eoa_mnemonic_path.exists():
+            return None
+
+        encrypted_mnemonic = EncryptedData.load(eoa_mnemonic_path)
+        mnemonic = encrypted_mnemonic.decrypt(password).decode("utf-8")
+        return mnemonic.split()
+
     def update_password(self, new_password: str) -> None:
         """Updates password."""
         create_backup(self.path / self._key)
@@ -635,8 +680,9 @@ class EthereumMasterWallet(MasterWallet):
         rpc: t.Optional[str] = None,
     ) -> t.Optional[str]:
         """Create safe."""
-        if chain in self.safe_chains:
-            return None
+        if chain in self.safes:
+            raise ValueError(f"Wallet already has a Safe on chain {chain}.")
+
         safe, self.safe_nonce, tx_hash = create_gnosis_safe(
             ledger_api=self.ledger_api(chain=chain, rpc=rpc),
             crypto=self.crypto,
@@ -658,9 +704,9 @@ class EthereumMasterWallet(MasterWallet):
     ) -> bool:
         """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."""
         ledger_api = self.ledger_api(chain=chain, rpc=rpc)
-        if chain not in self.safes:  # type: ignore
-            raise ValueError(f"Safes not created for chain {chain}!")
-        safe = t.cast(str, self.safes[chain])  # type: ignore
+        if chain not in self.safes:
+            raise ValueError(f"Wallet does not have a Safe on chain {chain}.")
+        safe = t.cast(str, self.safes[chain])
         owners = get_owners(ledger_api=ledger_api, safe=safe)
 
         if len(owners) > 2:
@@ -715,39 +761,35 @@ class EthereumMasterWallet(MasterWallet):
     def extended_json(self) -> t.Dict:
         """Get JSON representation with extended information (e.g., safe owners)."""
         rpc = None
-        tokens = (OLAS, USDC)
         wallet_json = self.json
 
-        if not self.safes:
-            return wallet_json
-
+        balances: t.Dict[str, t.Dict[str, t.Dict[str, int]]] = {}
         owner_sets = set()
         for chain, safe in self.safes.items():
+            chain_str = chain.value
             ledger_api = self.ledger_api(chain=chain, rpc=rpc)
             owners = get_owners(ledger_api=ledger_api, safe=safe)
             owners.remove(self.address)
 
-            balances: t.Dict[str, int] = {}
-            balances[ZERO_ADDRESS] = ledger_api.get_balance(safe) or 0
-            for token in tokens:
-                balance = (
-                    registry_contracts.erc20.get_instance(
-                        ledger_api=ledger_api,
-                        contract_address=token[chain],
-                    )
-                    .functions.balanceOf(safe)
-                    .call()
-                )
-                balances[token[chain]] = balance
+            balances[chain_str] = {self.address: {}, safe: {}}
 
+            assets = [token[chain] for token in ERC20_TOKENS.values()] + [ZERO_ADDRESS]
+            for asset in assets:
+                balances[chain_str][self.address][asset] = self.get_balance(
+                    chain=chain, asset=asset, from_safe=False
+                )
+                balances[chain_str][safe][asset] = self.get_balance(
+                    chain=chain, asset=asset, from_safe=True
+                )
             wallet_json["safes"][chain.value] = {
                 wallet_json["safes"][chain.value]: {
                     "backup_owners": owners,
-                    "balances": balances,
+                    "balances": balances[chain_str][safe],
                 }
             }
             owner_sets.add(frozenset(owners))
 
+        wallet_json["balances"] = balances
         wallet_json["extended_json"] = True
         wallet_json["consistent_safe_address"] = len(set(self.safes.values())) == 1
         wallet_json["consistent_backup_owner"] = len(owner_sets) == 1
@@ -847,12 +889,10 @@ class MasterWalletManager:
     def __init__(
         self,
         path: Path,
-        logger: logging.Logger,
         password: t.Optional[str] = None,
     ) -> None:
         """Initialize master wallet manager."""
         self.path = path
-        self.logger = logger
         self._password = password
 
     @property