olas-operate-middleware 0.1.0rc59__py3-none-any.whl → 0.13.2__py3-none-any.whl

operate/utils/gnosis.py

@@ -20,25 +20,33 @@
 """Safe helpers."""
 
 import binascii
+import itertools
 import secrets
+import time
 import typing as t
 from enum import Enum
 
 from aea.crypto.base import Crypto, LedgerApi
+from aea.helpers.logging import setup_logger
 from autonomy.chain.base import registry_contracts
 from autonomy.chain.config import ChainType as ChainProfile
+from autonomy.chain.exceptions import ChainInteractionError
 from autonomy.chain.tx import TxSettler
+from web3 import Web3
 
 from operate.constants import (
     ON_CHAIN_INTERACT_RETRIES,
     ON_CHAIN_INTERACT_SLEEP,
     ON_CHAIN_INTERACT_TIMEOUT,
+    ZERO_ADDRESS,
 )
+from operate.ledger import get_default_ledger_api
+from operate.operate_types import Chain
 
 
-NULL_ADDRESS: str = "0x" + "0" * 40
+logger = setup_logger(name="operate.utils.gnosis")
 MAX_UINT256 = 2**256 - 1
-ZERO_ETH = 0
+SENTINEL_OWNERS = "0x0000000000000000000000000000000000000001"
 
 
 class SafeOperation(Enum):
@@ -65,8 +73,8 @@ def hash_payload_to_hex(  # pylint: disable=too-many-arguments,too-many-locals
     operation: int = SafeOperation.CALL.value,
     base_gas: int = 0,
     safe_gas_price: int = 0,
-    gas_token: str = NULL_ADDRESS,
-    refund_receiver: str = NULL_ADDRESS,
+    gas_token: str = ZERO_ADDRESS,
+    refund_receiver: str = ZERO_ADDRESS,
     use_flashbots: bool = False,
     gas_limit: int = 0,
     raise_on_failed_simulation: bool = False,
@@ -156,9 +164,9 @@ def _get_nonce() -> int:
 def create_safe(
     ledger_api: LedgerApi,
     crypto: Crypto,
-    owner: t.Optional[str] = None,
+    backup_owner: t.Optional[str] = None,
     salt_nonce: t.Optional[int] = None,
-) -> t.Tuple[str, int]:
+) -> t.Tuple[str, int, str]:
     """Create gnosis safe."""
     salt_nonce = salt_nonce or _get_nonce()
 
@@ -168,7 +176,7 @@ def create_safe(
         tx = registry_contracts.gnosis_safe.get_deploy_transaction(
             ledger_api=ledger_api,
             deployer_address=crypto.address,
-            owners=[crypto.address] if owner is None else [crypto.address, owner],
+            owners=([crypto.address]),
             threshold=1,
             salt_nonce=salt_nonce,
         )
@@ -193,12 +201,37 @@ def create_safe(
         contract="",
         kwargs={},
     )
+    tx_hash = receipt.get("transactionHash", "").hex()
     instance = registry_contracts.gnosis_safe_proxy_factory.get_instance(
         ledger_api=ledger_api,
         contract_address="0xa6b71e26c5e0845f74c812102ca7114b6a896ab2",
     )
     (event,) = instance.events.ProxyCreation().process_receipt(receipt)
-    return event["args"]["proxy"], salt_nonce
+    safe_address = event["args"]["proxy"]
+
+    if backup_owner is not None:
+        retry_delays = [0, 60, 120, 180, 240]
+        for attempt in range(1, len(retry_delays) + 1):
+            try:
+                add_owner(
+                    ledger_api=ledger_api,
+                    crypto=crypto,
+                    safe=safe_address,
+                    owner=backup_owner,
+                )
+                break  # success
+            except Exception as e:  # pylint: disable=broad-except
+                if attempt == len(retry_delays):
+                    raise RuntimeError(
+                        f"Failed to add backup owner {backup_owner} after {len(retry_delays)} attempts: {e}"
+                    ) from e
+                next_delay = retry_delays[attempt]
+                logger.error(
+                    f"Retry add owner {attempt}/{len(retry_delays)} in {next_delay} seconds due to error: {e}"
+                )
+                time.sleep(next_delay)
+
+    return safe_address, salt_nonce, tx_hash
 
 
 def get_owners(ledger_api: LedgerApi, safe: str) -> t.List[str]:
@@ -214,49 +247,68 @@ def send_safe_txs(
     safe: str,
     ledger_api: LedgerApi,
     crypto: Crypto,
-) -> None:
+    to: t.Optional[str] = None,
+) -> t.Optional[str]:
     """Send internal safe transaction."""
     owner = ledger_api.api.to_checksum_address(
         crypto.address,
     )
