From f1477d5e7d4380114970bd7ca9f4fc7fc4e49ab1 Mon Sep 17 00:00:00 2001 From: Gus Date: Fri, 16 Aug 2024 16:50:40 -0400 Subject: [PATCH] feat: return error message instead raising exception --- bittensor/commands/delegates.py | 8 +++++- bittensor/commands/identity.py | 9 ++++++- bittensor/commands/senate.py | 26 ++++++++++++++++--- bittensor/extrinsics/delegation.py | 38 ++++++++++++++++++++++++---- bittensor/extrinsics/network.py | 8 +++++- bittensor/extrinsics/registration.py | 24 +++++++++++++++--- bittensor/extrinsics/root.py | 16 ++++++++++-- bittensor/extrinsics/senate.py | 18 +++++++++++-- bittensor/extrinsics/staking.py | 8 +++++- bittensor/extrinsics/transfer.py | 11 ++++++-- bittensor/extrinsics/unstaking.py | 24 +++++++++++++++--- bittensor/keyfile.py | 6 ++++- tests/unit_tests/test_keyfile.py | 17 +++++++++++++ 13 files changed, 188 insertions(+), 25 deletions(-) diff --git a/bittensor/commands/delegates.py b/bittensor/commands/delegates.py index 4d03b289e4..cfba3526d2 100644 --- a/bittensor/commands/delegates.py +++ b/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): diff --git a/bittensor/commands/identity.py b/bittensor/commands/identity.py index 15232c4440..4f74548495 100644 --- a/bittensor/commands/identity.py +++ b/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( diff --git a/bittensor/commands/senate.py b/bittensor/commands/senate.py index 03a73cde5b..37f2d79585 100644 --- a/bittensor/commands/senate.py +++ b/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: diff --git a/bittensor/extrinsics/delegation.py b/bittensor/extrinsics/delegation.py index 5d31855cdb..e61a97efb4 100644 --- a/bittensor/extrinsics/delegation.py +++ b/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( diff --git a/bittensor/extrinsics/network.py b/bittensor/extrinsics/network.py index 16cbc0ed26..5aecaa459a 100644 --- a/bittensor/extrinsics/network.py +++ b/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: diff --git a/bittensor/extrinsics/registration.py b/bittensor/extrinsics/registration.py index e82add8383..40bde3fc89 100644 --- a/bittensor/extrinsics/registration.py +++ b/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( diff --git a/bittensor/extrinsics/root.py b/bittensor/extrinsics/root.py index 8a7e9e3863..c0a4fcabd1 100644 --- a/bittensor/extrinsics/root.py +++ b/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): diff --git a/bittensor/extrinsics/senate.py b/bittensor/extrinsics/senate.py index 043233996c..f586cec399 100644 --- a/bittensor/extrinsics/senate.py +++ b/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: diff --git a/bittensor/extrinsics/staking.py b/bittensor/extrinsics/staking.py index 686b7d60e5..93ac1ae2de 100644 --- a/bittensor/extrinsics/staking.py +++ b/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: diff --git a/bittensor/extrinsics/transfer.py b/bittensor/extrinsics/transfer.py index 91ef3237eb..aa340ab406 100644 --- a/bittensor/extrinsics/transfer.py +++ b/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): diff --git a/bittensor/extrinsics/unstaking.py b/bittensor/extrinsics/unstaking.py index 57329915eb..a5de71b7d7 100644 --- a/bittensor/extrinsics/unstaking.py +++ b/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 = [] diff --git a/bittensor/keyfile.py b/bittensor/keyfile.py index f1b2ad622e..d2c75c1041 100644 --- a/bittensor/keyfile.py +++ b/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) diff --git a/tests/unit_tests/test_keyfile.py b/tests/unit_tests/test_keyfile.py index 8db105c3bd..0f3b69cacf 100644 --- a/tests/unit_tests/test_keyfile.py +++ b/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)