From c10de67c4ee4fdbe3c87f9a1bd01546d59d320b2 Mon Sep 17 00:00:00 2001
From: Benjamin Himes <benhimes@opentensor.dev>
Date: Thu, 7 Nov 2024 22:24:47 +0200
Subject: [PATCH 1/5] Handles the new PasswordError from bt-wallet

---
 .../src/bittensor/extrinsics/registration.py  | 10 +++----
 .../src/bittensor/extrinsics/root.py          | 11 ++-----
 .../src/bittensor/extrinsics/transfer.py      |  6 ++--
 bittensor_cli/src/bittensor/utils.py          | 26 ++++++++++++++++-
 bittensor_cli/src/commands/root.py            | 29 ++++---------------
 .../src/commands/stake/children_hotkeys.py    | 11 +++----
 bittensor_cli/src/commands/stake/stake.py     | 23 ++++-----------
 bittensor_cli/src/commands/subnets.py         |  6 ++--
 bittensor_cli/src/commands/sudo.py            |  6 ++--
 bittensor_cli/src/commands/wallets.py         | 18 ++++--------
 10 files changed, 63 insertions(+), 83 deletions(-)

diff --git a/bittensor_cli/src/bittensor/extrinsics/registration.py b/bittensor_cli/src/bittensor/extrinsics/registration.py
index 584d1291c..9b7ca42e1 100644
--- a/bittensor_cli/src/bittensor/extrinsics/registration.py
+++ b/bittensor_cli/src/bittensor/extrinsics/registration.py
@@ -20,7 +20,7 @@
 
 import backoff
 from bittensor_wallet import Wallet
-from bittensor_wallet.errors import KeyFileError
+from bittensor_wallet.errors import KeyFileError, PasswordError
 from Crypto.Hash import keccak
 import numpy as np
 from rich.prompt import Confirm
@@ -37,6 +37,7 @@
     get_human_readable,
     print_verbose,
     print_error,
+    unlock_key,
 )
 
 if typing.TYPE_CHECKING:
@@ -728,6 +729,8 @@ async def run_faucet_extrinsic(
     # Unlock coldkey
     try:
         wallet.unlock_coldkey()
+    except PasswordError:
+        return False, "Incorrect unlock password"
     except KeyFileError:
         return False, "There was an error unlocking your coldkey"
 
@@ -1639,10 +1642,7 @@ async def swap_hotkey_extrinsic(
         )
         return False
 
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet):
         return False
 
     if prompt:
diff --git a/bittensor_cli/src/bittensor/extrinsics/root.py b/bittensor_cli/src/bittensor/extrinsics/root.py
index 295ee640b..362b820e8 100644
--- a/bittensor_cli/src/bittensor/extrinsics/root.py
+++ b/bittensor_cli/src/bittensor/extrinsics/root.py
@@ -37,6 +37,7 @@
     u16_normalized_float,
     print_verbose,
     format_error_message,
+    unlock_key,
 )
 
 if TYPE_CHECKING:
