htcli 1.1.0__py3-none-any.whl

src/htcli/utils/wallet/core.py

@@ -0,0 +1,1069 @@
+"""
+Wallet management utilities for the Hypertensor CLI.
+"""
+
+import hashlib
+import json
+import os
+from enum import Enum
+from pathlib import Path
+from typing import Optional
+
+import typer
+from substrateinterface import Keypair
+
+from src.htcli.ui.display import print_success
+
+from ...utils.logging import get_logger
+from .crypto import public_key_to_evm_address
+
+# Use centralized logging framework
+logger = get_logger(__name__)
+
+
+class WalletType(Enum):
+    COLDKEY = 0
+    INDEPENDENT_HOTKEY = 1
+    OWNED_HOTKEY = 2
+
+
+def get_wallet_directory() -> Path:
+    """Get the wallet directory from config or use default."""
+    try:
+        from ...dependencies import get_config
+
+        config = get_config()
+        if config and config.wallet and config.wallet.path:
+            # Expand ~ and other path variables
+            wallet_path = Path(config.wallet.path).expanduser()
+            return wallet_path
+    except Exception:
+        # Fallback to default if config is not available
+        pass
+
+    # Default fallback
+    return Path.home() / ".htcli" / "wallets"
+
+
+def generate_mnemonic() -> str:
+    """
+    Generate a new mnemonic phrase.
+
+    Returns:
+        str: A new mnemonic phrase
+    """
+    return Keypair.generate_mnemonic()
+
+
+def prompt_for_wallet_selection(
+    message: str = "Select wallet",
+    wallet_type: str = "hotkey",
+    auto_password: bool = True,
+) -> str:
+    """
+    Prompt user to select a wallet by name from available wallets.
+
+    Args:
+        message: The prompt message to display
+        wallet_type: Type of wallet to show - "hotkey", "coldkey", or "all"
+
+    Returns:
+        str: The selected wallet name
+
+    Raises:
+        ValueError: If no wallets are available
+    """
+    from ...dependencies import get_client
+    from ...ui.prompts import select_prompt
+
+    try:
+        # Get available wallets from the client
+        client = get_client()
+        response = client.offchain.wallet.list_wallets()
+        wallets = response.get("data", [])
+
+        if not wallets:
+            raise ValueError(
+                "No wallets found. Please create a wallet first using 'htcli wallet create'"
+            )
+
+        # Filter wallets based on type
+        if wallet_type == "hotkey":
+            filtered_wallets = [w for w in wallets if w.get("is_hotkey", True)]
+        elif wallet_type == "coldkey":
+            filtered_wallets = [w for w in wallets if not w.get("is_hotkey", False)]
+        else:  # "all"
+            filtered_wallets = wallets
+
+        if not filtered_wallets:
+            wallet_type_display = (
+                "hotkey"
+                if wallet_type == "hotkey"
+                else "coldkey"
+                if wallet_type == "coldkey"
+                else ""
+            )
+            from ...ui.display import print_error
+
+            print_error(
+                f"No {wallet_type_display} wallets found. Please create a {wallet_type_display} wallet first using 'htcli wallet create --type {wallet_type}'"
+            )
+            raise typer.Abort()
+
+        # Create address to name mapping for owner display
+        address_to_name = {}
+        for wallet in wallets:
+            if not wallet.get("is_hotkey", False):  # This is a coldkey
+                address_to_name[wallet.get("ss58_address")] = wallet.get("name")
+
+        # Extract wallet names and create selection choices
+        wallet_choices = []
+        for wallet in filtered_wallets:
+            name = wallet.get("name", "Unknown")
+            address = wallet.get("ss58_address", "Unknown")
+
+            if wallet_type == "hotkey":
+                # For hotkeys, show owner information
+                owner_address = wallet.get("owner_address", "N/A")
+                if owner_address != "N/A" and owner_address in address_to_name:
+                    owner_display = f"{address_to_name[owner_address]}"
+                    display_text = f"{name} (Owner: {owner_display})"
+                else:
+                    display_text = (
+                        f"{name} (Owner: {owner_address[:8]}...{owner_address[-8:]})"
+                    )
+            else:
+                # For coldkeys or all, show address
+                display_text = f"{name} ({address[:8]}...{address[-8:]})"
+
+            wallet_choices.append((name, display_text, name))
+
+        return select_prompt(message, wallet_choices)
+
+    except Exception as e:
+        from ...ui.display import print_error
+
+        print_error(f"Failed to load wallets: {str(e)}")
+        raise
+
+
+def prompt_for_wallet_and_keypair(
+    message: str = "Select wallet", wallet_type: str = "hotkey"
+) -> tuple[str, "Keypair"]:
+    """
+    Prompt user to select a wallet and return both name and keypair.
+    Handles password prompting automatically for encrypted wallets.
+
+    Args:
+        message: The prompt message to display
+        wallet_type: Type of wallet to show - "hotkey", "coldkey", or "all"
+
+    Returns:
+        tuple: (wallet_name, keypair)
+
+    Raises:
+        typer.Abort: If no wallets are available or user cancels
+    """
+    from ...dependencies import get_client
+    from ...ui.prompts import password_prompt
+
+    # Get wallet name from selection
+    wallet_name = prompt_for_wallet_selection(message, wallet_type, auto_password=False)
+
+    # Get client and load wallet
+    client = get_client()
+    response = client.offchain.wallet.list_wallets()
+    wallets = response.get("data", [])
+
+    # Find the selected wallet
+    selected_wallet = None
+    for wallet in wallets:
+        if wallet.get("name") == wallet_name:
+            selected_wallet = wallet
+            break
+
+    if not selected_wallet:
+        raise ValueError(f"Wallet '{wallet_name}' not found")
+
+    # Extract wallet metadata for disambiguation
+    is_hotkey = selected_wallet.get("is_hotkey", False)
+    owner_address = selected_wallet.get("owner_address") if is_hotkey else None
+
+    # Check if wallet is encrypted and prompt for password if needed
+    if selected_wallet.get("is_encrypted", False):
+        password = password_prompt(f"Enter password for wallet '{wallet_name}'")
+        keypair = client.offchain.wallet.get_keypair(
+            wallet_name, password, is_hotkey=is_hotkey, owner_address=owner_address
+        )
+    else:
+        # For unencrypted wallets, get keypair directly
+        keypair = client.offchain.wallet.get_keypair(
+            wallet_name, None, is_hotkey=is_hotkey, owner_address=owner_address
+        )
+
+    return wallet_name, keypair
+
+
+def retrieve_wallet_with_validation(
+    wallet_type: str = "coldkey", 
+    purpose: str = "sign the transaction",
+    only_existing_wallets: bool = False,
+) -> tuple[str, "Keypair"]:
+    """
+    Comprehensive wallet retrieval mechanism that handles:
+    - Address vs wallet name selection
+    - Wallet type filtering (hotkey/coldkey)
+    - Password handling with retry attempts
+    - Option to create new wallet if needed
+
+    Args:
+        wallet_type: Type of wallet to retrieve - "hotkey" or "coldkey"
+        purpose: Description of what the wallet will be used for
+        only_existing_wallets: If True, skip initial selection and only allow selecting from existing wallets.
+                             This is useful for owner-only operations that require an existing wallet.
+
+    Returns:
+        tuple: (wallet_name_or_address, keypair)
+
+    Raises:
+        typer.Abort: If user cancels or max retries exceeded
+    """
+    from ...dependencies import get_client
+    from ...ui.display import print_error
+    from ...ui.prompts import confirm_prompt, select_prompt, text_prompt
+
+    try:
+        client = get_client()
+    except Exception as e:
+        print_error(f"Failed to initialize client: {str(e)}")
+        print_error("Please check your configuration and try again.")
+        raise typer.Abort() from e
+
+    # If only_existing_wallets is True, skip the initial selection and go directly to wallet selection
+    if only_existing_wallets:
+        return _select_existing_wallet(client, wallet_type, purpose, allow_create_new=False)
+
+    # Step 1: Ask if user wants to provide address or wallet name
+    input_type = select_prompt(
+        f"How do you want to provide the {wallet_type} for {purpose}?",
+        [
+            ("address", f"Provide {wallet_type} address directly"),
+            ("wallet", f"Select from existing {wallet_type} wallets"),
+            ("create", f"Create new {wallet_type} wallet"),
+        ],
+    )
+
+    if input_type == "address":
+        # Step 2a: Get address directly
+        address = text_prompt(f"Enter {wallet_type} address")
+
+        # Validate address format
+        if wallet_type == "hotkey":
+            if not client.offchain.wallet.validate_ethereum_address(address):
+                print_error(
+                    f"Invalid {wallet_type} address format. Expected Bytes20 format (0x + 40 hex characters)"
+                )
+                raise typer.Abort()
+        else:  # coldkey
+            # For coldkey, we need to find a wallet with this address
+            from ...utils.wallet.crypto import list_keys
+
+            wallets = list_keys()
+
+            matching_wallet = None
+            for wallet in wallets:
+                # Check both ss58_address and derived address from public_key
+                if not wallet.get("is_hotkey", False):
+                    # Match by ss58_address (primary)
+                    if wallet.get("ss58_address") == address:
+                        matching_wallet = wallet
+                        break
+                    # Also match by public_key (keypair derives 0x + public_key as address)
+                    if wallet.get("public_key"):
+                        derived_addr = "0x" + wallet.get("public_key")
+                        if derived_addr.lower() == address.lower():
+                            matching_wallet = wallet
+                            break
+
+            if not matching_wallet:
+                print_error(f"No {wallet_type} wallet found with address: {address}")
+                if confirm_prompt("Would you like to create a new wallet instead?"):
+                    return _create_new_wallet(client, wallet_type)
+                raise typer.Abort()
+
+            # Get keypair for the matching wallet
+            return _get_keypair_with_password_retry(
+                client, matching_wallet["name"], matching_wallet, 3
+            )
+
+        # For hotkey addresses, we don't need a keypair (just the address)
+        return address, None
+
+    elif input_type == "wallet":
+        # Step 2b: Select from existing wallets
+        return _select_existing_wallet(client, wallet_type, purpose, allow_create_new=True)
+
+    else:  # create
+        # Step 2c: Create new wallet
+        return _create_new_wallet(client, wallet_type)
+
+
+def _select_existing_wallet(
+    client, wallet_type: str, purpose: str, allow_create_new: bool = True
+) -> tuple[str, "Keypair"]:
+    """Select from existing wallets of specified type.
+    
+    Args:
+        client: Client instance
+        wallet_type: Type of wallet - "hotkey" or "coldkey"
+        purpose: Description of what the wallet will be used for
+        allow_create_new: If False, do not show option to create new wallet
+    """
+    from ...ui.display import print_error
+    from ...ui.prompts import confirm_prompt, select_prompt
+    from ...utils.wallet.crypto import list_keys
+
+    try:
+        # Get available wallets with proper structure including is_hotkey
+        wallets = list_keys()
+
+        # Ensure wallets is a list and not None
+        if wallets is None:
+            print_error("Failed to retrieve wallet information.")
+            if allow_create_new and confirm_prompt(f"Would you like to create a new {wallet_type} wallet?"):
+                return _create_new_wallet(client, wallet_type)
+            raise typer.Abort()
+
+        # Filter by wallet type
+        if wallet_type == "hotkey":
+            filtered_wallets = [w for w in wallets if w.get("is_hotkey", False)]
+        else:  # coldkey
+            filtered_wallets = [w for w in wallets if not w.get("is_hotkey", False)]
+
+        # Check if we have any wallets of the requested type
+        if not filtered_wallets:
+            print_error(f"No {wallet_type} wallets found.")
+            if allow_create_new and confirm_prompt(f"Would you like to create a new {wallet_type} wallet?"):
+                return _create_new_wallet(client, wallet_type)
+            raise typer.Abort()
+
+        # Create selection choices
+        wallet_choices = []
+        
+        for wallet in filtered_wallets:
+            name = wallet.get("name", "Unknown")
+            address = wallet.get("ss58_address", "Unknown")
+
+            if wallet_type == "hotkey":
+                # For hotkeys, always show owner information to disambiguate
+                owner_address = wallet.get("owner_address")
+                if owner_address:
+                    # Find owner wallet name
+                    owner_display = None
+                    for w in wallets:
+                        if (
+                            w.get("ss58_address") == owner_address
+                            or w.get("evm_address", "").lower() == owner_address.lower()
+                        ):
+                            owner_display = w.get("name", "Unknown")
+                            break
+                    
+                    if owner_display:
+                        display_text = f"{name} (Owner: {owner_display})"
+                    else:
+                        display_text = f"{name} (Owner: {owner_address[:8]}...{owner_address[-8:]})"
+                else:
+                    display_text = f"{name} (No owner)"
+            else:
+                # For coldkeys, show address
+                display_text = f"{name} ({address[:8]}...{address[-8:]})"
+
+            # Create unique key for this wallet (for hotkeys, include owner_address to disambiguate)
+            if wallet_type == "hotkey":
+                owner_addr = wallet.get("owner_address") or ""
+                wallet_key = f"{name}__owner_{owner_addr}"
+            else:
+                wallet_key = f"{name}__address_{address}"
+
+            # Store the full wallet dict as the value so we can use it directly
+            wallet_choices.append((wallet_key, display_text, wallet))
+
+        # Add option to create new wallet only if allowed
+        if allow_create_new:
+            wallet_choices.append(
+                ("CREATE_NEW", f"Create new {wallet_type} wallet", "CREATE_NEW")
+            )
+
+        prompt_message = f"Select {wallet_type} wallet for {purpose}"
+        selected_wallet = select_prompt(prompt_message, wallet_choices)
+
+        if selected_wallet == "CREATE_NEW":
+            return _create_new_wallet(client, wallet_type)
+
+        # The selected_wallet should be the wallet dict (stored as value in choices tuple)
+        if not isinstance(selected_wallet, dict):
+            # Fallback: if for some reason we got a string key instead of dict
+            # Find the wallet from filtered_wallets by matching the key
+            selected_key = str(selected_wallet)
+            selected_wallet = None
+            for wallet in filtered_wallets:
+                name = wallet.get("name", "Unknown")
+                if wallet_type == "hotkey":
+                    owner_addr = wallet.get("owner_address") or ""
+                    wallet_key = f"{name}__owner_{owner_addr}"
+                else:
+                    address = wallet.get("ss58_address", "Unknown")
+                    wallet_key = f"{name}__address_{address}"
+                
+                if wallet_key == selected_key:
+                    selected_wallet = wallet
+                    break
+
+        if not selected_wallet or not isinstance(selected_wallet, dict):
+            raise ValueError(f"Selected wallet not found or invalid")
+
+        # Get wallet name from the selected wallet dict
+        selected_name = selected_wallet.get("name")
+
+        # Get keypair with password retry
+        # selected_wallet already has is_hotkey and owner_address from the wallet data
+        # _get_keypair_with_password_retry will extract these and pass them to get_keypair
+        return _get_keypair_with_password_retry(
+            client, selected_name, selected_wallet, 3
+        )
+
+    except Exception as e:
+        print_error(f"Error retrieving wallets: {str(e)}")
+        # Propagate actual error upward so callers can surface context
+        raise
+
+
+def _get_keypair_with_password_retry(
+    client, wallet_name: str, wallet_info: dict, max_attempts: int = 3
+) -> tuple[str, "Keypair"]:
+    """Get keypair with password retry mechanism."""
+    from ...ui.display import print_error
+    from ...ui.prompts import password_prompt
+
+    # Extract wallet type information for disambiguation
+    is_hotkey = wallet_info.get("is_hotkey", False)
+    owner_address = wallet_info.get("owner_address")
+
+    # Check if wallet is encrypted
+    if not wallet_info.get("is_encrypted", False):
+        # Unencrypted wallet - no password needed
+        try:
+            keypair = client.offchain.wallet.get_keypair(
+                wallet_name, None, is_hotkey=is_hotkey, owner_address=owner_address
+            )
+            return wallet_name, keypair
+        except Exception as e:
+            raise ValueError(f"Failed to unlock wallet '{wallet_name}': {str(e)}") from e
+
+    # Encrypted wallet - need password
+    for attempt in range(max_attempts):
+        try:
+            remaining = max_attempts - attempt
+            password = password_prompt(
+                f"Enter password for wallet '{wallet_name}' (attempt {attempt + 1}/{max_attempts})"
+            )
+            keypair = client.offchain.wallet.get_keypair(
+                wallet_name, password, is_hotkey=is_hotkey, owner_address=owner_address
+            )
+            print_success(f"Successfully unlocked wallet '{wallet_name}'")
+            return wallet_name, keypair
+        except ValueError as e:
+            if "Invalid password" in str(e):
+                if attempt < max_attempts - 1:
+                    print_error(
+                        f"Invalid password. {remaining - 1} attempts remaining."
+                    )
+                else:
+                    print_error("Invalid password. Maximum attempts exceeded.")
+                    raise
+            else:
+                # Don't print error here - let the caller handle it
+                raise
+        except Exception as e:
+            error_message = str(e)
+            if "Invalid password" in error_message:
+                if attempt < max_attempts - 1:
+                    print_error(
+                        f"Invalid password. {remaining - 1} attempts remaining."
+                    )
+                    continue
+                print_error("Invalid password. Maximum attempts exceeded.")
+                raise ValueError(
+                    f"Failed to unlock wallet '{wallet_name}': Invalid password"
+                ) from e
+            raise ValueError(
+                f"Failed to unlock wallet '{wallet_name}': {error_message}"
+            ) from e
+
+
+def resolve_coldkey_and_get_keypair(
+    coldkey_input: str,
+) -> tuple[str, "Keypair"]:
+    """
+    Resolve coldkey input (address or wallet name) and get the keypair.
+
+    Args:
+        coldkey_input: Either a coldkey address (ss58 or 0x format) or wallet name
+
+    Returns:
+        tuple: (wallet_name, keypair)
+
+    Raises:
+        ValueError: If coldkey cannot be resolved or wallet not found
+    """
+    from ...dependencies import get_client
+    from ...utils.wallet.crypto import get_wallet_info_by_name, list_keys
+
+    coldkey_input = coldkey_input.strip()
+    client = get_client()
+
+    # First, try to find by wallet name
+    try:
+        wallet_info = get_wallet_info_by_name(coldkey_input, is_hotkey=False)
+        # If found and it's a coldkey (not hotkey)
+        if not wallet_info.get("is_hotkey", False):
+            # Get keypair for this wallet
+            return _get_keypair_with_password_retry(
+                client, coldkey_input, wallet_info, 3
+            )
+    except FileNotFoundError:
+        # Not a wallet name, try as address
+        pass
+    except Exception as e:
+        raise ValueError(f"Failed to resolve coldkey '{coldkey_input}': {str(e)}") from e
+
+    # If not found by name, try to find by address
+    wallets = list_keys()
+    matching_wallet = None
+
+    for wallet in wallets:
+        if wallet.get("is_hotkey", False):
+            continue  # Skip hotkeys
+
+        # Match by ss58_address (primary)
+        if wallet.get("ss58_address") and wallet.get("ss58_address") == coldkey_input:
+            matching_wallet = wallet
+            break
+
+        # Match by evm_address
+        if wallet.get("evm_address") and wallet.get("evm_address").lower() == coldkey_input.lower():
+            matching_wallet = wallet
+            break
+
+        # Match by address field
+        if wallet.get("address") and wallet.get("address").lower() == coldkey_input.lower():
+            matching_wallet = wallet
+            break
+
+        # Also match by public_key (keypair derives 0x + public_key as address)
+        if wallet.get("public_key"):
+            derived_addr = "0x" + wallet.get("public_key")
+            if derived_addr.lower() == coldkey_input.lower():
+                matching_wallet = wallet
+                break
+
+    if not matching_wallet:
+        raise ValueError(
+            f"No coldkey wallet found with address or name: {coldkey_input}"
+        )
+
+    # Get keypair for the matching wallet
+    return _get_keypair_with_password_retry(
+        client, matching_wallet["name"], matching_wallet, 3
+    )
+
+
+def resolve_coldkey_address(coldkey_input: str) -> str:
+    """
+    Resolve a coldkey wallet input (name or address) to a usable address string.
+
+    This is a lightweight resolver that does NOT unlock the wallet. 
+    For ECDSA keys, always returns EVM address format.
+
+    Args:
+        coldkey_input: Wallet name or address.
+
+    Returns:
+        Resolved EVM address string (0x... format).
+
+    Raises:
+        ValueError: If the input is empty or cannot be resolved.
+    """
+    from ...utils.wallet.crypto import get_wallet_info_by_name, list_keys
+
+    if coldkey_input is None:
+        raise ValueError("Coldkey value cannot be empty")
+
+    value = coldkey_input.strip()
+    if not value:
+        raise ValueError("Coldkey value cannot be empty")
+
+    # Direct address passthrough (EVM format)
+    if value.startswith("0x") and len(value) == 42:
+        return value.lower()
+    
+    # FIX: Don't accept SS58 addresses - only EVM format
+    # if len(value) in (48, 49) and value.startswith("5"):
+    #     return value
+
+    # Try resolving by wallet name (coldkey only)
+    try:
+        wallet_info = get_wallet_info_by_name(value, is_hotkey=False)
+        if wallet_info and not wallet_info.get("is_hotkey", False):
+            # FIX: Always use evm_address for ECDSA keys
+            key_type = wallet_info.get("key_type", "").lower()
+            if key_type == "ecdsa":
+                resolved = wallet_info.get("evm_address") or wallet_info.get("address")
+            else:
+                resolved = (
+                    wallet_info.get("evm_address")
+                    or wallet_info.get("address")
+                    or wallet_info.get("ss58_address")
+                )
+            if resolved:
+                return resolved
+    except FileNotFoundError:
+        pass
+    except Exception as e:
+        raise ValueError(f"Failed to resolve coldkey '{value}': {str(e)}") from e
+
+    # Fallback: scan known wallets for matching name or address
+    wallets = list_keys()
+    for wallet in wallets:
+        if wallet.get("is_hotkey", False):
+            continue
+        
+        # FIX: Match by evm_address first (for ECDSA)
+        if wallet.get("evm_address") and wallet.get("evm_address").lower() == value.lower():
+            return wallet.get("evm_address")
+        
+        # Match by address field (should be EVM for ECDSA)
+        if wallet.get("address") and wallet.get("address").lower() == value.lower():
+            return wallet.get("address")
+        
+        # Match by name
+        if wallet.get("name") == value:
+            key_type = wallet.get("key_type", "").lower()
+            if key_type == "ecdsa":
+                return wallet.get("evm_address") or wallet.get("address")
+            else:
+                return (
+                    wallet.get("evm_address")
+                    or wallet.get("address")
+                    or wallet.get("ss58_address")
+                )
+
+    raise ValueError(f"No coldkey wallet found with address or name: {value}")
+
+
+def _create_new_wallet(client, wallet_type: str) -> tuple[str, "Keypair"]:
+    """Create a new wallet of specified type."""
+    from ...ui.display import print_error, print_info, print_success
+    from ...ui.prompts import confirm_prompt, password_prompt, text_prompt
+
+    print_info(f"Creating new {wallet_type} wallet...")
+
+    # Get wallet name
+    wallet_name = text_prompt(f"Enter name for new {wallet_type} wallet")
+
+    # Ask for encryption
+    encrypt = confirm_prompt(
+        f"Encrypt the {wallet_type} wallet with a password?", default=True
+    )
+
+    password = None
+    if encrypt:
+        password = password_prompt(
+            f"Enter password for new {wallet_type} wallet", confirm=True
+        )
+
+    try:
+        if wallet_type == "hotkey":
+            # For hotkeys, we need an owner (coldkey)
+            print_info("Hotkeys must have an owner (coldkey wallet).")
+
+            # Get available coldkeys
+            try:
+                wallets_response = client.offchain.wallet.list_wallets()
+                all_wallets = wallets_response.get("data", [])
+
+                # Filter for coldkeys only
+                coldkey_wallets = [
+                    w for w in all_wallets if not w.get("is_hotkey", False)
+                ]
+
+                if coldkey_wallets:
+                    # Create selection choices
+                    coldkey_choices = []
+                    for _i, wallet in enumerate(coldkey_wallets):
+                        name = wallet.get("name", "Unknown")
+                        address = wallet.get("ss58_address", "Unknown")
+                        display_text = f"{name} ({address[:8]}...{address[-8:]})"
+                        coldkey_choices.append((name, display_text, name))
+
+                    # Add option to enter address manually
+                    coldkey_choices.append(
+                        ("MANUAL", "Enter coldkey address manually", "MANUAL")
+                    )
+
+                    # Let user select from available coldkeys
+                    from ...ui.prompts import select_prompt
+
+                    selected_owner = select_prompt(
+                        "Select owner coldkey wallet", coldkey_choices
+                    )
+
+                    if selected_owner == "MANUAL":
+                        # Fallback to manual entry
+                        owner_input = text_prompt("Enter coldkey wallet address")
+                        owner_address = owner_input
+                        print_info(f"Using provided address as owner: {owner_address}")
+                    else:
+                        # Find the selected coldkey
+                        selected_wallet = None
+                        for wallet in coldkey_wallets:
+                            if wallet.get("name") == selected_owner:
+                                selected_wallet = wallet
+                                break
+
+                        if selected_wallet:
+                            owner_address = selected_wallet["ss58_address"]
+                            print_info(
+                                f"Using coldkey wallet '{selected_owner}' as owner: {owner_address}"
+                            )
+                        else:
+                            raise ValueError(
+                                f"Selected coldkey '{selected_owner}' not found"
+                            )
+                else:
+                    # No coldkeys available, ask for manual entry
+                    print_info(
+                        "No coldkey wallets found. Please provide a coldkey address manually."
+                    )
+                    owner_input = text_prompt("Enter owner coldkey wallet address")
+                    owner_address = owner_input
+                    print_info(f"Using provided address as owner: {owner_address}")
+
+            except Exception as e:
+                # Fallback to manual entry if listing fails
+                print_info(f"Could not list coldkey wallets: {str(e)}")
+                owner_input = text_prompt("Enter owner coldkey wallet address")
+                owner_address = owner_input
+                print_info(f"Using provided address as owner: {owner_address}")
+
+            # Create hotkey wallet with owner
+            result = client.offchain.wallet.create_hotkey_wallet(
+                wallet_name, owner_address, password
+            )
+        else:
+            # Create coldkey wallet
+            result = client.offchain.wallet.create_coldkey_wallet(wallet_name, password)
+
+        if result.get("success"):
+            print_success(f"Successfully created {wallet_type} wallet '{wallet_name}'")
+            # Get the keypair for the newly created wallet
+            # For hotkeys, pass is_hotkey=True and owner_address to help locate the wallet
+            if wallet_type == "hotkey":
+                keypair = client.offchain.wallet.get_keypair(
+                    wallet_name, password, is_hotkey=True, owner_address=owner_address
+                )
+            else:
+                keypair = client.offchain.wallet.get_keypair(wallet_name, password)
+            return wallet_name, keypair
+        else:
+            raise Exception(result.get("message", "Failed to create wallet"))
+
+    except Exception as e:
+        print_error(f"Failed to create {wallet_type} wallet: {str(e)}")
+        raise
+
+
+def create_keypair_from_mnemonic(mnemonic: str) -> Keypair:
+    """
+    Create a keypair from a mnemonic phrase.
+
+    Args:
+        mnemonic: The mnemonic phrase to create the keypair from
+
+    Returns:
+        Keypair: The generated keypair
+
+    Raises:
+        RuntimeError: If the mnemonic is invalid
+    """
+    try:
+        # Match mesh-template pattern - no ss58_format for ECDSA
+        keypair = Keypair.create_from_mnemonic(mnemonic, crypto_type=2)
+        logger.info("Created keypair from mnemonic:")
+
+        return keypair
+    except Exception as e:
+        raise RuntimeError(f"Failed to create keypair from mnemonic: {e}") from e
+
+
+def get_account_id(public_key: bytes) -> str:
+    """
+    Generate the account ID from a public key.
+    This is a 32-byte blake2b hash of the public key.
+
+    Args:
+        public_key: The public key bytes
+
+    Returns:
+        str: The account ID in hex format with 0x prefix
+    """
+    # Use blake2b hash function with 32-byte output
+    blake2b = hashlib.blake2b(digest_size=32)
+    blake2b.update(public_key)
+    return "0x" + blake2b.hexdigest()
+
+
+def encrypt_data(data: bytes, password: str) -> bytes:
+    """
+    Encrypt data using XOR with the password and add a password hash for validation.
+
+    Args:
+        data: The data to encrypt
+        password: The password to use for encryption
+
+    Returns:
+        bytes: The encrypted data with password hash
+    """
+    # For empty password, return data as is (unencrypted)
+    if not password:
+        return data
+
+    # Generate password hash for validation
+    password_hash = hashlib.sha256(password.encode()).digest()
+
+    # Encrypt the data
+    key_bytes = password.encode("utf-8")
+    encrypted = bytes(data[i] ^ key_bytes[i % len(key_bytes)] for i in range(len(data)))
+
+    # Combine password hash and encrypted data
+    return password_hash + encrypted
+
+
+def decrypt_data(encrypted_data: bytes, password: str) -> bytes:
+    """
+    Decrypt data that was encrypted with XOR using the password.
+    Also validates the password using the stored hash.
+
+    Args:
+        encrypted_data: The encrypted data with password hash
+        password: The password used for encryption
+
+    Returns:
+        bytes: The decrypted data
+
+    Raises:
+        RuntimeError: If the password is incorrect
+    """
+    if not password:
+        return encrypted_data
+
+    # Extract password hash and encrypted data
+    stored_hash = encrypted_data[:32]  # First 32 bytes are the hash
+    encrypted = encrypted_data[32:]  # Rest is the encrypted data
+
+    # Verify password
+    password_hash = hashlib.sha256(password.encode()).digest()
+    if password_hash != stored_hash:
+        raise RuntimeError("Invalid password")
+
+    # Decrypt the data
+    key_bytes = password.encode("utf-8")
+    return bytes(
+        encrypted[i] ^ key_bytes[i % len(key_bytes)] for i in range(len(encrypted))
+    )
+
+
+def create_wallet(
+    name: str,  # This will be the file name
+    wallet_dir: Optional[Path] = None,  # Accept the full target directory path
+    is_hotkey: bool = False,  # True for hotkey, False for coldkey
+    password: Optional[str] = None,  # Password for encryption
+    owner_address: Optional[str] = None,  # Required for owned hotkeys
+    mnemonic: Optional[str] = None,  # Optional mnemonic for regeneration,
+    force: bool = False,  # Skip confirmation prompt or overwrite existing wallet
+) -> tuple[str, str, str]:  # Return main wallet file path, ss58_address, and mnemonic
+    """
+    Create a Hypertensor-compatible wallet in the specified directory.
+    Generates a new keypair and saves it as a JSON file with encrypted private key.
+
+    Args:
+        name: The desired file name.
+        wallet_dir: The full Path object for the directory where the wallet files should be created.
+        is_hotkey: True for hotkey, False for coldkey.
+        password: Optional password for encryption.
+        owner_address: Required for owned hotkeys, specifies the coldkey address that owns this hotkey.
+        mnemonic: Optional mnemonic phrase for regeneration. If provided, uses this instead of generating a new one.
+
+    Returns:
+        tuple of (main_wallet_file_path, ss58_address, mnemonic)
+    """
+    if not wallet_dir:
+        wallet_dir = get_wallet_directory()
+
+    # Create directory if it doesn't exist
+    wallet_dir.mkdir(parents=True, exist_ok=True)
+
+    main_wallet_file_path = wallet_dir / f"{name}.json"
+
+    # Check if wallet file already exists
+    if main_wallet_file_path.exists():
+        if not force:
+            raise ValueError(
+                f"Wallet file '{name}' already exists in {wallet_dir}. Use --force to overwrite existing file or choose another name."
+            )
+        else:
+            # Remove existing file
+            logger.info(
+                f"Wallet file '{name}' already exists in {wallet_dir}. Overwriting existing file."
+            )
+            logger.info(f"Removing existing file '{name}'...")
+            main_wallet_file_path.unlink()
+            logger.info(f"Removed existing file '{name}'.")
+
+    # Validate owner_address for owned hotkeys
+    if is_hotkey and owner_address:
+        if not (owner_address.startswith("5") or owner_address.startswith("0x")):
+            raise ValueError(
+                "Owner address must be a valid SS58 address starting with '5' or EVM address starting with '0x'"
+            )
+
+    # Generate or use provided mnemonic
+    if not mnemonic:
+        mnemonic = Keypair.generate_mnemonic()
+
+    # Create keypair from mnemonic
+    keypair = create_keypair_from_mnemonic(mnemonic)
+    evm_address = public_key_to_evm_address(keypair.public_key)
+
+    # Prepare wallet data
+    wallet_data = {
+        "name": name,
+        "key_type": "ecdsa",  # Default to ecdsa for EVM compatibility
+        "public_key": keypair.public_key.hex(),
+        "ss58_address": keypair.ss58_address,
+        "evm_address": evm_address,
+        "is_hotkey": is_hotkey,
+        "owner_address": owner_address,
+        "is_encrypted": password is not None,
+    }
+
+    # Store private key
+    if password:
+        encrypted_private_key = encrypt_data(keypair.private_key, password)
+        wallet_data["private_key"] = "0x" + encrypted_private_key.hex()
+    else:
+        wallet_data["private_key"] = "0x" + keypair.private_key.hex()
+
+    try:
+        with open(main_wallet_file_path, "w") as f:
+            json.dump(wallet_data, f, indent=2)
+        os.chmod(main_wallet_file_path, 0o600)  # Secure file permissions
+    except Exception as e:
+        if main_wallet_file_path.exists():
+            os.remove(main_wallet_file_path)
+        raise RuntimeError(f"Failed to save wallet file: {e}") from e
+
+    return str(main_wallet_file_path), evm_address, mnemonic
+
+
+def import_wallet(
+    name: str, wallet_dir: Optional[Path] = None, password: Optional[str] = None
+) -> Keypair:
+    """
+    Import a wallet and return its keypair.
+
+    Args:
+        name (str): Name of the wallet
+        wallet_dir (Path): Directory containing the wallet file
+        password (str, optional): Password for encrypted wallets
+
+    Returns:
+        Keypair: The wallet's keypair
+
+    Raises:
+        ValueError: If wallet file is not found or invalid
+        RuntimeError: If password is incorrect or wallet is corrupted
+    """
+    if not wallet_dir:
+        wallet_dir = get_wallet_directory()
+
+    wallet_path = wallet_dir / f"{name}.json"
+    if not wallet_path.exists():
+        raise ValueError(f"Wallet file not found at {wallet_path}")
+
+    try:
+        with open(wallet_path) as f:
+            wallet_data = json.load(f)
+    except json.JSONDecodeError as j:
+        raise ValueError("Invalid wallet file: not a valid JSON file") from j
+
+    # Get private key
+    private_key_hex = wallet_data.get("private_key", "").replace("0x", "")
+    if not private_key_hex:
+        raise ValueError("Invalid wallet file: missing private key")
+
+    try:
+        private_key_bytes = bytes.fromhex(private_key_hex)
+    except ValueError as v:
+        raise ValueError("Invalid wallet file: private key is not valid hex") from v
+
+    # Check if wallet is encrypted
+    is_encrypted = wallet_data.get("is_encrypted", False)
+
+    # For unencrypted wallets
+    if not is_encrypted:
+        if password:
+            raise RuntimeError("Invalid password: This wallet is not encrypted")
+        try:
+            # Match mesh-template pattern - no ss58_format for ECDSA
+            return Keypair.create_from_private_key(private_key_bytes, crypto_type=2)
+        except Exception as e:
+            raise RuntimeError(f"Failed to create keypair: {str(e)}") from e
+
+    # For encrypted wallets
+    if not password:
+        raise RuntimeError(
+            "Invalid password: Wallet is encrypted but no password provided"
+        )
+
+    try:
+        private_key_bytes = decrypt_data(private_key_bytes, password)
+        if not is_valid_private_key(private_key_bytes):
+            raise RuntimeError("Invalid password: Failed to decrypt private key")
+    except Exception as e:
+        raise RuntimeError("Invalid password: Failed to decrypt private key") from e
+
+    try:
+        # Match mesh-template pattern - no ss58_format for ECDSA
+        return Keypair.create_from_private_key(private_key_bytes, crypto_type=2)
+    except Exception as e:
+        raise RuntimeError(f"Failed to create keypair: {str(e)}") from e
+
+
+def is_valid_private_key(private_key_bytes: bytes) -> bool:
+    """
+    Validate if the decrypted private key is valid.
+
+    Args:
+        private_key_bytes (bytes): The decrypted private key bytes
+
+    Returns:
+        bool: True if the private key is valid, False otherwise
+    """
+    try:
+        # Try to create a keypair with the private key
+        # Match mesh-template pattern - no ss58_format for ECDSA
+        Keypair.create_from_private_key(private_key_bytes, crypto_type=2)
+        return True
+    except Exception:
+        return False