Merge pull request #2244 from opentensor/feat/gus/update-err-msg-pw

feat: return error message instead of raising exception

bittensor/commands/delegates.py

@@ -752,7 +752,13 @@ def _run(cli: "bittensor.cli", subtensor: "bittensor.subtensor"):
 
         # Unlock the wallet.
         wallet.hotkey
-        wallet.coldkey
+        try:
+            wallet.coldkey
+        except bittensor.KeyFileError:
+            bittensor.__console__.print(
+                ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+            )
+            return
 
         # Check if the hotkey is already a delegate.
         if subtensor.is_hotkey_delegate(wallet.hotkey.ss58_address):

bittensor/commands/identity.py

@@ -115,7 +115,14 @@ def _run(cli: "bittensor.cli", subtensor: "bittensor.subtensor"):
             console.print(":cross_mark: Aborted!")
             exit(0)
 
-        wallet.coldkey  # unlock coldkey
+        try:
+            wallet.coldkey  # unlock coldkey
+        except bittensor.KeyFileError:
+            bittensor.__console__.print(
+                ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+            )
+            return
+
         with console.status(":satellite: [bold green]Updating identity on-chain..."):
             try:
                 subtensor.update_identity(

bittensor/commands/senate.py

@@ -432,7 +432,13 @@ def _run(cli: "bittensor.cli", subtensor: "bittensor.subtensor"):
 
         # Unlock the wallet.
         wallet.hotkey
-        wallet.coldkey
+        try:
+            wallet.coldkey
+        except bittensor.KeyFileError:
+            bittensor.__console__.print(
+                ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+            )
+            return
 
         # Check if the hotkey is a delegate.
         if not subtensor.is_hotkey_delegate(wallet.hotkey.ss58_address):
@@ -514,7 +520,13 @@ def _run(cli: "bittensor.cli", subtensor: "bittensor.cli"):
 
         # Unlock the wallet.
         wallet.hotkey
-        wallet.coldkey
+        try:
+            wallet.coldkey
+        except bittensor.KeyFileError:
+            bittensor.__console__.print(
+                ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+            )
+            return
 
         if not subtensor.is_senate_member(hotkey_ss58=wallet.hotkey.ss58_address):
             console.print(
@@ -603,7 +615,15 @@ def _run(cli: "bittensor.cli", subtensor: "bittensor.subtensor"):
 
         # Unlock the wallet.
         wallet.hotkey
-        wallet.coldkey
+        try:
+            wallet.coldkey
+        except bittensor.KeyFileError:
+            bittensor.__console__.print(
+                ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+            )
+            return
+
+        vote_data = subtensor.get_vote_data(proposal_hash)
 
         vote_data = subtensor.get_vote_data(proposal_hash)
         if vote_data == None:

bittensor/extrinsics/delegation.py

@@ -47,9 +47,17 @@ def nominate_extrinsic(
         success (bool): ``True`` if the transaction was successful.
     """
     # Unlock the coldkey.
-    wallet.coldkey
-    wallet.hotkey
 
+    try:
+        wallet.coldkey
+
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
+
+    wallet.hotkey
     # Check if the hotkey is already a delegate.
     if subtensor.is_hotkey_delegate(wallet.hotkey.ss58_address):
         logger.error(
@@ -133,7 +141,13 @@ def delegate_extrinsic(
         NotDelegateError: If the hotkey is not a delegate on the chain.
     """
     # Decrypt keys,
-    wallet.coldkey
+    try:
+        wallet.coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
     if not subtensor.is_hotkey_delegate(delegate_ss58):
         raise NotDelegateError("Hotkey: {} is not a delegate.".format(delegate_ss58))
 
@@ -394,7 +408,14 @@ def decrease_take_extrinsic(
         success (bool): ``True`` if the transaction was successful.
     """
     # Unlock the coldkey.
-    wallet.coldkey
+    try:
+        wallet.coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
+
     wallet.hotkey
 
     with bittensor.__console__.status(
@@ -454,7 +475,14 @@ def increase_take_extrinsic(
         success (bool): ``True`` if the transaction was successful.
     """
     # Unlock the coldkey.
-    wallet.coldkey
+    try:
+        wallet.coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
+
     wallet.hotkey
 
     with bittensor.__console__.status(

bittensor/extrinsics/network.py

@@ -87,7 +87,13 @@ def register_subnetwork_extrinsic(
         ):
             return False
 
-    wallet.coldkey  # unlock coldkey
+    try:
+        wallet.coldkey  # unlock coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
 
     with bittensor.__console__.status(":satellite: Registering subnet..."):
         with subtensor.substrate as substrate:

bittensor/extrinsics/registration.py

@@ -259,7 +259,13 @@ def burned_register_extrinsic(
         )
         return False
 
-    wallet.coldkey  # unlock coldkey
+    try:
+        wallet.coldkey  # unlock coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
     with bittensor.__console__.status(
         f":satellite: Checking Account on [bold]subnet:{netuid}[/bold]..."
     ):
@@ -394,7 +400,13 @@ def run_faucet_extrinsic(
         return False, "Requires torch"
 
     # Unlock coldkey
-    wallet.coldkey
+    try:
+        wallet.coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False, ""
 
     # Get previous balance.
     old_balance = subtensor.get_balance(wallet.coldkeypub.ss58_address)
@@ -497,7 +509,13 @@ def swap_hotkey_extrinsic(
     wait_for_finalization: bool = True,
     prompt: bool = False,
 ) -> bool:
-    wallet.coldkey  # unlock coldkey
+    try:
+        wallet.coldkey  # unlock coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
     if prompt:
         # Prompt user for confirmation.
         if not Confirm.ask(

bittensor/extrinsics/root.py

@@ -54,7 +54,13 @@ def root_register_extrinsic(
             Flag is ``true`` if extrinsic was finalized or uncluded in the block. If we did not wait for finalization / inclusion, the response is ``true``.
     """
 
-    wallet.coldkey  # unlock coldkey
+    try:
+        wallet.coldkey  # unlock coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
 
     is_registered = subtensor.is_hotkey_registered(
         netuid=0, hotkey_ss58=wallet.hotkey.ss58_address
@@ -131,7 +137,13 @@ def set_root_weights_extrinsic(
             Flag is ``true`` if extrinsic was finalized or uncluded in the block. If we did not wait for finalization / inclusion, the response is ``true``.
     """
 
-    wallet.coldkey  # unlock coldkey
+    try:
+        wallet.coldkey  # unlock coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
 
     # First convert types.
     if isinstance(netuids, list):

bittensor/extrinsics/senate.py

@@ -46,7 +46,14 @@ def register_senate_extrinsic(
         success (bool):
             Flag is ``true`` if extrinsic was finalized or included in the block. If we did not wait for finalization / inclusion, the response is ``true``.
     """
-    wallet.coldkey  # unlock coldkey
+    try:
+        wallet.coldkey  # unlock coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
+
     wallet.hotkey  # unlock hotkey
 
     if prompt:
@@ -121,7 +128,14 @@ def leave_senate_extrinsic(
         success (bool):
             Flag is ``true`` if extrinsic was finalized or included in the block. If we did not wait for finalization / inclusion, the response is ``true``.
     """
-    wallet.coldkey  # unlock coldkey
+    try:
+        wallet.coldkey  # unlock coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
+
     wallet.hotkey  # unlock hotkey
 
     if prompt:

bittensor/extrinsics/staking.py

@@ -86,7 +86,13 @@ def add_stake_extrinsic(
             If the hotkey is not a delegate on the chain.
     """
     # Decrypt keys,
-    wallet.coldkey
+    try:
+        wallet.coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
 
     # Default to wallet's own hotkey if the value is not passed.
     if hotkey_ss58 is None:

bittensor/extrinsics/transfer.py

@@ -68,8 +68,15 @@ def transfer_extrinsic(
         # Convert bytes to hex string.
         dest = "0x" + dest.hex()
 
-    # Unlock wallet coldkey.
-    wallet.coldkey
+    try:
+        # Unlock wallet coldkey.
+        wallet.coldkey
+
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
 
     # Convert to bittensor.Balance
     if not isinstance(amount, bittensor.Balance):

bittensor/extrinsics/unstaking.py

@@ -58,7 +58,13 @@ def __do_remove_stake_single(
 
     """
     # Decrypt keys,
-    wallet.coldkey
+    try:
+        wallet.coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
 
     success = subtensor._do_unstake(
         wallet=wallet,
@@ -126,7 +132,13 @@ def unstake_extrinsic(
             Flag is ``true`` if extrinsic was finalized or uncluded in the block. If we did not wait for finalization / inclusion, the response is ``true``.
     """
     # Decrypt keys,
-    wallet.coldkey
+    try:
+        wallet.coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
 
     if hotkey_ss58 is None:
         hotkey_ss58 = wallet.hotkey.ss58_address  # Default to wallet's own hotkey.
@@ -304,7 +316,13 @@ def unstake_multiple_extrinsic(
             return True
 
     # Unlock coldkey.
-    wallet.coldkey
+    try:
+        wallet.coldkey
+    except bittensor.KeyFileError:
+        bittensor.__console__.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid[/red]:[bold white]\n  [/bold white]"
+        )
+        return False
 
     old_stakes = []
     own_hotkeys = []

bittensor/keyfile.py

@@ -33,6 +33,7 @@
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 from nacl import pwhash, secret
+from nacl.exceptions import CryptoError
 from password_strength import PasswordPolicy
 from substrateinterface.utils.ss58 import ss58_encode
 from termcolor import colored
@@ -321,7 +322,10 @@ def decrypt_keyfile_data(
                     memlimit=pwhash.argon2i.MEMLIMIT_SENSITIVE,
                 )
                 box = secret.SecretBox(key)
-                decrypted_keyfile_data = box.decrypt(keyfile_data[len("$NACL") :])
+                try:
+                    decrypted_keyfile_data = box.decrypt(keyfile_data[len("$NACL") :])
+                except CryptoError:
+                    raise bittensor.KeyFileError("Invalid password")
             # Ansible decrypt.
             elif keyfile_data_is_encrypted_ansible(keyfile_data):
                 vault = Vault(password)

tests/unit_tests/test_keyfile.py

@@ -624,3 +624,20 @@ def test_get_coldkey_password_from_environment(monkeypatch):
         assert get_coldkey_password_from_environment(wallet) == password
 
     assert get_coldkey_password_from_environment("non_existent_wallet") is None
+
+
+def test_keyfile_error_incorrect_password(keyfile_setup_teardown):
+    """
+    Test case for attempting to decrypt a keyfile with an incorrect password.
+    """
+    root_path = keyfile_setup_teardown
+    keyfile = bittensor.keyfile(path=os.path.join(root_path, "keyfile"))
+
+    # Ensure the keyfile is encrypted
+    assert keyfile.is_encrypted()
+
+    # Attempt to decrypt with an incorrect password
+    with pytest.raises(bittensor.KeyFileError) as excinfo:
+        keyfile.get_keypair(password="incorrect_password")
+
+    assert "Invalid password" in str(excinfo.value)