Generate missing aliases in out of sync Sui aliases and keystore

fix_alias.py

#
# -*- coding: utf-8 -*-
"""Script to generate missing aliases in Sui Config."""

import os
import base64
import json
from pathlib import Path
import yaml
from pysui import SuiConfig
from pysui.abstracts.client_keypair import KeyPair
import pysui.sui.sui_crypto as crypto
import pysui.sui.sui_constants as const
import pysui.sui.sui_utils as utils


def _build_keyref(keys_file: Path) -> dict[str, KeyPair]:
    """Get and generate alias ready public key map."""
    keys_list: list[str] = json.loads(keys_file.read_text("utf-8"))
    keys_map: dict[str, KeyPair] = {}
    for key_string in keys_list:
        kp = crypto.keypair_from_keystring(key_string)
        puks = base64.b64encode(kp.public_key.scheme_and_key()).decode()
        keys_map[puks] = kp
    return keys_map


def _get_aliases(alias_file: Path) -> dict[str, dict]:
    """Get the aliases."""
    alias_list: list[dict] = json.loads(alias_file.read_text("utf-8"))
    return {x["public_key_base64"]: x for x in alias_list}


def _patch_diffs(diffs: list[str], aliascfg: Path, alias_map: list[dict[str, dict]]):
    """Gen and patch up aliases."""
    alias_list: list[dict] = [x for x in alias_map.values()]
    spare_list = alias_list.copy()
    # for each diff, gen a new alias
    for diff in diffs:
        # Get the alias and pluck the first one
        parts1 = list(utils.partition(crypto.gen_mnemonic_phrase(12).split(" "), 6))
        parts2 = [k + "-" + v for k, v in zip(*parts1)][0]
        print(f"Assigning key {diff} to alias {parts2}")
        alias_list.append({"alias": parts2, "public_key_base64": diff})
    # Write the results and verify by trying to load config
    try:
        aliascfg.write_text(json.dumps(alias_list, indent=2))
        _ = SuiConfig.default_config()
        print("Success")
    # If exception, write back the original list
    except Exception as exc:
        print(f"_patch_diffs failed {exc.args}, restoring original aliases")
        aliascfg.write_text(json.dumps(spare_list, indent=2))


def main():
    """Evaluate and patch missing aliases.

    Assumes:
        client config (yaml) is '~/.sui/sui_config/client.yaml'
        keystore (json) is '~/.sui/sui_config/sui.keystore'
        aliases  (json) is '~/.sui/sui_config/sui.aliases'
    """
    aliascfg: Path = Path(os.path.expanduser(const.DEFAULT_ALIAS_PATH_STRING))
    keyscfg: Path = (
        Path(os.path.expanduser(const.DEFAULT_DEVNET_PATH_STRING)).parent
        / "sui.keystore"
    )
    if aliascfg.exists() and keyscfg.exists():
        print(f"Using {aliascfg} and {keyscfg}")
        # Load files and convert to ease of use collections
        key_map = _build_keyref(keyscfg)
        alias_map = _get_aliases(aliascfg)
        key_keys = set(key_map.keys())
        alias_keys = set(alias_map.keys())
        # If lens match then exit
        if len(key_keys) == len(alias_keys):
            print("Sui keys and aliases are synched")
        # Otherwise, get the delta and generate
        else:
            diffs = list(key_keys.difference(alias_keys))
            print(f"Missing aliases for {diffs}... generating\n")
            _patch_diffs(diffs, aliascfg, alias_map)
    else:
        raise ValueError("Currupt default Sui configuration")


if __name__ == "__main__":
    main()