-    safe_tx_hash = registry_contracts.gnosis_safe.get_raw_safe_transaction_hash(
-        ledger_api=ledger_api,
-        contract_address=safe,
-        value=0,
-        safe_tx_gas=0,
-        to_address=safe,
-        data=txd,
-        operation=SafeOperation.CALL.value,
-    ).get("tx_hash")
-    safe_tx_bytes = binascii.unhexlify(
-        safe_tx_hash[2:],
-    )
-    signatures = {
-        owner: crypto.sign_message(
-            message=safe_tx_bytes,
-            is_deprecated_mode=True,
-        )[2:]
-    }
-    transaction = registry_contracts.gnosis_safe.get_raw_safe_transaction(
+    to_address = to or safe
+
+    def _build_tx(  # pylint: disable=unused-argument
+        *args: t.Any, **kwargs: t.Any
+    ) -> t.Optional[str]:
+        safe_tx_hash = registry_contracts.gnosis_safe.get_raw_safe_transaction_hash(
+            ledger_api=ledger_api,
+            contract_address=safe,
+            value=0,
+            safe_tx_gas=0,
+            to_address=to_address,
+            data=txd,
+            operation=SafeOperation.CALL.value,
+        ).get("tx_hash")
+        safe_tx_bytes = binascii.unhexlify(
+            safe_tx_hash[2:],
+        )
+        signatures = {
+            owner: crypto.sign_message(
+                message=safe_tx_bytes,
+                is_deprecated_mode=True,
+            )[2:]
+        }
+        return registry_contracts.gnosis_safe.get_raw_safe_transaction(
+            ledger_api=ledger_api,
+            contract_address=safe,
+            sender_address=owner,
+            owners=(owner,),  # type: ignore
+            to_address=to_address,
+            value=0,
+            data=txd,
+            safe_tx_gas=0,
+            signatures_by_owner=signatures,
+            operation=SafeOperation.CALL.value,
+            nonce=ledger_api.api.eth.get_transaction_count(owner),
+        )
+
+    tx_settler = TxSettler(
         ledger_api=ledger_api,
-        contract_address=safe,
-        sender_address=owner,
-        owners=(owner,),  # type: ignore
-        to_address=safe,
-        value=0,
-        data=txd,
-        safe_tx_gas=0,
-        signatures_by_owner=signatures,
-        operation=SafeOperation.CALL.value,
-        nonce=ledger_api.api.eth.get_transaction_count(owner),
+        crypto=crypto,
+        chain_type=Chain.from_id(
+            ledger_api._chain_id  # pylint: disable=protected-access
+        ),
+        timeout=ON_CHAIN_INTERACT_TIMEOUT,
+        retries=ON_CHAIN_INTERACT_RETRIES,
+        sleep=ON_CHAIN_INTERACT_SLEEP,
     )
-    ledger_api.get_transaction_receipt(
-        ledger_api.send_signed_transaction(
-            crypto.sign_transaction(
-                transaction,
-            ),
-        )
+    setattr(tx_settler, "build", _build_tx)  # noqa: B010
+    tx_receipt = tx_settler.transact(
+        method=lambda: {},
+        contract="",
+        kwargs={},
+        dry_run=False,
     )
+    tx_hash = tx_receipt.get("transactionHash", "").hex()
+    return tx_hash
 
 
 def add_owner(
@@ -285,14 +337,81 @@ def add_owner(
     )
 
 
-def swap_owner(  # pylint: disable=unused-argument
+def get_prev_owner(ledger_api: LedgerApi, safe: str, owner: str) -> str:
+    """Retrieve the previous owner in the owners list of the Safe."""
+
+    owners = get_owners(ledger_api=ledger_api, safe=safe)
+
+    try:
+        index = owners.index(owner) - 1
+    except ValueError as e:
+        raise ValueError(
+            f"Owner {owner} not found in the owners' list of the Safe."
+        ) from e
+
+    if index < 0:
+        return SENTINEL_OWNERS
+    return owners[index]
+
+
+def swap_owner(
     ledger_api: LedgerApi,
     crypto: Crypto,
     safe: str,
     old_owner: str,
     new_owner: str,
 ) -> None:
-    """Swap owner on a safe."""
+    """Swap owner of a safe."""
+
+    prev_owner = get_prev_owner(ledger_api=ledger_api, safe=safe, owner=old_owner)
+    instance = registry_contracts.gnosis_safe.get_instance(
+        ledger_api=ledger_api,
+        contract_address=safe,
+    )
+    txd = instance.encodeABI(
+        fn_name="swapOwner",
+        args=[
+            prev_owner,
+            old_owner,
+            new_owner,
+        ],
+    )
+    send_safe_txs(
+        txd=bytes.fromhex(txd[2:]),
+        safe=safe,
+        ledger_api=ledger_api,
+        crypto=crypto,
+    )
+
+
+def remove_owner(
+    ledger_api: LedgerApi,
+    crypto: Crypto,
+    safe: str,
+    owner: str,
+    threshold: int,
+) -> None:
+    """Remove owner from a safe."""
+
+    prev_owner = get_prev_owner(ledger_api=ledger_api, safe=safe, owner=owner)
+    instance = registry_contracts.gnosis_safe.get_instance(
+        ledger_api=ledger_api,
+        contract_address=safe,
+    )
+    txd = instance.encodeABI(
+        fn_name="removeOwner",
+        args=[
+            prev_owner,
+            owner,
+            threshold,
+        ],
+    )
+    send_safe_txs(
+        txd=bytes.fromhex(txd[2:]),
+        safe=safe,
+        ledger_api=ledger_api,
+        crypto=crypto,
+    )
 
 
 def transfer(
@@ -301,47 +420,255 @@ def transfer(
     safe: str,
     to: str,
     amount: t.Union[float, int],
-) -> None:
+) -> t.Optional[str]:
     """Transfer assets from safe to given address."""
     amount = int(amount)
     owner = ledger_api.api.to_checksum_address(
         crypto.address,
     )
-    safe_tx_hash = registry_contracts.gnosis_safe.get_raw_safe_transaction_hash(
+
+    def _build_tx(  # pylint: disable=unused-argument
+        *args: t.Any, **kwargs: t.Any
+    ) -> t.Optional[str]:
+        safe_tx_hash = registry_contracts.gnosis_safe.get_raw_safe_transaction_hash(
+            ledger_api=ledger_api,
+            contract_address=safe,
+            value=amount,
+            safe_tx_gas=0,
+            to_address=to,
+            data=b"",
+            operation=SafeOperation.CALL.value,
+        ).get("tx_hash")
+        safe_tx_bytes = binascii.unhexlify(
+            safe_tx_hash[2:],
+        )
+        signatures = {
+            owner: crypto.sign_message(
+                message=safe_tx_bytes,
+                is_deprecated_mode=True,
+            )[2:]
+        }
+        return registry_contracts.gnosis_safe.get_raw_safe_transaction(
+            ledger_api=ledger_api,
+            contract_address=safe,
+            sender_address=owner,
+            owners=(owner,),  # type: ignore
+            to_address=to,
+            value=amount,
+            data=b"",
+            safe_tx_gas=0,
+            signatures_by_owner=signatures,
+            operation=SafeOperation.CALL.value,
+            nonce=ledger_api.api.eth.get_transaction_count(owner),
+        )
+
+    tx_settler = TxSettler(
         ledger_api=ledger_api,
-        contract_address=safe,
-        value=amount,
-        safe_tx_gas=0,
-        to_address=to,
-        data=b"",
-        operation=SafeOperation.CALL.value,
-    ).get("tx_hash")
-    safe_tx_bytes = binascii.unhexlify(
-        safe_tx_hash[2:],
+        crypto=crypto,
+        chain_type=Chain.from_id(
+            ledger_api._chain_id  # pylint: disable=protected-access
+        ),
+        timeout=ON_CHAIN_INTERACT_TIMEOUT,
+        retries=ON_CHAIN_INTERACT_RETRIES,
+        sleep=ON_CHAIN_INTERACT_SLEEP,
+    )
+    setattr(tx_settler, "build", _build_tx)  # noqa: B010
+    tx_receipt = tx_settler.transact(
+        method=lambda: {},
+        contract="",
+        kwargs={},
+        dry_run=False,
+    )
+    tx_hash = tx_receipt.get("transactionHash", "").hex()
+    return tx_hash
+
+
+def transfer_erc20_from_safe(
+    ledger_api: LedgerApi,
+    crypto: Crypto,
+    safe: str,
+    token: str,
+    to: str,
+    amount: t.Union[float, int],
+) -> t.Optional[str]:
+    """Transfer ERC20 assets from safe to given address."""
+    amount = int(amount)
+    instance = registry_contracts.erc20.get_instance(
+        ledger_api=ledger_api,
+        contract_address=token,
     )
-    signatures = {
-        owner: crypto.sign_message(
-            message=safe_tx_bytes,
-            is_deprecated_mode=True,
-        )[2:]
-    }
-    transaction = registry_contracts.gnosis_safe.get_raw_safe_transaction(
+    txd = instance.encodeABI(
+        fn_name="transfer",
+        args=[
+            to,
+            amount,
+        ],
+    )
+    return send_safe_txs(
+        txd=bytes.fromhex(txd[2:]),
+        safe=safe,
         ledger_api=ledger_api,
-        contract_address=safe,
-        sender_address=owner,
-        owners=(owner,),  # type: ignore
-        to_address=to,
-        value=amount,
-        data=b"",
-        safe_tx_gas=0,
-        signatures_by_owner=signatures,
-        operation=SafeOperation.CALL.value,
-        nonce=ledger_api.api.eth.get_transaction_count(owner),
+        crypto=crypto,
+        to=token,
     )
-    ledger_api.get_transaction_receipt(
-        ledger_api.send_signed_transaction(
-            crypto.sign_transaction(
-                transaction,
-            ),
-        )
+
+
+def estimate_transfer_tx_fee(chain: Chain, sender_address: str, to: str) -> int:
+    """Estimate transfer transaction fee."""
+    ledger_api = get_default_ledger_api(chain)
+    tx = ledger_api.get_transfer_transaction(
+        sender_address=sender_address,
+        destination_address=to,
+        amount=0,
+        tx_fee=0,
+        tx_nonce="0x",
+        chain_id=chain.id,
+        raise_on_try=True,
+    )
+    tx = ledger_api.update_with_gas_estimate(
+        transaction=tx,
+        raise_on_try=False,
+    )
+    chain_fee = tx["gas"] * tx["maxFeePerGas"]
+    if chain in (
+        Chain.ARBITRUM_ONE,
+        Chain.BASE,
+        Chain.OPTIMISM,
+        Chain.MODE,
+    ):
+        chain_fee += ledger_api.get_l1_data_fee(tx)
+    return chain_fee
+
+
+def drain_eoa(
+    ledger_api: LedgerApi,
+    crypto: Crypto,
+    withdrawal_address: str,
+    chain_id: int,
+) -> t.Optional[str]:
+    """Drain all the native tokens from the crypto wallet."""
+    tx_helper = TxSettler(
+        ledger_api=ledger_api,
+        crypto=crypto,
+        chain_type=ChainProfile.CUSTOM,
+        timeout=ON_CHAIN_INTERACT_TIMEOUT,
+        retries=ON_CHAIN_INTERACT_RETRIES,
+        sleep=ON_CHAIN_INTERACT_SLEEP,
     )
+
+    def _build_tx(  # pylint: disable=unused-argument
+        *args: t.Any, **kwargs: t.Any
+    ) -> t.Dict:
+        """Build transaction"""
+        chain_fee = estimate_transfer_tx_fee(
+            chain=Chain.from_id(chain_id),
+            sender_address=crypto.address,
+            to=withdrawal_address,
+        )
+
+        amount = ledger_api.get_balance(crypto.address) - chain_fee
+        if amount <= 0:
+            raise ChainInteractionError(
+                f"No balance to drain from wallet: {crypto.address}"
+            )
+
+        tx = ledger_api.get_transfer_transaction(
+            sender_address=crypto.address,
+            destination_address=withdrawal_address,
+            amount=amount,
+            tx_fee=0,
+            tx_nonce="0x",
+            chain_id=chain_id,
+            raise_on_try=True,
+        )
+        empty_tx = tx.copy()
+        empty_tx["value"] = 0
+        empty_tx = ledger_api.update_with_gas_estimate(
+            transaction=empty_tx,
+            raise_on_try=False,
+        )
+        tx["gas"] = empty_tx["gas"]
+
+        logger.info(
+            f"Draining {tx['value']} native units from wallet: {crypto.address}"
+        )
+
+        return tx
+
+    setattr(tx_helper, "build", _build_tx)  # noqa: B010
+    try:
+        tx_receipt = tx_helper.transact(
+            method=lambda: {},
+            contract="",
+            kwargs={},
+            dry_run=False,
+        )
+    except ChainInteractionError as e:
+        if "No balance to drain from wallet" in str(e):
+            logger.warning(f"Failed to drain wallet {crypto.address} with error: {e}.")
+            return None
+
+        raise e
+
+    tx_hash = tx_receipt.get("transactionHash", None)
+    if tx_hash is not None:
+        return tx_hash.hex()
+
+    return None
+
+
+def get_asset_balance(
+    ledger_api: LedgerApi,
+    asset_address: str,
+    address: str,
+    raise_on_invalid_address: bool = True,
+) -> int:
+    """
+    Get the balance of a native asset or ERC20 token.
+
+    If contract address is a zero address, return the native balance.
+    """
+    if not Web3.is_address(address):
+        if raise_on_invalid_address:
+            raise ValueError(f"Invalid address: {address}")
+        return 0
+
+    try:
+        if asset_address == ZERO_ADDRESS:
+            return ledger_api.get_balance(address, raise_on_try=True)
+        return (
+            registry_contracts.erc20.get_instance(
+                ledger_api=ledger_api,
+                contract_address=asset_address,
+            )
+            .functions.balanceOf(address)
+            .call()
+        )
+    except Exception as e:
+        raise RuntimeError(
+            f"Cannot get balance of {address=} {asset_address=} rpc={ledger_api._api.provider.endpoint_uri}."  # pylint: disable=protected-access
+        ) from e
+
+
+def get_assets_balances(
+    ledger_api: LedgerApi,
+    asset_addresses: t.Set[str],
+    addresses: t.Set[str],
+    raise_on_invalid_address: bool = True,
+) -> t.Dict[str, t.Dict[str, int]]:
+    """
+    Get the balances of a list of native assets or ERC20 tokens.
+
+    If asset address is a zero address, return the native balance.
+    """
+    output: t.Dict[str, t.Dict[str, int]] = {}
+
+    for asset, address in itertools.product(asset_addresses, addresses):
+        output.setdefault(address, {})[asset] = get_asset_balance(
+            ledger_api=ledger_api,
+            asset_address=asset,
+            address=address,
+            raise_on_invalid_address=raise_on_invalid_address,
+        )
+
+    return output