communex.cli.key

  1import re
  2from enum import Enum
  3from typing import Any, Optional, cast
  4
  5import typer
  6from substrateinterface import Keypair
  7from typeguard import check_type
  8from typer import Context
  9
 10import communex.compat.key as comx_key
 11from communex._common import BalanceUnit, format_balance
 12from communex.cli._common import (
 13    make_custom_context,
 14    print_table_from_plain_dict,
 15    print_table_standardize,
 16)
 17from communex.compat.key import (
 18    classic_store_key,
 19    local_key_addresses,
 20)
 21from communex.key import check_ss58_address, generate_keypair, is_ss58_address
 22from communex.misc import (
 23    local_keys_allbalance,
 24    local_keys_to_freebalance,
 25    local_keys_to_stakedbalance,
 26)
 27
 28key_app = typer.Typer(no_args_is_help=True)
 29
 30
 31class SortBalance(str, Enum):
 32    all = "all"
 33    free = "free"
 34    staked = "staked"
 35
 36
 37@key_app.command()
 38def create(ctx: Context, name: str, password: str = typer.Option(None)):
 39    """
 40    Generates a new key and stores it on a disk with the given name.
 41    """
 42    context = make_custom_context(ctx)
 43
 44    keypair = generate_keypair()
 45    address = keypair.ss58_address
 46
 47    context.info(f"Generated key with public address '{address}'.")
 48
 49    classic_store_key(keypair, name, password)
 50
 51    context.info(f"Key successfully stored with name '{name}'.")
 52
 53
 54@key_app.command()
 55def regen(
 56    ctx: Context, name: str, key_input: str, password: Optional[str] = None
 57):
 58    """
 59    Stores the given key on a disk. Works with private key or mnemonic.
 60    """
 61    context = make_custom_context(ctx)
 62    # TODO: secret input from env var and stdin
 63
 64    # Determine the input type based on the presence of spaces.
 65    if re.search(r"\s", key_input):
 66        # If mnemonic (contains spaces between words).
 67        keypair = Keypair.create_from_mnemonic(key_input)
 68        key_type = "mnemonic"
 69    else:
 70        # If private key (assumes no spaces).
 71        keypair = Keypair.create_from_private_key(key_input, ss58_format=42)
 72        key_type = "private key"
 73        # Substrate does not return these.
 74        keypair.mnemonic = ""  # type: ignore
 75        keypair.seed_hex = ""
 76
 77    address = keypair.ss58_address
 78    context.info(f"Loaded {key_type} with public address `{address}`.")
 79
 80    classic_store_key(keypair, name, password)
 81
 82    context.info(f"Key stored with name `{name}` successfully.")
 83
 84
 85@key_app.command()
 86def show(
 87    ctx: Context,
 88    key: str,
 89    show_private: bool = False,
 90    password: Optional[str] = None,
 91):
 92    """
 93    Show information about a key.
 94    """
 95    context = make_custom_context(ctx)
 96
 97    keypair = context.load_key(key, password)
 98    key_dict = comx_key.to_classic_dict(keypair, path=key)
 99
