olas-operate-middleware 0.10.6__py3-none-any.whl → 0.10.8__py3-none-any.whl

operate/services/manage.py

@@ -37,8 +37,16 @@ from aea_ledger_ethereum import LedgerApi
 from autonomy.chain.base import registry_contracts
 from autonomy.chain.config import CHAIN_PROFILES, ChainType
 from autonomy.chain.metadata import IPFS_URI_PREFIX
-
-from operate.constants import IPFS_ADDRESS, ZERO_ADDRESS
+from web3 import Web3
+
+from operate.constants import (
+    AGENT_LOG_DIR,
+    AGENT_LOG_ENV_VAR,
+    AGENT_PERSISTENT_STORAGE_DIR,
+    AGENT_PERSISTENT_STORAGE_ENV_VAR,
+    IPFS_ADDRESS,
+    ZERO_ADDRESS,
+)
 from operate.data import DATA_DIR
 from operate.data.contracts.mech_activity.contract import MechActivityContract
 from operate.data.contracts.requester_activity_checker.contract import (
@@ -624,6 +632,8 @@ class ServiceManager:
         # TODO fix this
         os.environ["CUSTOM_CHAIN_RPC"] = ledger_config.rpc
 
+        self._enable_recovery_module(service_config_id=service_config_id, chain=chain)
+
         current_agent_id = None
         on_chain_state = OnChainState.NON_EXISTENT
         if chain_data.token > -1:
@@ -726,8 +736,8 @@ class ServiceManager:
 
         # Set environment variables for the service
         for dir_name, env_var_name in (
-            ("persistent_data", "STORE_PATH"),
-            ("benchmarks", "LOG_DIR"),
+            (AGENT_PERSISTENT_STORAGE_DIR, AGENT_PERSISTENT_STORAGE_ENV_VAR),
+            (AGENT_LOG_DIR, AGENT_LOG_ENV_VAR),
         ):
             dir_path = service.path / dir_name
             dir_path.mkdir(parents=True, exist_ok=True)
@@ -829,6 +839,11 @@ class ServiceManager:
                 self._get_on_chain_state(service=service, chain=chain)
                 == OnChainState.PRE_REGISTRATION
             ):
+                self.logger.info("Execute recovery module operations")
+                self._execute_recovery_module_flow_from_safe(
+                    service_config_id=service_config_id, chain=chain
+                )
+
                 self.logger.info("Updating service")
                 receipt = (
                     sftxb.new_tx()
@@ -1047,6 +1062,7 @@ class ServiceManager:
                 )
             ).settle()
 
+        # Deploy service
         if (
             self._get_on_chain_state(service=service, chain=chain)
             == OnChainState.FINISHED_REGISTRATION
@@ -1055,15 +1071,27 @@ class ServiceManager:
 
             reuse_multisig = True
             info = sftxb.info(token_id=chain_data.token)
-            if info["multisig"] == ZERO_ADDRESS:
+            service_safe_address = info["multisig"]
+            if service_safe_address == ZERO_ADDRESS:
                 reuse_multisig = False
 
             self.logger.info(f"{reuse_multisig=}")
 
+            is_recovery_module_enabled = (
+                registry_contracts.gnosis_safe.is_module_enabled(
+                    ledger_api=sftxb.ledger_api,
+                    contract_address=service_safe_address,
+                    module_address=CONTRACTS[Chain(chain)]["recovery_module"],
+                ).get("enabled")
+            )
+
+            self.logger.info(f"{is_recovery_module_enabled=}")
+
             messages = sftxb.get_deploy_data_from_safe(
                 service_id=chain_data.token,
                 reuse_multisig=reuse_multisig,
                 master_safe=safe,
+                use_recovery_module=is_recovery_module_enabled,
             )
             tx = sftxb.new_tx()
             for message in messages:
@@ -1198,6 +1226,9 @@ class ServiceManager:
         wallet = self.wallet_manager.load(ledger_config.chain.ledger_type)
         safe = wallet.safes[Chain(chain)]  # type: ignore
 
+        if withdrawal_address:
+            withdrawal_address = Web3.to_checksum_address(withdrawal_address)
+
         # TODO fixme
         os.environ["CUSTOM_CHAIN_RPC"] = ledger_config.rpc
 
@@ -1289,19 +1320,21 @@ class ServiceManager:
                     },
                 )
 