@@ -306,10 +307,7 @@ async def root_register_extrinsic(
              the response is `True`.
     """
 
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet):
         return False
 
     print_verbose(f"Checking if hotkey ({wallet.hotkey_str}) is registered on root")
@@ -427,10 +425,7 @@ async def _do_set_weights():
         err_console.print("Your hotkey is not registered to the root network")
         return False
 
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet):
         return False
 
     # First convert types.
diff --git a/bittensor_cli/src/bittensor/extrinsics/transfer.py b/bittensor_cli/src/bittensor/extrinsics/transfer.py
index eada5ec70..391d6f288 100644
--- a/bittensor_cli/src/bittensor/extrinsics/transfer.py
+++ b/bittensor_cli/src/bittensor/extrinsics/transfer.py
@@ -16,6 +16,7 @@
     get_explorer_url_for_network,
     is_valid_bittensor_address_or_public_key,
     print_error,
+    unlock_key,
 )
 
 
@@ -115,10 +116,7 @@ async def do_transfer() -> tuple[bool, str, str]:
         return False
     console.print(f"[dark_orange]Initiating transfer on network: {subtensor.network}")
     # Unlock wallet coldkey.
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet):
         return False
 
     # Check balance.
diff --git a/bittensor_cli/src/bittensor/utils.py b/bittensor_cli/src/bittensor/utils.py
index 2a7f2b89e..f2f97a694 100644
--- a/bittensor_cli/src/bittensor/utils.py
+++ b/bittensor_cli/src/bittensor/utils.py
@@ -9,7 +9,7 @@
 
 from bittensor_wallet import Wallet, Keypair
 from bittensor_wallet.utils import SS58_FORMAT
-from bittensor_wallet.errors import KeyFileError
+from bittensor_wallet.errors import KeyFileError, PasswordError
 from bittensor_wallet import utils
 from jinja2 import Template
 from markupsafe import Markup
@@ -977,3 +977,27 @@ def retry_prompt(
             return var
         else:
             err_console.print(rejection_text)
+
+
+def unlock_key(wallet: Wallet, unlock_type="cold") -> bool:
+    if unlock_type == "cold":
+        unlocker = "unlock_coldkey"
+    elif unlock_type == "hot":
+        unlocker = "unlock_hotkey"
+    else:
+        raise ValueError(
+            f"Invalid unlock type provided: {unlock_type}. Must be 'cold' or 'hot'."
+        )
+    try:
+        getattr(wallet, unlocker)()
+        return True
+    except PasswordError:
+        err_console.print(
+            ":cross_mark: [red]The password used to decrypt your Keyfile is invalid.[/red]"
+        )
+        return False
+    except KeyFileError:
+        err_console.print(
+            ":cross_mark: [red]Keyfile is corrupt, non-writable, or non-readable.[/red]"
+        )
+        return False
diff --git a/bittensor_cli/src/commands/root.py b/bittensor_cli/src/commands/root.py
index 2401eb002..e874d7c1b 100644
--- a/bittensor_cli/src/commands/root.py
+++ b/bittensor_cli/src/commands/root.py
@@ -42,6 +42,7 @@
     ss58_to_vec_u8,
     update_metadata_table,
     group_subnets,
+    unlock_key,
 )
 
 if TYPE_CHECKING:
@@ -280,10 +281,7 @@ async def burned_register_extrinsic(
              finalization/inclusion, the response is `True`.
     """
 
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet):
         return False
 
     with console.status(
@@ -537,10 +535,7 @@ async def get_stake_for_coldkey_and_hotkey(
     delegate_string = "delegate" if delegate else "undelegate"
 
     # Decrypt key
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet):
         return False
 
     print_verbose("Checking if hotkey is a delegate")
@@ -1098,11 +1093,7 @@ async def senate_vote(
         return False
 
     # Unlock the wallet.
-    try:
-        wallet.unlock_hotkey()
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet) and unlock_key(wallet, "hot"):
         return False
 
     console.print(f"Fetching proposals in [dark_orange]network: {subtensor.network}")
@@ -1322,11 +1313,7 @@ async def _do_set_take() -> bool:
 
     console.print(f"Setting take on [dark_orange]network: {subtensor.network}")
     # Unlock the wallet.
-    try:
-        wallet.unlock_hotkey()
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet) and unlock_key(wallet, "hot"):
         return False
 
     result_ = await _do_set_take()
@@ -1724,11 +1711,7 @@ async def nominate(wallet: Wallet, subtensor: SubtensorInterface, prompt: bool):
 
     console.print(f"Nominating on [dark_orange]network: {subtensor.network}")
     # Unlock the wallet.
-    try:
-        wallet.unlock_hotkey()
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet) and unlock_key(wallet, "hot"):
         return False
 
     print_verbose(f"Checking hotkey ({wallet.hotkey_str}) is a delegate")
diff --git a/bittensor_cli/src/commands/stake/children_hotkeys.py b/bittensor_cli/src/commands/stake/children_hotkeys.py
index 97f8e5fc9..0a3afe5da 100644
--- a/bittensor_cli/src/commands/stake/children_hotkeys.py
+++ b/bittensor_cli/src/commands/stake/children_hotkeys.py
@@ -2,7 +2,7 @@
 from typing import Optional
 
 from bittensor_wallet import Wallet
