Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: return error message instead of raising exception #2244

Merged
merged 2 commits into from
Aug 17, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion bittensor/commands/delegates.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
9 changes: 8 additions & 1 deletion bittensor/commands/identity.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
26 changes: 23 additions & 3 deletions bittensor/commands/senate.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand Down
38 changes: 33 additions & 5 deletions bittensor/extrinsics/delegation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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))

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
8 changes: 7 additions & 1 deletion bittensor/extrinsics/network.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
24 changes: 21 additions & 3 deletions bittensor/extrinsics/registration.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]..."
):
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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(
Expand Down
16 changes: 14 additions & 2 deletions bittensor/extrinsics/root.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down
18 changes: 16 additions & 2 deletions bittensor/extrinsics/senate.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
8 changes: 7 additions & 1 deletion bittensor/extrinsics/staking.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
11 changes: 9 additions & 2 deletions bittensor/extrinsics/transfer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
24 changes: 21 additions & 3 deletions bittensor/extrinsics/unstaking.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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 = []
Expand Down
6 changes: 5 additions & 1 deletion bittensor/keyfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
17 changes: 17 additions & 0 deletions tests/unit_tests/test_keyfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Loading