+            self._enable_recovery_module(
+                service_config_id=service_config_id, chain=chain
+            )
             self.logger.info("Swapping Safe owners")
-            sftxb.swap(  # noqa: E800
-                service_id=chain_data.token,  # noqa: E800
+            owner_crypto = self.keys_manager.get_crypto_instance(
+                address=current_safe_owners[0]
+            )
+            sftxb.swap(
+                service_id=chain_data.token,
                 multisig=chain_data.multisig,  # TODO this can be read from the registry
-                owner_key=str(
-                    self.keys_manager.get(
-                        key=current_safe_owners[0]
-                    ).private_key  # TODO allow multiple owners
-                ),  # noqa: E800
+                owner_cryptos=[owner_crypto],  # TODO allow multiple owners
                 new_owner_address=(
                     safe if safe else wallet.crypto.address
                 ),  # TODO it should always be safe address
-            )  # noqa: E800
+            )
 
         if withdrawal_address is not None:
             ethereum_crypto = KeysManager().get_crypto_instance(
@@ -1318,6 +1351,148 @@ class ServiceManager:
             )
             self.logger.info(f"{service.name} signer drained")
 
+    def _execute_recovery_module_flow_from_safe(  # pylint: disable=too-many-locals
+        self,
+        service_config_id: str,
+        chain: str,
+    ) -> None:
+        """Execute recovery module operations from Safe"""
+        self.logger.info(f"_execute_recovery_module_operations_from_safe {chain=}")
+        service = self.load(service_config_id=service_config_id)
+        chain_config = service.chain_configs[chain]
+        chain_data = chain_config.chain_data
+        ledger_config = chain_config.ledger_config
+        wallet = self.wallet_manager.load(ledger_config.chain.ledger_type)
+        sftxb = self.get_eth_safe_tx_builder(ledger_config=ledger_config)
+        safe = wallet.safes[Chain(chain)]
+
+        if chain_data.token == NON_EXISTENT_TOKEN:
+            self.logger.info("Service is not minted.")
+            return
+
+        info = sftxb.info(token_id=chain_data.token)
+        service_safe_address = info["multisig"]
+        on_chain_state = OnChainState(info["service_state"])
+
+        if service_safe_address == ZERO_ADDRESS:
+            self.logger.info("Service Safe is not deployed.")
+            return
+
+        recovery_module_address = CONTRACTS[Chain(chain)]["recovery_module"]
+        is_recovery_module_enabled = registry_contracts.gnosis_safe.is_module_enabled(
+            ledger_api=sftxb.ledger_api,
+            contract_address=service_safe_address,
+            module_address=recovery_module_address,
+        ).get("enabled")
+
+        service_safe_owners = sftxb.get_service_safe_owners(service_id=chain_data.token)
+        master_safe_is_service_safe_owner = service_safe_owners == [safe]
+
+        self.logger.info(f"{is_recovery_module_enabled=}")
+        self.logger.info(f"{master_safe_is_service_safe_owner=}")
+
+        if not is_recovery_module_enabled and not master_safe_is_service_safe_owner:
+            self.logger.info(
+                "Recovery module is not enabled and Master Safe is not service Safe owner. Skipping recovery operations."
+            )
+            return
+
+        if not is_recovery_module_enabled:
+            self._enable_recovery_module(
+                service_config_id=service_config_id, chain=chain
+            )
+
+        if (
+            not master_safe_is_service_safe_owner
+            and on_chain_state == OnChainState.PRE_REGISTRATION
+        ):
+            self.logger.info("Recovering service Safe access through recovery module.")
+            sftxb.new_tx().add(
+                sftxb.get_recover_access_data(
+                    service_id=chain_data.token,
+                )
+            ).settle()
+            self.logger.info("Recovering service Safe done.")
+
+    def _enable_recovery_module(  # pylint: disable=too-many-locals
+        self,
+        service_config_id: str,
+        chain: str,
+    ) -> None:
+        """Enable recovery module"""
+        self.logger.info(f"_enable_recovery_module {chain=}")
+        service = self.load(service_config_id=service_config_id)
+        chain_config = service.chain_configs[chain]
+        chain_data = chain_config.chain_data
+        ledger_config = chain_config.ledger_config
+        wallet = self.wallet_manager.load(ledger_config.chain.ledger_type)
+        sftxb = self.get_eth_safe_tx_builder(ledger_config=ledger_config)
+        safe = wallet.safes[Chain(chain)]
+
+        if chain_data.token == NON_EXISTENT_TOKEN:
+            self.logger.info("Service is not minted.")
+            return
+
+        info = sftxb.info(token_id=chain_data.token)
+        service_safe_address = info["multisig"]
+
+        if service_safe_address == ZERO_ADDRESS:
+            self.logger.info("Service Safe is not deployed.")
+            return
+
+        recovery_module_address = CONTRACTS[Chain(chain)]["recovery_module"]
+        is_recovery_module_enabled = registry_contracts.gnosis_safe.is_module_enabled(
+            ledger_api=sftxb.ledger_api,
+            contract_address=service_safe_address,
+            module_address=recovery_module_address,
+        ).get("enabled")
+
+        if is_recovery_module_enabled:
+            self.logger.info("Recovery module is already enabled in service Safe.")
+            return
+
+        self.logger.info("Recovery module is not enabled.")
+
+        # NOTE Recovery from agent only works for single-agent services
+        agent_address = service.agent_addresses[0]
+        service_safe_owners = sftxb.get_service_safe_owners(service_id=chain_data.token)
+        agent_is_service_safe_owner = service_safe_owners == [agent_address]
+        master_safe_is_service_safe_owner = service_safe_owners == [safe]
+
+        if agent_is_service_safe_owner:
+            self.logger.info("(Agent) Enabling recovery module in service Safe.")
+            try:
+                crypto = self.keys_manager.get_crypto_instance(address=agent_address)
+                EthSafeTxBuilder._new_tx(  # pylint: disable=protected-access
+                    ledger_api=sftxb.ledger_api,
+                    crypto=crypto,
+                    chain_type=ChainType(chain),
+                    safe=service_safe_address,
+                ).add(
+                    sftxb.get_enable_module_data(
+                        module_address=recovery_module_address,
+                        safe_address=service_safe_address,
+                    )
+                ).settle()
+                self.logger.info(
+                    "(Agent) Recovery module enabled successfully in service Safe."
+                )
+            except Exception as e:  # pylint: disable=broad-except
+                self.logger.error(
+                    f"Failed to enable recovery module in service Safe. Exception {e}: {traceback.format_exc()}"
+                )
+        elif master_safe_is_service_safe_owner:
+            # TODO Enable recovery module when Safe owner = master Safe.
+            # This should be similar to the above code, but
+            # requires implement a transaction where the owner is another Safe.
+            self.logger.info(
+                "(Service owner) Enabling recovery module in service Safe. [Not implemented]"
+            )
+        else:
+            self.logger.error(
+                f"Cannot enable recovery module. Safe {service_safe_address} has inconsistent owners."
+            )
+
     def _get_current_staking_program(
         self, service: Service, chain: str
     ) -> t.Optional[str]:
