Handle new PasswordError from btwallet (#2406)

* Handles new PasswordError from btwallet.

* Ruff, docstrings.

* Fixes hotkey unlock rather than coldkey

* Added unit test.

* Added unit test. Add integration test.

* Removed coldkeypub

* Opinions.

* More tests

* grammar

* Update test_utils.py

* bump up btwallet version

* fix wording

* fix wording

* Update tests/unit_tests/utils/test_utils.py

Co-authored-by: Cameron Fairchild <[email protected]>

---------

Co-authored-by: Roman <[email protected]>
Co-authored-by: Roman <[email protected]>
Co-authored-by: Cameron Fairchild <[email protected]>

Author: 4 people

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,14 @@
 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,
+    unlock_key,
 )
 from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
@@ -116,10 +116,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, "hotkey")).success:
+        logging.error(unlock.message)
+        return False
+
     params: "AxonServeCallParams" = {
         "version": version_as_int,
         "ip": net.ip_to_int(ip),
@@ -192,8 +195,9 @@ 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, "hotkey")).success:
+        logging.error(unlock.message)
+        return False
     external_port = axon.external_port
 
     # ---- Get external ip ----
@@ -254,7 +258,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, "hotkey")).success:
+        logging.error(unlock.message)
+        return False
 
     with self.substrate as substrate:
         call = substrate.compose_call(

bittensor/core/extrinsics/transfer.py

@@ -23,6 +23,7 @@
     get_explorer_url_for_network,
     format_error_message,
     is_valid_bittensor_address_or_public_key,
+    unlock_key,
 )
 from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
@@ -120,8 +121,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,30 @@ 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="coldkey") -> "UnlockStatus":
+    """
+    Attempts to decrypt a wallet's coldkey or hotkey
+    Args:
+        wallet: a Wallet object
+        unlock_type: the key type, 'coldkey' or 'hotkey'
+    Returns: UnlockStatus for success status of unlock, with error message if unsuccessful
+    """
+    if unlock_type == "coldkey":
+        unlocker = "unlock_coldkey"
+    elif unlock_type == "hotkey":
+        unlocker = "unlock_hotkey"
+    else:
+        raise ValueError(
+            f"Invalid unlock type provided: {unlock_type}. Must be 'coldkey' or 'hotkey'."
+        )
+    try:
+        getattr(wallet, unlocker)()
+        return UnlockStatus(True, "")
+    except PasswordError:
+        err_msg = f"The password used to decrypt your {unlock_type.capitalize()} keyfile is invalid."
+        return UnlockStatus(False, err_msg)
+    except KeyFileError:
+        err_msg = f"{unlock_type.capitalize()} keyfile is corrupt, non-writable, or non-readable, or non-existent."
+        return UnlockStatus(False, err_msg)

requirements/prod.txt

@@ -25,4 +25,4 @@ scalecodec==1.2.11
 substrate-interface~=1.7.9
 uvicorn
 websockets>12.0
-bittensor-wallet>=2.0.2
+bittensor-wallet>=2.1.0

tests/integration_tests/utils/test_init.py

@@ -0,0 +1,42 @@
+import os
+import shutil
+
+from bittensor_wallet import Wallet, Keyfile, Keypair
+import pytest
+
+from bittensor import utils
+
+
+def test_unlock_key(monkeypatch):
+    # Ensure path is clean before we run the tests
+    if os.path.exists("/tmp/bittensor-tests-wallets"):
+        shutil.rmtree("/tmp/bittensor-tests-wallets")
+
+    wallet = Wallet(path="/tmp/bittensor-tests-wallets")
+    cold_kf = Keyfile("/tmp/bittensor-tests-wallets/default/coldkey", name="default")
+    kp = Keypair.create_from_mnemonic(
+        "stool feel open east woman high can denial forget screen trust salt"
+    )
+    cold_kf.set_keypair(kp, False, False)
+    cold_kf.encrypt("1234password1234")
+    hot_kf = Keyfile("/tmp/bittensor-tests-wallets/default/hotkey", name="default")
+    hkp = Keypair.create_from_mnemonic(
+        "stool feel open east woman high can denial forget screen trust salt"
+    )
+    hot_kf.set_keypair(hkp, False, False)
+    hot_kf.encrypt("1234hotkey1234")
+    monkeypatch.setattr("getpass.getpass", lambda _: "badpassword1234")
+    result = utils.unlock_key(wallet)
+    assert result.success is False
+    monkeypatch.setattr("getpass.getpass", lambda _: "1234password1234")
+    result = utils.unlock_key(wallet)
+    assert result.success is True
+    monkeypatch.setattr("getpass.getpass", lambda _: "badpassword1234")
+    result = utils.unlock_key(wallet, "hotkey")
+    assert result.success is False
+    with pytest.raises(ValueError):
+        utils.unlock_key(wallet, "mycoldkey")
+
+    # Ensure test wallets path is deleted after running tests
+    if os.path.exists("/tmp/bittensor-tests-wallets"):
+        shutil.rmtree("/tmp/bittensor-tests-wallets")