Handles new PasswordError from btwallet.

bittensor/core/extrinsics/async_registration.py

@@ -21,13 +21,12 @@
 import numpy as np
 from Crypto.Hash import keccak
 from bittensor_wallet import Wallet
-from bittensor_wallet.errors import KeyFileError
 from rich.console import Console
 from rich.status import Status
 from substrateinterface.exceptions import SubstrateRequestException
 
 from bittensor.core.chain_data import NeuronInfo
-from bittensor.utils import format_error_message
+from bittensor.utils import format_error_message, unlock_key
 from bittensor.utils.btlogging import logging
 from bittensor.utils.formatting import millify, get_human_readable
 
@@ -673,10 +672,8 @@ async def run_faucet_extrinsic(
         return False, "Requires torch"
 
     # Unlock coldkey
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        return False, "There was an error unlocking your coldkey"
+    if not (unlock := unlock_key(wallet)).success:
+        return False, unlock.message
 
     # Get previous balance.
     old_balance = await subtensor.get_balance(wallet.coldkeypub.ss58_address)

bittensor/core/extrinsics/async_root.py

@@ -4,11 +4,10 @@
 
 import numpy as np
 from bittensor_wallet import Wallet
-from bittensor_wallet.errors import KeyFileError
 from numpy.typing import NDArray
 from substrateinterface.exceptions import SubstrateRequestException
 
-from bittensor.utils import u16_normalized_float, format_error_message
+from bittensor.utils import u16_normalized_float, format_error_message, unlock_key
 from bittensor.utils.btlogging import logging
 from bittensor.utils.weight_utils import (
     normalize_max_weight,
@@ -63,10 +62,8 @@ async def root_register_extrinsic(
         `True` if extrinsic was finalized or included in the block. If we did not wait for finalization/inclusion, the response is `True`.
     """
 
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        logging.error("Error decrypting coldkey (possibly incorrect password)")
+    if not (unlock := unlock_key(wallet)).success:
+        logging.error(unlock.message)
         return False
 
     logging.debug(
@@ -182,10 +179,8 @@ async def _do_set_weights():
         logging.error("Your hotkey is not registered to the root network.")
         return False
 
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        logging.error("Error decrypting coldkey (possibly incorrect password).")
+    if not (unlock := unlock_key(wallet)).success:
+        logging.error(unlock.message)
         return False
 
     # First convert types.

bittensor/core/extrinsics/async_transfer.py

@@ -2,14 +2,13 @@
 from typing import TYPE_CHECKING
 
 from bittensor_wallet import Wallet
-from bittensor_wallet.errors import KeyFileError
 from substrateinterface.exceptions import SubstrateRequestException
 
 from bittensor.core.settings import NETWORK_EXPLORER_MAP
 from bittensor.utils import (
     format_error_message,
     get_explorer_url_for_network,
-    is_valid_bittensor_address_or_public_key,
+    is_valid_bittensor_address_or_public_key, unlock_key,
 )
 from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
@@ -116,10 +115,8 @@ async def do_transfer() -> tuple[bool, str, str]:
         return False
     logging.info(f"Initiating transfer on network: {subtensor.network}")
     # Unlock wallet coldkey.
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        logging.error("Error decrypting coldkey (possibly incorrect password)")
+    if not (unlock := unlock_key(wallet)).success:
+        logging.error(unlock.message)
         return False
 
     # Check balance.

bittensor/core/extrinsics/registration.py

@@ -18,9 +18,8 @@
 import time
 from typing import Union, Optional, TYPE_CHECKING
 
-from bittensor_wallet.errors import KeyFileError
 
-from bittensor.utils import format_error_message
+from bittensor.utils import format_error_message, unlock_key
 from bittensor.utils.btlogging import logging
 from bittensor.utils.networking import ensure_connected
 from bittensor.utils.registration import (
@@ -347,13 +346,10 @@ def burned_register_extrinsic(
         )
         return False
 
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        logging.error(
-            ":cross_mark: <red>Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid.</red>"
-        )
+    if not (unlock := unlock_key(wallet)).success:
+        logging.error(unlock.message)
         return False
+
     logging.info(
         f":satellite: <magenta>Checking Account on subnet</magenta> <blue>{netuid}</blue><magenta> ...</magenta>"
     )

bittensor/core/extrinsics/root.py

@@ -2,11 +2,10 @@
 from typing import Optional, Union, TYPE_CHECKING
 
 import numpy as np
-from bittensor_wallet.errors import KeyFileError
 from numpy.typing import NDArray
 
 from bittensor.core.settings import version_as_int
-from bittensor.utils import format_error_message, weight_utils
+from bittensor.utils import format_error_message, weight_utils, unlock_key
 from bittensor.utils.btlogging import logging
 from bittensor.utils.networking import ensure_connected
 from bittensor.utils.registration import torch, legacy_torch_api_compat
@@ -71,12 +70,8 @@ def root_register_extrinsic(
         success (bool): Flag is ``true`` if extrinsic was finalized or uncluded in the block. If we did not wait for finalization / inclusion, the response is ``true``.
     """
 
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        logging.error(
-            "<red>Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid.</red>"
-        )
+    if not (unlock := unlock_key(wallet)).success:
+        logging.error(unlock.message)
         return False
 
     is_registered = subtensor.is_hotkey_registered(
@@ -199,13 +194,8 @@ def set_root_weights_extrinsic(
     Returns:
         success (bool): Flag is ``true`` if extrinsic was finalized or uncluded in the block. If we did not wait for finalization / inclusion, the response is ``true``.
     """
-
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        logging.error(
-            ":cross_mark: <red>Keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid.</red>"
-        )
+    if not (unlock := unlock_key(wallet)).success:
+        logging.error(unlock.message)
         return False
 
     # First convert types.

bittensor/core/extrinsics/serving.py

@@ -20,7 +20,7 @@
 from bittensor.core.errors import MetadataError
 from bittensor.core.extrinsics.utils import submit_extrinsic
 from bittensor.core.settings import version_as_int
-from bittensor.utils import format_error_message, networking as net
+from bittensor.utils import format_error_message, networking as net, unlock_key
 from bittensor.utils.btlogging import logging
 from bittensor.utils.networking import ensure_connected
 
@@ -109,7 +109,10 @@ def serve_extrinsic(
         success (bool): 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 hotkey
-    wallet.unlock_hotkey()
+    if not (unlock := unlock_key(wallet, "hot")).success:
+        logging.error(unlock.message)
+        return False
+
     params: "AxonServeCallParams" = {
         "version": version_as_int,
         "ip": net.ip_to_int(ip),
@@ -192,8 +195,12 @@ def serve_axon_extrinsic(
     Returns:
         success (bool): Flag is ``true`` if extrinsic was finalized or uncluded in the block. If we did not wait for finalization / inclusion, the response is ``true``.
     """
-    axon.wallet.unlock_hotkey()
-    axon.wallet.unlock_coldkeypub()
+    if not (unlock := unlock_key(axon.wallet)).success:
+        logging.error(unlock.message)
+        return False
+    if not (unlock := unlock_key(axon.wallet, "coldkeypub")).success:
+        logging.error(unlock.message)
+        return False
     external_port = axon.external_port
 
     # ---- Get external ip ----
@@ -254,7 +261,9 @@ def publish_metadata(
         MetadataError: If there is an error in submitting the extrinsic or if the response from the blockchain indicates failure.
     """
 
-    wallet.unlock_hotkey()
+    if not (unlock := unlock_key(wallet, "hot")).success:
+        logging.error(unlock.message)
+        return False
 
     with self.substrate as substrate:
         call = substrate.compose_call(

bittensor/core/extrinsics/transfer.py

@@ -22,7 +22,7 @@
 from bittensor.utils import (
     get_explorer_url_for_network,
     format_error_message,
-    is_valid_bittensor_address_or_public_key,
+    is_valid_bittensor_address_or_public_key, unlock_key,
 )
 from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
@@ -120,8 +120,9 @@ def transfer_extrinsic(
         # Convert bytes to hex string.
         dest = "0x" + dest.hex()
 
-    # Unlock wallet coldkey.
-    wallet.unlock_coldkey()
+    if not (unlock := unlock_key(wallet)).success:
+        logging.error(unlock.message)
+        return False
 
     # Convert to bittensor.Balance
     if not isinstance(amount, Balance):

bittensor/utils/__init__.py

@@ -15,29 +15,35 @@
 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 # DEALINGS IN THE SOFTWARE.
 
-from urllib.parse import urlparse
 import ast
+from collections import namedtuple
 import hashlib
 from typing import Any, Literal, Union, Optional, TYPE_CHECKING
+from urllib.parse import urlparse
 
 import scalecodec
 from bittensor_wallet import Keypair
 from substrateinterface.utils import ss58
 
 from bittensor.core.settings import SS58_FORMAT
 from bittensor.utils.btlogging import logging
+from bittensor_wallet.errors import KeyFileError, PasswordError
 from .registration import torch, use_torch
 from .version import version_checking, check_version, VersionCheckError
 
 if TYPE_CHECKING:
     from bittensor.utils.async_substrate_interface import AsyncSubstrateInterface
     from substrateinterface import SubstrateInterface
+    from bittensor_wallet import Wallet
 
 RAOPERTAO = 1e9
 U16_MAX = 65535
 U64_MAX = 18446744073709551615
 
 
+UnlockStatus = namedtuple("UnlockStatus", ["success", "message"])
+
+
 def ss58_to_vec_u8(ss58_address: str) -> list[int]:
     ss58_bytes: bytes = ss58_address_to_bytes(ss58_address)
     encoded_address: list[int] = [int(byte) for byte in ss58_bytes]
@@ -370,3 +376,35 @@ def validate_chain_endpoint(endpoint_url: str) -> tuple[bool, str]:
     if not parsed.netloc:
         return False, "Invalid URL passed as the endpoint"
     return True, ""
+
+
+def unlock_key(
+    wallet: "Wallet", unlock_type="cold"
+) -> "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":
+        unlocker = "unlock_hotkey"
+    elif unlock_type == "coldkeypub":
+        unlocker = "unlock_coldkeypub"
+    else:
+        raise ValueError(
+            f"Invalid unlock type provided: {unlock_type}. Must be 'cold', 'hot', or 'coldkeypub'."
+        )
+    try:
+        getattr(wallet, unlocker)()
+        return UnlockStatus(True, "")
+    except PasswordError:
+        err_msg = f"The password used to decrypt your {unlock_type.capitalize()}key Keyfile is invalid."
+        return UnlockStatus(False, err_msg)
+    except KeyFileError:
+        err_msg = f"{unlock_type.capitalize()}key Keyfile is corrupt, non-writable, or non-readable, or non-existent."
+        return UnlockStatus(False, err_msg)