Integrate walterriddler contribution: Allow claiming olas staking rewards for password protected keys

README.md

@@ -114,12 +114,12 @@ To stop your agent, use:
 ./stop_service.sh
 ```
 
-### Claim accrued OLAS
+### Claim accrued OLAS staking rewards
 
-If your service is staked, you can claim accrued OLAS through the script
+If your service is staked, you can claim accrued OLAS staking rewards through the script
 
 ```bash
-./claim_olas.sh
+./claim_staking_rewards.sh
 ```
 
 The accrued OLAS will be transferred to your service Safe without having to unstake your service.

claim_olas.sh → claim_staking_rewards.sh

@@ -18,4 +18,4 @@
 #
 # ------------------------------------------------------------------------------
 
-cd trader; poetry run python "../scripts/claim.py"; cd ..
+cd trader; poetry run python "../scripts/claim_staking_rewards.py"; cd ..

scripts/claim.py → scripts/claim_staking_rewards.py

@@ -23,9 +23,12 @@
 import json
 import os
 import sys
+from getpass import getpass
 from pathlib import Path
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
+from aea.crypto.helpers import DecryptError, KeyIsIncorrect
+from aea_ledger_ethereum.ethereum import EthereumCrypto
 from dotenv import dotenv_values
 from web3 import Web3
 
@@ -36,7 +39,7 @@
 RPC_PATH = Path(STORE_PATH, "rpc.txt")
 SERVICE_ID_PATH = Path(STORE_PATH, "service_id.txt")
 SERVICE_SAFE_ADDRESS_PATH = Path(STORE_PATH, "service_safe_address.txt")
-OWNER_KEYS_JSON_PATH = Path(STORE_PATH, "operator_keys.json")
+OPERATOR_PKEY_PATH = Path(STORE_PATH, "operator_pkey.txt")
 DEFAULT_ENCODING = "utf-8"
 
 OLAS_TOKEN_ADDRESS_GNOSIS = "0xcE11e14225575945b8E6Dc0D4F2dD4C570f79d9f"
@@ -70,6 +73,15 @@
 )
 
 
+def _is_keystore(pkeypath: Path) -> bool:
+    try:
+        with open(pkeypath, "r", encoding="utf-8") as f:
+            json.load(f)
+        return True
+    except json.JSONDecodeError:
+        return False
+
+
 def _load_abi_from_file(path: Path) -> Dict[str, Any]:
     if not os.path.exists(path):
         print(
@@ -87,7 +99,7 @@ def _erc20_balance(
     address: str,
     token_address: str = OLAS_TOKEN_ADDRESS_GNOSIS,
     token_name: str = "OLAS",
-    decimal_precision: int = 2
+    decimal_precision: int = 2,
 ) -> str:
     """Get ERC20 balance"""
     rpc = RPC_PATH.read_text(encoding=DEFAULT_ENCODING).strip()
@@ -98,10 +110,14 @@ def _erc20_balance(
     return f"{balance / 10**18:.{decimal_precision}f} {token_name}"
 
 
-def _claim_rewards() -> None:
-    service_safe_address = SERVICE_SAFE_ADDRESS_PATH.read_text(encoding=DEFAULT_ENCODING).strip()
+def _claim_rewards(  # pylint: disable=too-many-locals
+    password: Optional[str] = None,
+) -> None:
+    service_safe_address = SERVICE_SAFE_ADDRESS_PATH.read_text(
+        encoding=DEFAULT_ENCODING
+    ).strip()
     print(
-        f"OLAS Balance on service Safe {service_safe_address}: {_erc20_balance(service_safe_address)}"
+        f"OLAS Balance of service Safe {service_safe_address}: {_erc20_balance(service_safe_address)}"
     )
 
     env_file_vars = dotenv_values(DOTENV_PATH)
@@ -113,50 +129,73 @@ def _claim_rewards() -> None:
     abi = _load_abi_from_file(STAKING_TOKEN_IMPLEMENTATION_ABI_PATH)
     staking_token_contract = w3.eth.contract(address=staking_token_address, abi=abi)
 
-    owner_private_key = json.loads(OWNER_KEYS_JSON_PATH.read_text(encoding=DEFAULT_ENCODING))[0][
-        "private_key"
-    ]
-    owner_address = Web3.to_checksum_address(
-        w3.eth.account.from_key(owner_private_key).address
-    )
+    try:
+        ethereum_crypto = EthereumCrypto(OPERATOR_PKEY_PATH, password=password)
+        operator_address = ethereum_crypto.address
+        operator_pkey = ethereum_crypto.private_key
+    except DecryptError:
+        print(
+            f"Could not decrypt key {OPERATOR_PKEY_PATH}. Please verify if your key file is password-protected, and if the provided password is correct (passwords are case-sensitive)."
+        )
+        sys.exit(1)
+    except KeyIsIncorrect:
+        print(f"Error decoding key file {OPERATOR_PKEY_PATH}.")
+        sys.exit(1)
 
     function = staking_token_contract.functions.claim(service_id)
     claim_transaction = function.build_transaction(
         {
             "chainId": GNOSIS_CHAIN_ID,
             "gas": DEFAULT_GAS,
             "gasPrice": w3.to_wei("3", "gwei"),
-            "nonce": w3.eth.get_transaction_count(owner_address),
+            "nonce": w3.eth.get_transaction_count(operator_address),
         }
     )
 
-    signed_tx = w3.eth.account.sign_transaction(claim_transaction, owner_private_key)
+    signed_tx = w3.eth.account.sign_transaction(claim_transaction, operator_pkey)
     tx_hash = w3.eth.send_raw_transaction(signed_tx.rawTransaction)
     tx_receipt = w3.eth.wait_for_transaction_receipt(tx_hash)
     print(f"Claim transaction done. Hash: {tx_hash.hex()}")
 
     if "status" in tx_receipt and tx_receipt["status"] == 0:
         print(
-            "The transaction was reverted. This may be caused because your service does not have rewards to claim."
+            "WARNING: The transaction was reverted. This may be caused because your service does not have rewards to claim."
         )
     else:
         print("")
         print(f"Claimed OLAS transferred to your service Safe {service_safe_address}")
-        print(
-            f"You can use your Owner/Operator wallet (address {owner_address}) to connect your Safe at"
-        )
-        print(f"{SAFE_WEBAPP_URL}{service_safe_address}")
+
+    print("")
+    print(
+        f"You can use your Owner/Operator wallet (address {operator_address}) to connect your Safe at {SAFE_WEBAPP_URL}{service_safe_address}."
+    )
 
 
 def main() -> None:
-    "Main method"
-    print("This script will claim accrued OLAS on the current staking contract to your service Safe.")
-    response = input("Do you want to continue (yes/no)? ").strip().lower()
+    """Main method."""
+    print("---------------------")
+    print("Claim staking rewards")
+    print("---------------------")
+    print("")
+    print(
+        "This script will claim the OLAS staking rewards accrued in the current staking contract and transfer them to your service Safe."
+    )
+    _continue = input("Do you want to continue (yes/no)? ").strip().lower()
+
+    if _continue not in ("y", "yes"):
+        sys.exit(0)
 
-    if response not in ("y", "yes"):
-        return
+    print("")
+
+    password = None
+    if _is_keystore(OPERATOR_PKEY_PATH):
+        print("Enter your password")
+        print("-------------------")
+        print("Your key files are protected with a password.")
+        password = getpass("Please, enter your password: ").strip()
+        print("")
 
-    _claim_rewards()
+    _claim_rewards(password)
 
 
 if __name__ == "__main__":