@@ -2009,6 +2184,7 @@ class ServiceManager:
         wallet = self.wallet_manager.load(ledger_config.chain.ledger_type)
         ledger_api = wallet.ledger_api(chain=ledger_config.chain, rpc=ledger_config.rpc)
         ethereum_crypto = KeysManager().get_crypto_instance(service.agent_addresses[0])
+        withdrawal_address = Web3.to_checksum_address(withdrawal_address)
 
         # drain ERC20 tokens from service safe
         for token_name, token_address in (
@@ -2162,12 +2338,6 @@ class ServiceManager:
             deployment.delete()
         return deployment
 
-    def log_directories(self) -> None:
-        """Log directories."""
-        directories = [f"  - {str(p)}" for p in self.path.iterdir() if p.is_dir()]
-        directories_str = "\n".join(directories)
-        self.logger.info(f"Directories in {self.path}\n: {directories_str}")
-
     def update(
         self,
         service_config_id: str,

operate/services/protocol.py

@@ -34,19 +34,20 @@ from typing import Optional, Union, cast
 from aea.configurations.data_types import PackageType
 from aea.crypto.base import Crypto, LedgerApi
 from aea.helpers.base import IPFSHash, cd
-from aea_ledger_ethereum.ethereum import EthereumCrypto
 from autonomy.chain.base import registry_contracts
 from autonomy.chain.config import ChainConfigs, ChainType, ContractConfigs
 from autonomy.chain.constants import (
     GNOSIS_SAFE_PROXY_FACTORY_CONTRACT,
     GNOSIS_SAFE_SAME_ADDRESS_MULTISIG_CONTRACT,
     MULTISEND_CONTRACT,
+    RECOVERY_MODULE_CONTRACT,
+    SAFE_MULTISIG_WITH_RECOVERY_MODULE_CONTRACT,
 )
 from autonomy.chain.metadata import publish_metadata
 from autonomy.chain.service import (
     get_agent_instances,
     get_deployment_payload,
-    get_reuse_multisig_payload,
+    get_deployment_with_recovery_payload,
     get_service_info,
     get_token_deposit_amount,
 )
@@ -65,6 +66,7 @@ from operate.constants import (
 )
 from operate.data import DATA_DIR
 from operate.data.contracts.dual_staking_token.contract import DualStakingTokenContract
+from operate.data.contracts.recovery_module.contract import RecoveryModule
 from operate.data.contracts.staking_token.contract import StakingTokenContract
 from operate.operate_types import Chain as OperateChain
 from operate.operate_types import ContractAddresses
@@ -719,7 +721,11 @@ class _ChainUtil:
         ).get("owners", [])
 
     def swap(  # pylint: disable=too-many-arguments,too-many-locals
-        self, service_id: int, multisig: str, owner_key: str, new_owner_address: str
+        self,
+        service_id: int,
+        multisig: str,
+        owner_cryptos: t.List[Crypto],
+        new_owner_address: str,
     ) -> None:
         """Swap safe owner."""
         logging.info(f"Swapping safe for service {service_id} [{multisig}]...")
@@ -733,11 +739,6 @@ class _ChainUtil:
             retries=ON_CHAIN_INTERACT_RETRIES,
             sleep=ON_CHAIN_INTERACT_SLEEP,
         )
