From d0011f9a5ac0b3d11c1d5ff255ade6e9ea0ff5e2 Mon Sep 17 00:00:00 2001
From: Benjamin Himes <benhimes@opentensor.dev>
Date: Mon, 24 Feb 2025 16:09:29 +0200
Subject: [PATCH] Use `unlock_key` function.

---
 bittensor_cli/cli.py                          |  2 --
 bittensor_cli/src/commands/stake/add.py       |  6 ++----
 bittensor_cli/src/commands/stake/move.py      | 16 ++++------------
 bittensor_cli/src/commands/stake/remove.py    | 11 +++--------
 bittensor_cli/src/commands/subnets/subnets.py |  5 +----
 bittensor_cli/version.py                      |  2 ++
 6 files changed, 12 insertions(+), 30 deletions(-)

diff --git a/bittensor_cli/cli.py b/bittensor_cli/cli.py
index 2c5b70992..419fd8c32 100755
--- a/bittensor_cli/cli.py
+++ b/bittensor_cli/cli.py
@@ -74,8 +74,6 @@ class GitError(Exception):
         pass
 
 
-
-
 _epilog = "Made with [bold red]:heart:[/bold red] by The Openτensor Foundaτion"
 
 np.set_printoptions(precision=8, suppress=True, floatmode="fixed")
diff --git a/bittensor_cli/src/commands/stake/add.py b/bittensor_cli/src/commands/stake/add.py
index 4aba39e69..2950da876 100644
--- a/bittensor_cli/src/commands/stake/add.py
+++ b/bittensor_cli/src/commands/stake/add.py
@@ -17,6 +17,7 @@
     is_valid_ss58_address,
     print_error,
     print_verbose,
+    unlock_key,
 )
 from bittensor_wallet import Wallet
 from bittensor_wallet.errors import KeyFileError
@@ -338,10 +339,7 @@ async def stake_extrinsic(
     if prompt:
         if not Confirm.ask("Would you like to continue?"):
             raise typer.Exit()
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet).success:
         return False
 
     if safe_staking:
diff --git a/bittensor_cli/src/commands/stake/move.py b/bittensor_cli/src/commands/stake/move.py
index 964f0ed1b..876f500e6 100644
--- a/bittensor_cli/src/commands/stake/move.py
+++ b/bittensor_cli/src/commands/stake/move.py
@@ -17,6 +17,7 @@
     format_error_message,
     group_subnets,
     get_subnet_name,
+    unlock_key,
 )
 
 if TYPE_CHECKING:
@@ -621,10 +622,7 @@ async def move_stake(
             raise typer.Exit()
 
     # Perform moving operation.
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet).success:
         return False
     with console.status(
         f"\n:satellite: Moving [blue]{amount_to_move_as_balance}[/blue] from [blue]{origin_hotkey}[/blue] on netuid: [blue]{origin_netuid}[/blue] \nto "
@@ -777,10 +775,7 @@ async def transfer_stake(
             raise typer.Exit()
 
     # Perform transfer operation
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet).success:
         return False
 
     with console.status("\n:satellite: Transferring stake ..."):
@@ -933,10 +928,7 @@ async def swap_stake(
             raise typer.Exit()
 
     # Perform swap operation
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet).success:
         return False
 
     with console.status(
diff --git a/bittensor_cli/src/commands/stake/remove.py b/bittensor_cli/src/commands/stake/remove.py
index 048855fdc..fc298de47 100644
--- a/bittensor_cli/src/commands/stake/remove.py
+++ b/bittensor_cli/src/commands/stake/remove.py
@@ -21,6 +21,7 @@
     is_valid_ss58_address,
     format_error_message,
     group_subnets,
+    unlock_key,
 )
 
 if TYPE_CHECKING:
@@ -268,10 +269,7 @@ async def unstake(
             raise typer.Exit()
 
     # Execute extrinsics
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet).success:
         return False
 
     with console.status("\n:satellite: Performing unstaking operations...") as status:
@@ -465,10 +463,7 @@ async def unstake_all(
     ):
         return False
 
-    try:
-        wallet.unlock_coldkey()
-    except KeyFileError:
-        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+    if not unlock_key(wallet).success:
         return False
 
     with console.status("Unstaking all stakes...") as status:
diff --git a/bittensor_cli/src/commands/subnets/subnets.py b/bittensor_cli/src/commands/subnets/subnets.py
index f4e01486c..b419d81de 100644
--- a/bittensor_cli/src/commands/subnets/subnets.py
+++ b/bittensor_cli/src/commands/subnets/subnets.py
@@ -145,10 +145,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).success:
         return False
 
     with console.status(":satellite: Registering subnet...", spinner="earth"):
diff --git a/bittensor_cli/version.py b/bittensor_cli/version.py
index b47664c90..2daeadf01 100644
--- a/bittensor_cli/version.py
+++ b/bittensor_cli/version.py
@@ -1,5 +1,6 @@
 import re
 
+
 def version_as_int(version):
     _core_version = re.match(r"^\d+\.\d+\.\d+", version).group(0)
     _version_split = _core_version.split(".")
@@ -14,5 +15,6 @@ def version_as_int(version):
     __new_signature_version__ = 360
     return __version_as_int__
 
+
 __version__ = "9.0.0"
 __version_as_int__ = version_as_int(__version__)