100    if show_private is not True:
101        key_dict["private_key"] = "[SENSITIVE-MODE]"
102        key_dict["seed_hex"] = "[SENSITIVE-MODE]"
103        key_dict["mnemonic"] = "[SENSITIVE-MODE]"
104
105    key_dict = check_type(key_dict, dict[str, Any])
106
107    print_table_from_plain_dict(key_dict, ["Key", "Value"], context.console)
108
109
110@key_app.command()
111def balances(
112    ctx: Context,
113    unit: BalanceUnit = BalanceUnit.joule,
114    sort_balance: SortBalance = SortBalance.all,
115):
116    """
117    Gets balances of all keys.
118    """
119    context = make_custom_context(ctx)
120    client = context.com_client()
121
122    local_keys = local_key_addresses(context.password_manager)
123    with context.console.status(
124        "Getting balances of all keys, this might take a while..."
125    ):
126        key2freebalance, key2stake = local_keys_allbalance(client, local_keys)
127    key_to_freebalance = {
128        k: format_balance(v, unit) for k, v in key2freebalance.items()
129    }
130    key_to_stake = {k: format_balance(v, unit) for k, v in key2stake.items()}
131
132    key2balance = {k: v + key2stake[k] for k, v in key2freebalance.items()}
133    key_to_balance = {
134        k: format_balance(v, unit) for k, v in key2balance.items()
135    }
136
137    if sort_balance == SortBalance.all:
138        sorted_bal = {
139            k: v
140            for k, v in sorted(
141                key2balance.items(), key=lambda item: item[1], reverse=True
142            )
143        }
144    elif sort_balance == SortBalance.free:
145        sorted_bal = {
146            k: v
147            for k, v in sorted(
148                key2freebalance.items(), key=lambda item: item[1], reverse=True
149            )
150        }
151    elif sort_balance == SortBalance.staked:
152        sorted_bal = {
153            k: v
154            for k, v in sorted(
155                key2stake.items(), key=lambda item: item[1], reverse=True
156            )
157        }
158    else:
159        raise ValueError("Invalid sort balance option")
160
161    stake: list[str] = []
162    all_balance: list[str] = []
163    free: list[str] = []
164    keys: list[str] = []
165
166    for key, _ in sorted_bal.items():
167        keys.append(key)
168        free.append(key_to_freebalance[key])
169        stake.append(key_to_stake[key])
170        all_balance.append(key_to_balance[key])
171
172    pretty_dict = {
173        "key": keys,
174        "free": free,
175        "staked": stake,
176        "all": all_balance,
177    }
178
179    general_dict: dict[str, list[Any]] = cast(dict[str, list[Any]], pretty_dict)
180    print_table_standardize(general_dict, context.console)
181
182
183@key_app.command(name="list")
184def inventory(
185    ctx: Context,
186):
187    """
188    Lists all keys stored on disk.
189    """
190    context = make_custom_context(ctx)
191
192    key_to_address = local_key_addresses(context.password_manager)
193    general_key_to_address: dict[str, str] = cast(
194        dict[str, str], key_to_address
195    )
196    print_table_from_plain_dict(
197        general_key_to_address, ["Key", "Address"], context.console
198    )
199
200
201@key_app.command()
202def stakefrom(
203    ctx: Context,
204    key: str,
205    unit: BalanceUnit = BalanceUnit.joule,
206    password: Optional[str] = None,
207):
208    """
209    Gets what keys is key staked from.
210    """
211    context = make_custom_context(ctx)
212    client = context.com_client()
213
214    if is_ss58_address(key):
215        key_address = key
216    else:
217        keypair = context.load_key(key, password)
218        key_address = keypair.ss58_address
219        key_address = check_ss58_address(key_address)
220    with context.progress_status(
221        f"Getting stake-from map for {key_address}..."
222    ):
223        result = client.get_stakefrom(key=key_address)
224
225    result = {k: format_balance(v, unit) for k, v in result.items()}
226
227    print_table_from_plain_dict(result, ["Key", "Stake"], context.console)
228
229
230@key_app.command()
231def staketo(
232    ctx: Context,
233    key: str,
234    unit: BalanceUnit = BalanceUnit.joule,
235    password: Optional[str] = None,
236):
237    """
238    Gets stake to a key.
239    """
240    context = make_custom_context(ctx)
241    client = context.com_client()
242
243    if is_ss58_address(key):
244        key_address = key
245    else:
246        keypair = context.load_key(key, password)
247        key_address = keypair.ss58_address
248        key_address = check_ss58_address(key_address)
249
250    with context.progress_status(f"Getting stake-to of {key_address}..."):
251        result = client.get_staketo(key=key_address)
252
253    result = {k: format_balance(v, unit) for k, v in result.items()}
254
255    print_table_from_plain_dict(result, ["Key", "Stake"], context.console)
256
257
258@key_app.command()
259def total_free_balance(
260    ctx: Context,
261    unit: BalanceUnit = BalanceUnit.joule,
262    use_universal_password: Optional[str] = typer.Option(
263        False,
264        help="""
265        Password to decrypt all keys.
266        This will only work if all encrypted keys uses the same password.
267        If this is not the case, leave it blank and you will be prompted to give
268        every password.
269        """,
270    ),
271):
272    """
273    Returns total balance of all keys on a disk
274    """
275    context = make_custom_context(ctx)
276    client = context.com_client()
277
278    local_keys = local_key_addresses(context.password_manager)
279    with context.progress_status("Getting total free balance of all keys..."):
280        key2balance: dict[str, int] = local_keys_to_freebalance(
281            client, local_keys
282        )
283
284        balance_sum = sum(key2balance.values())
285
286        context.output(format_balance(balance_sum, unit=unit))
287
288
289@key_app.command()
290def total_staked_balance(
291    ctx: Context,
292    unit: BalanceUnit = BalanceUnit.joule,
293    use_universal_password: bool = typer.Option(
294        False,
295        help="""
296    Password to decrypt all keys.
297    This will only work if all encrypted keys uses the same password.
298    If this is not the case, leave it blank and you will be prompted to give
299    every password.
300    """,
301    ),
302):
303    """
304    Returns total stake of all keys on a disk
305    """
306    context = make_custom_context(ctx)
307    client = context.com_client()
308
309    local_keys = local_key_addresses(context.password_manager)
310    with context.progress_status("Getting total staked balance of all keys..."):
311        key2stake: dict[str, int] = local_keys_to_stakedbalance(
312            client,
313            local_keys,
314        )
315
316        stake_sum = sum(key2stake.values())
317
318        context.output(format_balance(stake_sum, unit=unit))
319
320
321@key_app.command()
322def total_balance(
323    ctx: Context,
324    unit: BalanceUnit = BalanceUnit.joule,
325    use_universal_password: bool = typer.Option(
326        False,
327        help="""
328    Password to decrypt all keys.
329    This will only work if all encrypted keys uses the same password.
330    If this is not the case, leave it blank and you will be prompted to give
331    every password.
332    """,
333    ),
334):
335    """
336    Returns total tokens of all keys on a disk
337    """
338    context = make_custom_context(ctx)
339    client = context.com_client()
340
341    local_keys = local_key_addresses(context.password_manager)
342    with context.progress_status("Getting total tokens of all keys..."):
343        key2balance, key2stake = local_keys_allbalance(client, local_keys)
344        key2tokens = {k: v + key2stake[k] for k, v in key2balance.items()}
345        tokens_sum = sum(key2tokens.values())
346
347        context.output(format_balance(tokens_sum, unit=unit))
348
349
350@key_app.command()
351def power_delegation(
352    ctx: Context,
353    key: Optional[str] = None,
354    enable: bool = typer.Option(True, "--disable"),
355):
356    """
357    Gets power delegation of a key.
358    """
359    context = make_custom_context(ctx)
360    client = context.com_client()
361    if key is None:
362        action = "enable" if enable else "disable"
363        confirm_message = (
364            f"Key was not set, this will {action} vote power delegation for all"
365            " keys on disk. Do you want to proceed?"
366        )
367        if not typer.confirm(confirm_message):
368            context.info("Aborted.")
369            exit(0)
370
371        local_keys = local_key_addresses(context.password_manager)
372    else:
373        local_keys = {key: None}
374    for key_name in local_keys.keys():
375        keypair = context.load_key(key_name, None)
376        if enable is True:
377            context.info(
378                f"Enabling vote power delegation on key {key_name} ..."
379            )
380            client.enable_vote_power_delegation(keypair)
381        else:
382            context.info(
383                f"Disabling vote power delegation on key {key_name} ..."
384            )
385            client.disable_vote_power_delegation(keypair)