communex.cli.key

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