-from bittensor_wallet.errors import KeyFileError
+from bittensor_wallet.errors import KeyFileError, PasswordError
 from rich.prompt import Confirm, Prompt, IntPrompt
 from rich.table import Table
 from rich.text import Text
@@ -19,6 +19,7 @@
     u64_to_float,
     is_valid_ss58_address,
     format_error_message,
+    unlock_key,
 )
 
 
@@ -72,10 +73,8 @@ async def set_children_extrinsic(
                 return False, "Operation Cancelled"
 
     # Decrypt coldkey.
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        return False, "There was an error unlocking your coldkey."
+    if not unlock_key(wallet):
+        return False
 
     with console.status(
         f":satellite: {operation} on [white]{subtensor.network}[/white] ..."
@@ -160,6 +159,8 @@ async def set_childkey_take_extrinsic(
     # Decrypt coldkey.
     try:
         wallet.unlock_coldkey()
+    except PasswordError:
+        return False, "Incorrect decrypt password"
     except KeyFileError:
         return False, "There was an error unlocking your coldkey."
 
diff --git a/bittensor_cli/src/commands/stake/stake.py b/bittensor_cli/src/commands/stake/stake.py
index 144e40f84..e147f0660 100644
--- a/bittensor_cli/src/commands/stake/stake.py
+++ b/bittensor_cli/src/commands/stake/stake.py
@@ -28,6 +28,7 @@
     render_tree,
     u16_normalized_float,
     validate_coldkey_presence,
+    unlock_key,
 )
 
 if TYPE_CHECKING:
@@ -103,10 +104,7 @@ async def add_stake_extrinsic(
     """
 
     # Decrypt keys,
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet):
         return False
 
     # Default to wallet's own hotkey if the value is not passed.
@@ -310,10 +308,7 @@ async def add_stake_multiple_extrinsic(
             return True
 
     # Decrypt coldkey.
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet):
         return False
 
     with console.status(
@@ -491,11 +486,8 @@ async def unstake_extrinsic(
     :return: success: `True` if extrinsic was finalized or included in the block. If we did not wait for
                       finalization/inclusion, the response is `True`.
     """
-    # Decrypt keys,
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    # Decrypt coldkey
+    if not unlock_key(wallet):
         return False
 
     if hotkey_ss58 is None:
@@ -663,10 +655,7 @@ async def unstake_multiple_extrinsic(
             return True
 
     # Unlock coldkey.
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet):
         return False
 
     with console.status(
diff --git a/bittensor_cli/src/commands/subnets.py b/bittensor_cli/src/commands/subnets.py
index dba96bfd2..000ed96fe 100644
--- a/bittensor_cli/src/commands/subnets.py
+++ b/bittensor_cli/src/commands/subnets.py
@@ -28,6 +28,7 @@
     millify,
     render_table,
     update_metadata_table,
+    unlock_key,
 )
 
 if TYPE_CHECKING:
@@ -100,10 +101,7 @@ async def _find_event_attributes_in_extrinsic_receipt(
         ):
             return False
 
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet):
         return False
 
     with console.status(":satellite: Registering subnet...", spinner="earth"):
diff --git a/bittensor_cli/src/commands/sudo.py b/bittensor_cli/src/commands/sudo.py
index 6ebbb0eca..f082bcda7 100644
--- a/bittensor_cli/src/commands/sudo.py
+++ b/bittensor_cli/src/commands/sudo.py
@@ -14,6 +14,7 @@
     print_error,
     print_verbose,
     normalize_hyperparameters,
+    unlock_key,
 )
 
 if TYPE_CHECKING:
@@ -101,10 +102,7 @@ async def set_hyperparameter_extrinsic(
         )
         return False
 
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet):
         return False
 
     extrinsic = HYPERPARAMS.get(parameter)
