diff --git a/bittensor_cli/cli.py b/bittensor_cli/cli.py index 4b3998344..dbf13be3b 100755 --- a/bittensor_cli/cli.py +++ b/bittensor_cli/cli.py @@ -716,6 +716,9 @@ def __init__(self): self.wallet_app.command( "new-coldkey", rich_help_panel=HELP_PANELS["WALLET"]["MANAGEMENT"] )(self.wallet_new_coldkey) + self.wallet_app.command( + "associate-hotkey", rich_help_panel=HELP_PANELS["WALLET"]["MANAGEMENT"] + )(self.wallet_associate_hotkey) self.wallet_app.command( "create", rich_help_panel=HELP_PANELS["WALLET"]["MANAGEMENT"] )(self.wallet_create_wallet) @@ -2265,6 +2268,74 @@ def wallet_new_hotkey( wallets.new_hotkey(wallet, n_words, use_password, uri, overwrite) ) + def wallet_associate_hotkey( + self, + wallet_name: Optional[str] = Options.wallet_name, + wallet_path: Optional[str] = Options.wallet_path, + wallet_hotkey: Optional[str] = Options.wallet_hotkey_ss58, + network: Optional[list[str]] = Options.network, + prompt: bool = Options.prompt, + quiet: bool = Options.quiet, + verbose: bool = Options.verbose, + ): + """ + Associate a hotkey with a wallet(coldkey). + + USAGE + + This command is used to associate a hotkey with a wallet(coldkey). + + EXAMPLE + + [green]$[/green] btcli wallet associate-hotkey --hotkey-name hotkey_name + [green]$[/green] btcli wallet associate-hotkey --hotkey-ss58 5DkQ4... + """ + self.verbosity_handler(quiet, verbose) + if not wallet_name: + wallet_name = Prompt.ask( + "Enter the [blue]wallet name[/blue] [dim](which you want to associate with the hotkey)[/dim]", + default=self.config.get("wallet_name") or defaults.wallet.name, + ) + if not wallet_hotkey: + wallet_hotkey = Prompt.ask( + "Enter the [blue]hotkey[/blue] name or " + "[blue]hotkey ss58 address[/blue] [dim](to associate with your coldkey)[/dim]" + ) + + hotkey_display = None + if is_valid_ss58_address(wallet_hotkey): + hotkey_ss58 = wallet_hotkey + wallet = self.wallet_ask( + wallet_name, + wallet_path, + None, + ask_for=[WO.NAME, WO.PATH], + validate=WV.WALLET, + ) + hotkey_display = ( + f"hotkey [{COLORS.GENERAL.HK}]{hotkey_ss58}[/{COLORS.GENERAL.HK}]" + ) + else: + wallet = self.wallet_ask( + wallet_name, + wallet_path, + wallet_hotkey, + ask_for=[WO.NAME, WO.PATH, WO.HOTKEY], + validate=WV.WALLET_AND_HOTKEY, + ) + hotkey_ss58 = wallet.hotkey.ss58_address + hotkey_display = f"hotkey [blue]{wallet_hotkey}[/blue] [{COLORS.GENERAL.HK}]({hotkey_ss58})[/{COLORS.GENERAL.HK}]" + + return self._run_command( + wallets.associate_hotkey( + wallet, + self.initialize_chain(network), + hotkey_ss58, + hotkey_display, + prompt, + ) + ) + def wallet_new_coldkey( self, wallet_name: Optional[str] = Options.wallet_name, diff --git a/bittensor_cli/src/commands/wallets.py b/bittensor_cli/src/commands/wallets.py index 61d0bcb6f..333e5e0ff 100644 --- a/bittensor_cli/src/commands/wallets.py +++ b/bittensor_cli/src/commands/wallets.py @@ -50,6 +50,73 @@ ) +async def associate_hotkey( + wallet: Wallet, + subtensor: SubtensorInterface, + hotkey_ss58: str, + hotkey_display: str, + prompt: bool = False, +): + """Associates a hotkey with a wallet""" + + owner_ss58 = await subtensor.get_hotkey_owner(hotkey_ss58) + if owner_ss58: + if owner_ss58 == wallet.coldkeypub.ss58_address: + console.print( + f":white_heavy_check_mark: {hotkey_display.capitalize()} is already " + f"associated with \nwallet [blue]{wallet.name}[/blue], " + f"SS58: [{COLORS.GENERAL.CK}]{owner_ss58}[/{COLORS.GENERAL.CK}]" + ) + return True + else: + owner_wallet = _get_wallet_by_ss58(wallet.path, owner_ss58) + wallet_name = owner_wallet.name if owner_wallet else "unknown wallet" + console.print( + f"[yellow]Warning[/yellow]: {hotkey_display.capitalize()} is already associated with \n" + f"wallet: [blue]{wallet_name}[/blue], SS58: [{COLORS.GENERAL.CK}]{owner_ss58}[/{COLORS.GENERAL.CK}]" + ) + return False + else: + console.print( + f"{hotkey_display.capitalize()} is not associated with any wallet" + ) + + if prompt and not Confirm.ask("Do you want to continue with the association?"): + return False + + if not unlock_key(wallet).success: + return False + + call = await subtensor.substrate.compose_call( + call_module="SubtensorModule", + call_function="try_associate_hotkey", + call_params={ + "hotkey": hotkey_ss58, + }, + ) + + with console.status(":satellite: Associating hotkey on-chain..."): + success, err_msg = await subtensor.sign_and_send_extrinsic( + call, + wallet, + wait_for_inclusion=True, + wait_for_finalization=False, + ) + + if not success: + console.print( + f"[red]:cross_mark: Failed to associate hotkey: {err_msg}[/red]" + ) + return False + + console.print( + f":white_heavy_check_mark: Successfully associated {hotkey_display} with \n" + f"wallet [blue]{wallet.name}[/blue], " + f"SS58: [{COLORS.GENERAL.CK}]{wallet.coldkeypub.ss58_address}[/{COLORS.GENERAL.CK}]" + ) + return True + + async def regen_coldkey( wallet: Wallet, mnemonic: Optional[str], @@ -257,6 +324,15 @@ def get_coldkey_wallets_for_path(path: str) -> list[Wallet]: return wallets +def _get_wallet_by_ss58(path: str, ss58_address: str) -> Optional[Wallet]: + """Find a wallet by its SS58 address in the given path.""" + ss58_addresses, wallet_names = _get_coldkey_ss58_addresses_for_path(path) + for wallet_name, addr in zip(wallet_names, ss58_addresses): + if addr == ss58_address: + return Wallet(path=path, name=wallet_name) + return None + + def _get_coldkey_ss58_addresses_for_path(path: str) -> tuple[list[str], list[str]]: """Get all coldkey ss58 addresses from path.""" @@ -1596,8 +1672,7 @@ async def find_coldkey_swap_extrinsic( """ current_block, genesis_block = await asyncio.gather( - subtensor.substrate.get_block_number(), - subtensor.substrate.get_block_hash(0) + subtensor.substrate.get_block_number(), subtensor.substrate.get_block_hash(0) ) if ( current_block - start_block > 300 diff --git a/tests/e2e_tests/test_wallet_interactions.py b/tests/e2e_tests/test_wallet_interactions.py index 076f03b3e..19cdcc22a 100644 --- a/tests/e2e_tests/test_wallet_interactions.py +++ b/tests/e2e_tests/test_wallet_interactions.py @@ -510,3 +510,112 @@ def test_wallet_identities(local_chain, wallet_setup): assert "Message signed successfully" in sign_using_coldkey.stdout print("โœ… Passed wallet set-id, get-id, sign command") + + +def test_wallet_associate_hotkey(local_chain, wallet_setup): + """ + Test the associating hotkeys and their different cases. + + Steps: + 1. Create wallets for Alice, Bob, and Charlie + 2. Associate a hotkey with Alice's wallet using hotkey name + 3. Verify the association is successful + 4. Try to associate Alice's hotkey with Bob's wallet (should fail) + 5. Try to associate Alice's hotkey again (should show already associated) + 6. Associate Charlie's hotkey with Bob's wallet using SS58 address + + Raises: + AssertionError: If any of the checks or verifications fail + """ + print("Testing wallet associate-hotkey command ๐Ÿงช") + + _, wallet_alice, wallet_path_alice, exec_command_alice = wallet_setup("//Alice") + _, wallet_bob, wallet_path_bob, exec_command_bob = wallet_setup("//Bob") + _, wallet_charlie, _, _ = wallet_setup("//Charlie") + + # Associate Alice's default hotkey with her wallet + result = exec_command_alice( + command="wallet", + sub_command="associate-hotkey", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--no-prompt", + ], + ) + + # Assert successful association + assert "Successfully associated hotkey" in result.stdout + assert wallet_alice.hotkey.ss58_address in result.stdout + assert wallet_alice.coldkeypub.ss58_address in result.stdout + assert wallet_alice.hotkey_str in result.stdout + + # Try to associate Alice's hotkey with Bob's wallet (should fail) + result = exec_command_bob( + command="wallet", + sub_command="associate-hotkey", + extra_args=[ + "--wallet-path", + wallet_path_bob, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_bob.name, + "--hotkey-ss58", + wallet_alice.hotkey.ss58_address, + "--no-prompt", + ], + ) + + assert "Warning" in result.stdout + assert "is already associated with" in result.stdout + + # Try to associate Alice's hotkey again with Alice's wallet + result = exec_command_alice( + command="wallet", + sub_command="associate-hotkey", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--no-prompt", + ], + ) + + assert "is already associated with" in result.stdout + assert "wallet" in result.stdout + assert wallet_alice.name in result.stdout + + # Associate Charlie's hotkey with Bob's wallet using SS58 address + result = exec_command_bob( + command="wallet", + sub_command="associate-hotkey", + extra_args=[ + "--wallet-path", + wallet_path_bob, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_bob.name, + "--hotkey-ss58", + wallet_charlie.hotkey.ss58_address, + "--no-prompt", + ], + ) + + assert "Successfully associated hotkey" in result.stdout + assert wallet_charlie.hotkey.ss58_address in result.stdout + assert wallet_bob.coldkeypub.ss58_address in result.stdout + + print("โœ… Passed wallet associate-hotkey command")