-        with tempfile.TemporaryDirectory() as temp_dir:
-            key_file = Path(temp_dir, "key.txt")
-            key_file.write_text(owner_key, encoding="utf-8")
-            owner_crypto = EthereumCrypto(private_key_path=str(key_file))
-        owner_cryptos: t.List[EthereumCrypto] = [owner_crypto]
         owners = [
             manager.ledger_api.api.to_checksum_address(owner_crypto.address)
             for owner_crypto in owner_cryptos
@@ -794,7 +795,7 @@ class _ChainUtil:
         tx = registry_contracts.gnosis_safe.get_raw_safe_transaction(
             ledger_api=manager.ledger_api,
             contract_address=multisig,
-            sender_address=owner_crypto.address,
+            sender_address=owner_cryptos[0].address,
             owners=tuple(owners),  # type: ignore
             to_address=tx_params["to_address"],
             value=tx_params["ether_value"],
@@ -803,7 +804,7 @@ class _ChainUtil:
             signatures_by_owner=owner_to_signature,
             operation=SafeOperation.DELEGATE_CALL.value,
         )
-        stx = owner_crypto.sign_transaction(tx)
+        stx = owner_cryptos[0].sign_transaction(tx)
         tx_digest = manager.ledger_api.send_signed_transaction(stx)
         receipt = manager.ledger_api.api.eth.wait_for_transaction_receipt(tx_digest)
         if receipt["status"] != 1:
@@ -1151,9 +1152,21 @@ class OnChainManager(_ChainUtil):
 class EthSafeTxBuilder(_ChainUtil):
     """Safe Transaction builder."""
 
-    def new_tx(self) -> GnosisSafeTransaction:
+    @classmethod
+    def _new_tx(
+        cls, ledger_api: LedgerApi, crypto: Crypto, chain_type: ChainType, safe: str
+    ) -> GnosisSafeTransaction:
         """Create a new GnosisSafeTransaction instance."""
         return GnosisSafeTransaction(
+            ledger_api=ledger_api,
+            crypto=crypto,
+            chain_type=chain_type,
+            safe=safe,
+        )
+
+    def new_tx(self) -> GnosisSafeTransaction:
+        """Create a new GnosisSafeTransaction instance."""
+        return EthSafeTxBuilder._new_tx(
             ledger_api=self.wallet.ledger_api(
                 chain=OperateChain.from_string(self.chain_type.value),
                 rpc=self.rpc,
@@ -1307,55 +1320,12 @@ class EthSafeTxBuilder(_ChainUtil):
             "value": cost_of_bond,
         }
 
-    def get_deploy_data(
-        self,
-        service_id: int,
-        reuse_multisig: bool = False,
-    ) -> t.Dict:
-        """Get deploy tx data."""
-        instance = registry_contracts.service_manager.get_instance(
-            ledger_api=self.ledger_api,
-            contract_address=self.contracts["service_manager"],
-        )
-        if reuse_multisig:
-            _deployment_payload, error = get_reuse_multisig_payload(
-                ledger_api=self.ledger_api,
-                crypto=self.crypto,
-                chain_type=self.chain_type,
-                service_id=service_id,
-            )
-            if _deployment_payload is None:
-                raise ValueError(error)
-            deployment_payload = _deployment_payload
-            gnosis_safe_multisig = ContractConfigs.get(
-                GNOSIS_SAFE_SAME_ADDRESS_MULTISIG_CONTRACT.name
-            ).contracts[self.chain_type]
-        else:
-            deployment_payload = get_deployment_payload()
-            gnosis_safe_multisig = ContractConfigs.get(
-                GNOSIS_SAFE_PROXY_FACTORY_CONTRACT.name
-            ).contracts[self.chain_type]
-
-        txd = instance.encodeABI(
-            fn_name="deploy",
-            args=[
-                service_id,
-                gnosis_safe_multisig,
-                deployment_payload,
-            ],
-        )
-        return {
-            "to": self.contracts["service_manager"],
-            "data": txd[2:],
-            "operation": MultiSendOperation.CALL,
-            "value": 0,
-        }
-
     def get_deploy_data_from_safe(
         self,
         service_id: int,
         master_safe: str,
         reuse_multisig: bool = False,
+        use_recovery_module: bool = True,
     ) -> t.List[t.Dict[str, t.Any]]:
         """Get the deploy data instructions for a safe"""
         registry_instance = registry_contracts.service_manager.get_instance(
@@ -1364,27 +1334,50 @@ class EthSafeTxBuilder(_ChainUtil):
         )
         approve_hash_message = None
         if reuse_multisig:
-            (
-                _deployment_payload,
-                approve_hash_message,
-                error,
-            ) = get_reuse_multisig_from_safe_payload(
-                ledger_api=self.ledger_api,
-                chain_type=self.chain_type,
-                service_id=service_id,
-                master_safe=master_safe,
-            )
-            if _deployment_payload is None:
-                raise ValueError(error)
-            deployment_payload = _deployment_payload
-            gnosis_safe_multisig = ContractConfigs.get(
-                GNOSIS_SAFE_SAME_ADDRESS_MULTISIG_CONTRACT.name
-            ).contracts[self.chain_type]
-        else:
-            deployment_payload = get_deployment_payload()
-            gnosis_safe_multisig = ContractConfigs.get(
-                GNOSIS_SAFE_PROXY_FACTORY_CONTRACT.name
-            ).contracts[self.chain_type]
+            if not use_recovery_module:
+                (
+                    _deployment_payload,
+                    approve_hash_message,
+                    error,
+                ) = get_reuse_multisig_from_safe_payload(
+                    ledger_api=self.ledger_api,
+                    chain_type=self.chain_type,
+                    service_id=service_id,
+                    master_safe=master_safe,
+                )
+                if _deployment_payload is None:
+                    raise ValueError(error)
+                deployment_payload = _deployment_payload
+                gnosis_safe_multisig = ContractConfigs.get(
+                    GNOSIS_SAFE_SAME_ADDRESS_MULTISIG_CONTRACT.name
+                ).contracts[self.chain_type]
+            else:
+                (
+                    _deployment_payload,
+                    error,
+                ) = get_reuse_multisig_with_recovery_from_safe_payload(
+                    ledger_api=self.ledger_api,
+                    chain_type=self.chain_type,
+                    service_id=service_id,
+                    master_safe=master_safe,
+                )
+                if _deployment_payload is None:
+                    raise ValueError(error)
+                deployment_payload = _deployment_payload
+                gnosis_safe_multisig = ContractConfigs.get(
+                    RECOVERY_MODULE_CONTRACT.name
+                ).contracts[self.chain_type]
+        else:  # Deploy a new multisig
+            if not use_recovery_module:
+                deployment_payload = get_deployment_payload()
+                gnosis_safe_multisig = ContractConfigs.get(
+                    GNOSIS_SAFE_PROXY_FACTORY_CONTRACT.name
+                ).contracts[self.chain_type]
+            else:
+                deployment_payload = get_deployment_with_recovery_payload()
+                gnosis_safe_multisig = ContractConfigs.get(
+                    SAFE_MULTISIG_WITH_RECOVERY_MODULE_CONTRACT.name
+                ).contracts[self.chain_type]
 
         deploy_data = registry_instance.encodeABI(
             fn_name="deploy",
@@ -1571,6 +1564,56 @@ class EthSafeTxBuilder(_ChainUtil):
         # TODO: Discuss implementation
         raise NotImplementedError()
 
+    def get_recover_access_data(self, service_id: int) -> t.Dict:
+        """Get recover access tx data."""
+        instance = t.cast(
+            RecoveryModule,
+            RecoveryModule.from_dir(
+                directory=str(DATA_DIR / "contracts" / "recovery_module"),
+            ),
+        ).get_instance(
+            ledger_api=self.ledger_api,
+            contract_address=self.contracts["recovery_module"],
+        )
+        # TODO Replace the line above by this one once the recovery_module is
+        # included in the release of OpenAutonomy.
+        # instance = registry_contracts.recovery_module.get_instance(  # noqa: E800
+        #     ledger_api=self.ledger_api,  # noqa: E800
+        #     contract_address=self.contracts["recovery_module"],  # noqa: E800
+        # )  # noqa: E800
+        txd = instance.encodeABI(
+            fn_name="recoverAccess",
+            args=[service_id],
+        )
+        return {
+            "to": self.contracts["recovery_module"],
+            "data": txd[2:],
+            "operation": MultiSendOperation.CALL,
+            "value": 0,
+        }
+
+    def get_enable_module_data(
+        self,
+        safe_address: str,
+        module_address: str,
+    ) -> t.Dict:
+        """Get enable module tx data"""
+        self._patch()
+        instance = registry_contracts.gnosis_safe.get_instance(
+            ledger_api=self.ledger_api,
+            contract_address=safe_address,
+        )
+        txd = instance.encodeABI(
+            fn_name="enableModule",
+            args=[module_address],
+        )
+        return {
+            "to": safe_address,
+            "data": txd[2:],
+            "operation": MultiSendOperation.CALL,
+            "value": 0,
+        }
+
 
 def get_packed_signature_for_approved_hash(owners: t.Tuple[str]) -> bytes:
     """Get the packed signatures."""
@@ -1730,3 +1773,37 @@ def get_reuse_multisig_from_safe_payload(  # pylint: disable=too-many-locals
     )
     payload = multisig_address + safe_exec_data[2:]
     return payload, approve_hash_message, None
+
+
+def get_reuse_multisig_with_recovery_from_safe_payload(  # pylint: disable=too-many-locals
+    ledger_api: LedgerApi,
+    chain_type: ChainType,
+    service_id: int,
+    master_safe: str,
+) -> t.Tuple[Optional[str], Optional[str]]:
+    """Reuse multisig."""
+    _, multisig_address, _, _, *_ = get_service_info(
+        ledger_api=ledger_api,
+        chain_type=chain_type,
+        token_id=service_id,
+    )
+    if multisig_address == ZERO_ADDRESS:
+        return None, "Cannot reuse multisig, No previous deployment exist!"
+
+    service_owner = master_safe
+
+    multisig_instance = registry_contracts.gnosis_safe.get_instance(
+        ledger_api=ledger_api,
+        contract_address=multisig_address,
+    )
+
+    # Verify if the service was terminated properly or not
+    old_owners = multisig_instance.functions.getOwners().call()
+    if len(old_owners) != 1 or service_owner not in old_owners:
+        return (
+            None,
+            "Service was not terminated properly, the service owner should be the only owner of the safe",
+        )
+
+    payload = "0x" + int(service_id).to_bytes(32, "big").hex()
+    return payload, None