diff --git a/bittensor_cli/src/commands/wallets.py b/bittensor_cli/src/commands/wallets.py
index acce3b1a4..a9e9956b2 100644
--- a/bittensor_cli/src/commands/wallets.py
+++ b/bittensor_cli/src/commands/wallets.py
@@ -55,6 +55,7 @@
     is_valid_ss58_address,
     validate_coldkey_presence,
     retry_prompt,
+    unlock_key,
 )
 
 
@@ -1616,10 +1617,7 @@ async def set_id(
             print_error(f":cross_mark: This wallet doesn't own subnet {subnet_netuid}.")
             return False
 
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet):
         return False
 
     with console.status(
@@ -1719,18 +1717,14 @@ async def check_coldkey_swap(wallet: Wallet, subtensor: SubtensorInterface):
 
 async def sign(wallet: Wallet, message: str, use_hotkey: str):
     """Sign a message using the provided wallet or hotkey."""
-
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_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]"
-        )
     if not use_hotkey:
+        if not unlock_key(wallet):
+            return False
         keypair = wallet.coldkey
         print_verbose(f"Signing using coldkey: {wallet.name}")
     else:
+        if not unlock_key(wallet, "hot"):
+            return False
         keypair = wallet.hotkey
         print_verbose(f"Signing using hotkey: {wallet.hotkey_str}")
 

From 2adedcb5577aebc80a146586aa8cbf3f04bef86b Mon Sep 17 00:00:00 2001
From: Benjamin Himes <benhimes@opentensor.dev>
Date: Thu, 7 Nov 2024 22:28:51 +0200
Subject: [PATCH 2/5] Addeda period

---
 bittensor_cli/src/commands/stake/children_hotkeys.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/bittensor_cli/src/commands/stake/children_hotkeys.py b/bittensor_cli/src/commands/stake/children_hotkeys.py
index 0a3afe5da..9d457cbbe 100644
--- a/bittensor_cli/src/commands/stake/children_hotkeys.py
+++ b/bittensor_cli/src/commands/stake/children_hotkeys.py
@@ -160,7 +160,7 @@ async def set_childkey_take_extrinsic(
     try:
         wallet.unlock_coldkey()
     except PasswordError:
-        return False, "Incorrect decrypt password"
+        return False, "Incorrect decrypt password."
     except KeyFileError:
         return False, "There was an error unlocking your coldkey."
 

From aba921a47ec2f94b38417182793a3e71eb522203 Mon Sep 17 00:00:00 2001
From: Benjamin Himes <benhimes@opentensor.dev>
Date: Thu, 7 Nov 2024 22:50:52 +0200
Subject: [PATCH 3/5] Updated unlock_coldkey to handle optional printing in
 place.

---
 .../src/bittensor/extrinsics/registration.py  | 11 ++----
 .../src/bittensor/extrinsics/root.py          |  5 ++-
 .../src/bittensor/extrinsics/transfer.py      |  3 +-
 bittensor_cli/src/bittensor/utils.py          | 35 +++++++++++++------
 bittensor_cli/src/commands/root.py            | 11 +++---
 .../src/commands/stake/children_hotkeys.py    | 13 +++----
 bittensor_cli/src/commands/stake/stake.py     |  9 +++--
 bittensor_cli/src/commands/subnets.py         |  3 +-
 bittensor_cli/src/commands/sudo.py            |  3 +-
 bittensor_cli/src/commands/wallets.py         |  6 ++--
 10 files changed, 49 insertions(+), 50 deletions(-)

diff --git a/bittensor_cli/src/bittensor/extrinsics/registration.py b/bittensor_cli/src/bittensor/extrinsics/registration.py
index 9b7ca42e1..d73837d2b 100644
--- a/bittensor_cli/src/bittensor/extrinsics/registration.py
+++ b/bittensor_cli/src/bittensor/extrinsics/registration.py
@@ -20,7 +20,6 @@
 
 import backoff
 from bittensor_wallet import Wallet
-from bittensor_wallet.errors import KeyFileError, PasswordError
 from Crypto.Hash import keccak
 import numpy as np
 from rich.prompt import Confirm
@@ -727,12 +726,8 @@ async def run_faucet_extrinsic(
         return False, "Requires torch"
 
     # Unlock coldkey
-    try:
-        wallet.unlock_coldkey()
-    except PasswordError:
-        return False, "Incorrect unlock password"
-    except KeyFileError:
-        return False, "There was an error unlocking your coldkey"
+    if not (unlock_status := unlock_key(wallet, print_out=False)).success:
+        return False, unlock_status.message
 
     # Get previous balance.
     old_balance = await subtensor.get_balance(wallet.coldkeypub.ss58_address)
@@ -1642,7 +1637,7 @@ async def swap_hotkey_extrinsic(
         )
         return False
 
-    if not unlock_key(wallet):
+    if not unlock_key(wallet).success:
         return False
 
     if prompt:
diff --git a/bittensor_cli/src/bittensor/extrinsics/root.py b/bittensor_cli/src/bittensor/extrinsics/root.py
index 362b820e8..22ed70a04 100644
--- a/bittensor_cli/src/bittensor/extrinsics/root.py
+++ b/bittensor_cli/src/bittensor/extrinsics/root.py
@@ -21,7 +21,6 @@
 from typing import Union, List, TYPE_CHECKING
 
 from bittensor_wallet import Wallet, Keypair
-from bittensor_wallet.errors import KeyFileError
 import numpy as np
 from numpy.typing import NDArray
 from rich.prompt import Confirm
@@ -307,7 +306,7 @@ async def root_register_extrinsic(
              the response is `True`.
     """
 
-    if not unlock_key(wallet):
+    if not unlock_key(wallet).success:
         return False
 
     print_verbose(f"Checking if hotkey ({wallet.hotkey_str}) is registered on root")
@@ -425,7 +424,7 @@ async def _do_set_weights():
         err_console.print("Your hotkey is not registered to the root network")
         return False
 
-    if not unlock_key(wallet):
+    if not unlock_key(wallet).success:
         return False
 
     # First convert types.
diff --git a/bittensor_cli/src/bittensor/extrinsics/transfer.py b/bittensor_cli/src/bittensor/extrinsics/transfer.py
index 391d6f288..3ca657cab 100644
--- a/bittensor_cli/src/bittensor/extrinsics/transfer.py
+++ b/bittensor_cli/src/bittensor/extrinsics/transfer.py
@@ -1,7 +1,6 @@
 import asyncio
 
 from bittensor_wallet import Wallet
-from bittensor_wallet.errors import KeyFileError
 from rich.prompt import Confirm
 from substrateinterface.exceptions import SubstrateRequestException
 
@@ -116,7 +115,7 @@ async def do_transfer() -> tuple[bool, str, str]:
         return False
     console.print(f"[dark_orange]Initiating transfer on network: {subtensor.network}")
     # Unlock wallet coldkey.
-    if not unlock_key(wallet):
+    if not unlock_key(wallet).success:
         return False
 
     # Check balance.
diff --git a/bittensor_cli/src/bittensor/utils.py b/bittensor_cli/src/bittensor/utils.py
index f2f97a694..ad469148d 100644
--- a/bittensor_cli/src/bittensor/utils.py
+++ b/bittensor_cli/src/bittensor/utils.py
@@ -1,4 +1,5 @@
 import ast
+from collections import namedtuple
 import math
 import os
 import sqlite3
@@ -35,6 +36,8 @@
 err_console = Console(stderr=True)
 verbose_console = Console(quiet=True)
 
+UnlockStatus = namedtuple("UnlockStatus", ["success", "message"])
+
 
 def print_console(message: str, colour: str, title: str, console: Console):
     console.print(
@@ -979,7 +982,19 @@ def retry_prompt(
             err_console.print(rejection_text)
 
 
-def unlock_key(wallet: Wallet, unlock_type="cold") -> bool:
+def unlock_key(
+    wallet: Wallet, unlock_type="cold", print_out: bool = True
+) -> "UnlockStatus":
+    """
+    Attempts to decrypt a wallet's coldkey or hotkey
+    Args:
+        wallet: a Wallet object
+        unlock_type: the key type, 'cold' or 'hot'
+        print_out:  whether to print out the error message to the err_console
+
+    Returns: UnlockStatus for success status of unlock, with error message if unsuccessful
+
+    """
     if unlock_type == "cold":
         unlocker = "unlock_coldkey"
     elif unlock_type == "hot":
@@ -990,14 +1005,14 @@ def unlock_key(wallet: Wallet, unlock_type="cold") -> bool:
         )
     try:
         getattr(wallet, unlocker)()
-        return True
+        return UnlockStatus(True, "")
     except PasswordError:
-        err_console.print(
-            ":cross_mark: [red]The password used to decrypt your Keyfile is invalid.[/red]"
-        )
-        return False
+        err_msg = "The password used to decrypt your Keyfile is invalid."
+        if print_out:
+            err_console.print(f":cross_mark: [red]{err_msg}[/red]")
+        return UnlockStatus(False, err_msg)
     except KeyFileError:
-        err_console.print(
-            ":cross_mark: [red]Keyfile is corrupt, non-writable, or non-readable.[/red]"
-        )
-        return False
+        err_msg = "Keyfile is corrupt, non-writable, or non-readable."
+        if print_out:
+            err_console.print(f":cross_mark: [red]{err_msg}[/red]")
+        return UnlockStatus(False, err_msg)
diff --git a/bittensor_cli/src/commands/root.py b/bittensor_cli/src/commands/root.py
index e874d7c1b..0af9c4ebb 100644
--- a/bittensor_cli/src/commands/root.py
+++ b/bittensor_cli/src/commands/root.py
@@ -3,7 +3,6 @@
 from typing import Optional, TYPE_CHECKING
 
 from bittensor_wallet import Wallet
-from bittensor_wallet.errors import KeyFileError
 import numpy as np
 from numpy.typing import NDArray
 from rich import box
@@ -281,7 +280,7 @@ async def burned_register_extrinsic(
              finalization/inclusion, the response is `True`.
     """
 
-    if not unlock_key(wallet):
+    if not unlock_key(wallet).success:
         return False
 
     with console.status(
@@ -535,7 +534,7 @@ async def get_stake_for_coldkey_and_hotkey(
     delegate_string = "delegate" if delegate else "undelegate"
 
     # Decrypt key
-    if not unlock_key(wallet):
+    if not unlock_key(wallet).success:
         return False
 
     print_verbose("Checking if hotkey is a delegate")
@@ -1093,7 +1092,7 @@ async def senate_vote(
         return False
 
     # Unlock the wallet.
-    if not unlock_key(wallet) and unlock_key(wallet, "hot"):
+    if not unlock_key(wallet).success and unlock_key(wallet, "hot").success:
         return False
 
     console.print(f"Fetching proposals in [dark_orange]network: {subtensor.network}")
@@ -1313,7 +1312,7 @@ async def _do_set_take() -> bool:
 
     console.print(f"Setting take on [dark_orange]network: {subtensor.network}")
     # Unlock the wallet.
-    if not unlock_key(wallet) and unlock_key(wallet, "hot"):
+    if not unlock_key(wallet).success and unlock_key(wallet, "hot").success:
         return False
 
     result_ = await _do_set_take()
@@ -1711,7 +1710,7 @@ async def nominate(wallet: Wallet, subtensor: SubtensorInterface, prompt: bool):
 
     console.print(f"Nominating on [dark_orange]network: {subtensor.network}")
     # Unlock the wallet.
-    if not unlock_key(wallet) and unlock_key(wallet, "hot"):
+    if not unlock_key(wallet).success and unlock_key(wallet, "hot").success:
         return False
 
     print_verbose(f"Checking hotkey ({wallet.hotkey_str}) is a delegate")
diff --git a/bittensor_cli/src/commands/stake/children_hotkeys.py b/bittensor_cli/src/commands/stake/children_hotkeys.py
index 9d457cbbe..c64d486f6 100644
--- a/bittensor_cli/src/commands/stake/children_hotkeys.py
+++ b/bittensor_cli/src/commands/stake/children_hotkeys.py
@@ -2,7 +2,6 @@
 from typing import Optional
 
 from bittensor_wallet import Wallet
-from bittensor_wallet.errors import KeyFileError, PasswordError
 from rich.prompt import Confirm, Prompt, IntPrompt
 from rich.table import Table
 from rich.text import Text
@@ -73,8 +72,8 @@ async def set_children_extrinsic(
                 return False, "Operation Cancelled"
 
     # Decrypt coldkey.
-    if not unlock_key(wallet):
-        return False
+    if not (unlock_status := unlock_key(wallet, print_out=False)).success:
+        return False, unlock_status.message
 
     with console.status(
         f":satellite: {operation} on [white]{subtensor.network}[/white] ..."
@@ -157,12 +156,8 @@ async def set_childkey_take_extrinsic(
             return False, "Operation Cancelled"
 
     # Decrypt coldkey.
-    try:
-        wallet.unlock_coldkey()
-    except PasswordError:
-        return False, "Incorrect decrypt password."
-    except KeyFileError:
-        return False, "There was an error unlocking your coldkey."
+    if not (unlock_status := unlock_key(wallet, print_out=False)).success:
+        return False, unlock_status.message
 
     with console.status(
         f":satellite: Setting childkey take on [white]{subtensor.network}[/white] ..."
diff --git a/bittensor_cli/src/commands/stake/stake.py b/bittensor_cli/src/commands/stake/stake.py
index e147f0660..7c7c83fe0 100644
--- a/bittensor_cli/src/commands/stake/stake.py
+++ b/bittensor_cli/src/commands/stake/stake.py
@@ -7,7 +7,6 @@
 from typing import TYPE_CHECKING, Optional, Sequence, Union, cast
 
 from bittensor_wallet import Wallet
-from bittensor_wallet.errors import KeyFileError
 from rich.prompt import Confirm
 from rich.table import Table, Column
 import typer
@@ -104,7 +103,7 @@ async def add_stake_extrinsic(
     """
 
     # Decrypt keys,
-    if not unlock_key(wallet):
+    if not unlock_key(wallet).success:
         return False
 
     # Default to wallet's own hotkey if the value is not passed.
@@ -308,7 +307,7 @@ async def add_stake_multiple_extrinsic(
             return True
 
     # Decrypt coldkey.
-    if not unlock_key(wallet):
+    if not unlock_key(wallet).success:
         return False
 
     with console.status(
@@ -487,7 +486,7 @@ async def unstake_extrinsic(
                       finalization/inclusion, the response is `True`.
     """
     # Decrypt coldkey
-    if not unlock_key(wallet):
+    if not unlock_key(wallet).success:
         return False
 
     if hotkey_ss58 is None:
@@ -655,7 +654,7 @@ async def unstake_multiple_extrinsic(
             return True
 
     # Unlock coldkey.
-    if not unlock_key(wallet):
+    if not unlock_key(wallet).success:
         return False
 
     with console.status(
diff --git a/bittensor_cli/src/commands/subnets.py b/bittensor_cli/src/commands/subnets.py
index 000ed96fe..1d4620b73 100644
--- a/bittensor_cli/src/commands/subnets.py
+++ b/bittensor_cli/src/commands/subnets.py
@@ -5,7 +5,6 @@
 from typing import TYPE_CHECKING, Optional, cast
 
 from bittensor_wallet import Wallet
-from bittensor_wallet.errors import KeyFileError
 from rich.prompt import Confirm
 from rich.table import Column, Table
 
@@ -101,7 +100,7 @@ async def _find_event_attributes_in_extrinsic_receipt(
         ):
             return False
 
-    if not unlock_key(wallet):
+    if not unlock_key(wallet).success:
         return False
 
     with console.status(":satellite: Registering subnet...", spinner="earth"):
diff --git a/bittensor_cli/src/commands/sudo.py b/bittensor_cli/src/commands/sudo.py
index f082bcda7..5711a3c27 100644
--- a/bittensor_cli/src/commands/sudo.py
+++ b/bittensor_cli/src/commands/sudo.py
@@ -2,7 +2,6 @@
 from typing import TYPE_CHECKING, Union
 
 from bittensor_wallet import Wallet
-from bittensor_wallet.errors import KeyFileError
 from rich import box
 from rich.table import Column, Table
 
@@ -102,7 +101,7 @@ async def set_hyperparameter_extrinsic(
         )
         return False
 
-    if not unlock_key(wallet):
+    if not unlock_key(wallet).success:
         return False
 
     extrinsic = HYPERPARAMS.get(parameter)
diff --git a/bittensor_cli/src/commands/wallets.py b/bittensor_cli/src/commands/wallets.py
index a9e9956b2..69188053e 100644
--- a/bittensor_cli/src/commands/wallets.py
+++ b/bittensor_cli/src/commands/wallets.py
@@ -1617,7 +1617,7 @@ async def set_id(
             print_error(f":cross_mark: This wallet doesn't own subnet {subnet_netuid}.")
             return False
 
-    if not unlock_key(wallet):
+    if not unlock_key(wallet).success:
         return False
 
     with console.status(
@@ -1718,12 +1718,12 @@ async def check_coldkey_swap(wallet: Wallet, subtensor: SubtensorInterface):
 async def sign(wallet: Wallet, message: str, use_hotkey: str):
     """Sign a message using the provided wallet or hotkey."""
     if not use_hotkey:
-        if not unlock_key(wallet):
+        if not unlock_key(wallet).success:
             return False
         keypair = wallet.coldkey
         print_verbose(f"Signing using coldkey: {wallet.name}")
     else:
-        if not unlock_key(wallet, "hot"):
+        if not unlock_key(wallet, "hot").success:
             return False
         keypair = wallet.hotkey
         print_verbose(f"Signing using hotkey: {wallet.hotkey_str}")

From cbd2977e2c36378f6ca8316534886f7da4e6768b Mon Sep 17 00:00:00 2001
From: Benjamin Himes <benhimes@opentensor.dev>
Date: Fri, 8 Nov 2024 15:46:31 +0200
Subject: [PATCH 4/5] Updated wording.

---
 bittensor_cli/src/bittensor/utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/bittensor_cli/src/bittensor/utils.py b/bittensor_cli/src/bittensor/utils.py
index ad469148d..a3954a85d 100644
--- a/bittensor_cli/src/bittensor/utils.py
+++ b/bittensor_cli/src/bittensor/utils.py
@@ -1007,12 +1007,12 @@ def unlock_key(
         getattr(wallet, unlocker)()
         return UnlockStatus(True, "")
     except PasswordError:
-        err_msg = "The password used to decrypt your Keyfile is invalid."
+        err_msg = f"The password used to decrypt your {unlock_type.capitalize()}key Keyfile is invalid."
         if print_out:
             err_console.print(f":cross_mark: [red]{err_msg}[/red]")
         return UnlockStatus(False, err_msg)
     except KeyFileError:
-        err_msg = "Keyfile is corrupt, non-writable, or non-readable."
+        err_msg = f"{unlock_type.capitalize()}key Keyfile is corrupt, non-writable, or non-readable, or non-existent."
         if print_out:
             err_console.print(f":cross_mark: [red]{err_msg}[/red]")
         return UnlockStatus(False, err_msg)

From 7fe4d7f8a715dd63c7a4b9661e83ef8a5c2606d0 Mon Sep 17 00:00:00 2001
From: Benjamin Himes <benhimes@opentensor.dev>
Date: Wed, 13 Nov 2024 16:44:28 +0200
Subject: [PATCH 5/5] Bump bt-wallet req

---
 requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index fbe806998..bb50a3a7a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -17,5 +17,5 @@ scalecodec==1.2.11
 substrate-interface~=1.7.9
 typer~=0.12
 websockets>=12.0
-bittensor-wallet>=2.0.2
+bittensor-wallet>=2.1.0
 bt-decode==0.2.0a0
\ No newline at end of file