Save Keys to Local Storage (#15)

* Init save keys to local

* Remove MIGRATE_LEGACY and fix linting

* Fix linting

* Fix linting

* Fix linting

* Extract get_beacon_client to eth2.py

* Update help info

* Add folder prompt

* Refactor

* Fix linting

* Fix linting

* Fix

* Fix PR

* Remove unused

* Update

* Fix linting

* Fix PR

* Fix linting

* Fix PR

* Update

* Fix linting

.flake8

@@ -1,3 +1,3 @@
 [flake8]
 max-line-length = 88
-ignore = E501, E203, W503
+ignore = E501, E203, W503, BLK100

.gitignore

@@ -7,3 +7,4 @@ build
 .idea
 local.env
 *.spec
+validator_keys/

operator_cli/commands/generate_proposal.py

@@ -139,7 +139,7 @@ def generate_proposal(chain: str, existing_vault: bool) -> None:
     specification += f"""
 
 - If the proposal will be approved, the operator must perform the following steps:
-    * Call `operator-cli sync-vault` with the same mnemonic as used for generating the proposal
+    * Call `operator-cli sync-vault` or `operator-cli sync-local` with the same mnemonic as used for generating the proposal
     * Create or update validators and make sure the new keys are added
     * Call `commitOperator` from the `{operator}` address
 """

operator_cli/commands/sync_local.py

@@ -0,0 +1,56 @@
+import os
+
+import click
+from eth2deposit.settings import MAINNET
+from requests.exceptions import ConnectionError, HTTPError
+
+from operator_cli.eth2 import get_beacon_client, validate_mnemonic
+from operator_cli.local_storage import LocalStorage
+from operator_cli.settings import SUPPORTED_CHAINS
+
+
+@click.command(help="Synchronizes validator keystores to the local folder")
+@click.option(
+    "--chain",
+    default=MAINNET,
+    help="The network of ETH2 you are targeting.",
+    prompt="Choose the (mainnet or testnet) network/chain name",
+    type=click.Choice(SUPPORTED_CHAINS.keys(), case_sensitive=False),
+)
+def sync_local(chain: str) -> None:
+    while True:
+        try:
+            beacon_client = get_beacon_client()
+            beacon_client.get_genesis()
+            break
+        except (ConnectionError, HTTPError):
+            pass
+
+        click.echo("Error: failed to connect to the ETH2 server with provided URL")
+
+    mnemonic = click.prompt(
+        'Enter your mnemonic separated by spaces (" ")',
+        value_proc=validate_mnemonic,
+        type=click.STRING,
+    )
+
+    folder = click.prompt(
+        "The folder to place the generated keystores and passwords in",
+        default=os.path.join(os.getcwd(), "validator_keys"),
+        type=click.STRING,
+    )
+
+    click.clear()
+    click.confirm(
+        "I confirm that this mnemonic is used only in one staking setup",
+        abort=True,
+    )
+
+    local_storage = LocalStorage(
+        beacon=beacon_client,
+        chain=chain,
+        mnemonic=mnemonic,
+        folder=folder,
+    )
+
+    local_storage.apply_local_changes()

operator_cli/commands/sync_vault.py

@@ -3,9 +3,8 @@
 from hvac import Client as VaultClient
 from hvac.exceptions import InvalidRequest
 from requests.exceptions import ConnectionError, HTTPError
-from web3.beacon import Beacon
 
-from operator_cli.eth2 import validate_mnemonic
+from operator_cli.eth2 import get_beacon_client, validate_mnemonic
 from operator_cli.settings import SUPPORTED_CHAINS, VAULT_VALIDATORS_MOUNT_POINT
 from operator_cli.vault import Vault
 
@@ -16,11 +15,6 @@ def get_vault_client() -> VaultClient:
     return VaultClient(url=url, token=token)
 
 
-def get_beacon_client() -> Beacon:
-    url = click.prompt("Please enter the ETH2 node URL", type=click.STRING)
-    return Beacon(base_url=url)
-
-
 def get_kubernetes_api_server() -> str:
     url = click.prompt(
         "Please enter host string, a host:port pair, or a URL to the base of the Kubernetes API server",

operator_cli/eth2.py

@@ -35,7 +35,6 @@
 from operator_cli.queries import REGISTRATIONS_QUERY
 from operator_cli.settings import (
     MAINNET_WITHDRAWAL_CREDENTIALS,
-    MIGRATE_LEGACY,
     PRATER_WITHDRAWAL_CREDENTIALS,
 )
 from operator_cli.typings import (
@@ -86,6 +85,11 @@ class ValidatorStatus(Enum):
 ]
 
 
+def get_beacon_client() -> Beacon:
+    url = click.prompt("Please enter the ETH2 node URL", type=click.STRING)
+    return Beacon(base_url=url)
+
+
 def validate_mnemonic(mnemonic) -> str:
     if verify_mnemonic(mnemonic, WORD_LISTS_PATH):
         return mnemonic
@@ -169,10 +173,7 @@ def get_mnemonic_signing_key(mnemonic: str, from_index: int) -> SigningKey:
     """Returns the signing key of the mnemonic at a specific index."""
     seed = get_seed(mnemonic=mnemonic, password="")
     private_key = BLSPrivkey(derive_master_SK(seed))
-    if MIGRATE_LEGACY:
-        signing_key_path = f"m/{PURPOSE}/{COIN_TYPE}/0/0/{from_index}"
-    else:
-        signing_key_path = f"m/{PURPOSE}/{COIN_TYPE}/{from_index}/0/0"
+    signing_key_path = f"m/{PURPOSE}/{COIN_TYPE}/{from_index}/0/0"
 
     for node in path_to_nodes(signing_key_path):
         private_key = BLSPrivkey(derive_child_SK(parent_SK=private_key, index=node))

operator_cli/local_storage.py

@@ -0,0 +1,237 @@
+import errno
+import json
+import time
+from functools import cached_property, lru_cache
+from os import listdir, makedirs
+from os.path import exists
+from typing import Dict, Set, Union
+
+import click
+from eth2deposit.key_handling.keystore import ScryptKeystore
+from eth_typing import ChecksumAddress, HexStr
+from py_ecc.bls import G2ProofOfPossession
+from web3 import Web3
+from web3.beacon import Beacon
+
+from operator_cli.eth1 import (
+    get_operators_deposit_data_merkle_proofs,
+    get_validator_operator_address,
+    is_validator_registered,
+)
+from operator_cli.eth2 import (
+    EXITED_STATUSES,
+    generate_password,
+    get_mnemonic_signing_key,
+    get_validators,
+)
+from operator_cli.ipfs import get_operator_deposit_datum
+from operator_cli.queries import get_stakewise_gql_client
+from operator_cli.typings import LocalKeystore, LocalState, SigningKey
+
+
+class LocalStorage(object):
+    def __init__(
+        self,
+        beacon: Beacon,
+        chain: str,
+        mnemonic: str,
+        folder: str,
+    ):
+        self.sw_gql_client = get_stakewise_gql_client(chain)
+        self.beacon = beacon
+        self.mnemonic = mnemonic
+        self.folder = folder
+
+    @cached_property
+    def all_operators_deposit_data_public_keys(self) -> Dict[HexStr, ChecksumAddress]:
+        """Fetches public keys and operators from deposit datum."""
+        deposit_data_merkle_proofs = get_operators_deposit_data_merkle_proofs(
+            self.sw_gql_client
+        )
+        result: Dict[HexStr, ChecksumAddress] = {}
+        with click.progressbar(
+            deposit_data_merkle_proofs.items(),
+            label="Fetching deposit datum\t\t",
+            show_percent=False,
+            show_pos=True,
+        ) as merkle_proofs:
+            for operator_addr, merkle_proofs_url in merkle_proofs:
+                deposit_datum = get_operator_deposit_datum(merkle_proofs_url)
+                for deposit_data in deposit_datum:
+                    public_key = deposit_data["public_key"]
+                    if public_key in result:
+                        raise click.ClickException(
+                            f"Public key {public_key} is presented in"
+                            f" deposit datum for {operator_addr} and {result[public_key]} operators"
+                        )
+                    result[public_key] = operator_addr
+
+        return result
+
+    @cached_property
+    def operator_address(self) -> Union[ChecksumAddress, None]:
+        """Returns local's operator address."""
+        signing_key = get_mnemonic_signing_key(self.mnemonic, 0)
+        first_public_key = Web3.toHex(
+            primitive=G2ProofOfPossession.SkToPk(signing_key.key)
+        )
+        operator_address = get_validator_operator_address(
+            self.sw_gql_client, first_public_key
+        )
+
+        if not operator_address:
+            return self.all_operators_deposit_data_public_keys.get(
+                first_public_key, None
+            )
+
+        return operator_address
+
+    @cached_property
+    def operator_deposit_data_public_keys(self) -> Set[HexStr]:
+        """Returns operator's deposit data public keys."""
+        return set(
+            [
+                pub_key
+                for pub_key, operator in self.all_operators_deposit_data_public_keys.items()
+                if operator == self.operator_address
+            ]
+        )
+
+    @cached_property
+    def generate_keystores(self) -> LocalState:
+        """
+        Returns ordered mapping of BLS public key to private key
+        that are in deposit data or active but are missing in the local.
+        """
+
+        missed_keypairs: Dict[HexStr, SigningKey] = {}
+        from_index = 0
+        while True:
+            signing_key = get_mnemonic_signing_key(self.mnemonic, from_index)
+            public_key = Web3.toHex(G2ProofOfPossession.SkToPk(signing_key.key))
+
+            if public_key in self.operator_deposit_data_public_keys:
+                missed_keypairs[public_key] = signing_key
+                from_index += 1
+                continue
+
+            is_registered = is_validator_registered(
+                gql_client=self.sw_gql_client, public_key=public_key
+            )
+            if is_registered:
+                missed_keypairs[public_key] = signing_key
+                from_index += 1
+                continue
+
+            break
+
+        if not missed_keypairs:
+            return missed_keypairs
+
+        click.confirm(
+            f"Fetched {len(missed_keypairs)} missing validator keys. Save them to the local storage?",
+            abort=True,
+        )
+
+        exited_public_keys: Set[HexStr] = set()
+        missed_keypairs_items = list(missed_keypairs.items())
+        missed_keypairs_count = len(missed_keypairs_items)
+        with click.progressbar(
+            length=missed_keypairs_count,
+            label="Checking local missing keys statuses\t\t",
+            show_percent=False,
+            show_pos=True,
+        ) as bar:
+            for i in range(0, missed_keypairs_count, 100):
+                keypairs_chunk = missed_keypairs_items[i : i + 100]
+                validators = get_validators(
+                    beacon=self.beacon,
+                    public_keys=[HexStr(keypair[0]) for keypair in keypairs_chunk],
+                    state_id="finalized",
+                )
+                for validator in validators:
+                    if validator["status"] in EXITED_STATUSES:
+                        public_key = validator["validator"]["pubkey"]
+                        exited_public_keys.add(public_key)
+
+                bar.update(len(keypairs_chunk))
+
+        for public_key in exited_public_keys:
+            del missed_keypairs[public_key]
+
+        new_state: Dict[int] = {}
+
+        # distribute missing keypairs across validators
+        with click.progressbar(
+            missed_keypairs,
+            label="Provisioning missing validator keys\t\t",
+            show_percent=False,
+            show_pos=True,
+        ) as missing_keypairs:
+            for public_key in missing_keypairs:
+                signing_key = missed_keypairs[public_key]
+                secret = signing_key.key.to_bytes(32, "big")
+                password = self.get_or_create_keystore_password()
+                keystore = ScryptKeystore.encrypt(
+                    secret=secret, password=password, path=signing_key.path
+                ).as_json()
+                new_state[public_key] = LocalKeystore(keystore=keystore)
+        return new_state
+
+    @lru_cache
+    def get_or_create_keystore_password(self) -> str:
+        """Retrieves validator keystore password if exists or creates a new one."""
+        try:
+            with open(f"{self.folder}/password/password.txt") as file:
+                password = file.readline()
+        except FileNotFoundError:
+            password = generate_password()
+            makedirs(f"{self.folder}/password", exist_ok=True)
+            with open(f"{self.folder}/password/password.txt", "w") as file:
+                file.write(password)
+
+        return password
+
+    def apply_local_changes(self) -> None:
+        """Updates local from current state to new state."""
+
+        if exists(self.folder) and len(listdir(self.folder)) > 1:
+            raise click.ClickException(f"{self.folder} already exist and not empty")
+
+        try:
+            makedirs(self.folder)
+        except OSError as e:
+            if e.errno != errno.EEXIST:
+                raise e
+
+        # sync keystores
+        self.sync_local_keystores()
+
+    def sync_local_keystores(self) -> None:
+        """Synchronizes local keystores."""
+        validators_keystores: Dict[str, str] = {}
+        for public_key, local_keystore in self.generate_keystores.items():
+            keystore = local_keystore["keystore"]
+            keystore_path = json.loads(keystore)["path"]
+
+            # generate unique keystore name
+            keystore_name = "keystore-%s-%i.json" % (
+                keystore_path.replace("/", "_"),
+                time.time(),
+            )
+
+            # save keystore
+            validators_keystores[keystore_name] = keystore
+
+        # sync keystores in the local storage
+        with click.progressbar(
+            validators_keystores.items(),
+            label="Syncing local keystores\t\t",
+            show_percent=False,
+            show_pos=True,
+        ) as _validators_keystores:
+            # for validator_name in _validators_keystores:
+            for name, keystore in _validators_keystores:
+                makedirs(f"{self.folder}/keystores", exist_ok=True)
+                with open(f"{self.folder}/keystores/{name}", "w") as file:
+                    file.write(keystore)

operator_cli/main.py

@@ -5,6 +5,7 @@
 warnings.filterwarnings("ignore")
 
 from operator_cli.commands.generate_proposal import generate_proposal  # noqa: E402
+from operator_cli.commands.sync_local import sync_local  # noqa: E402
 from operator_cli.commands.sync_vault import sync_vault  # noqa: E402
 
 
@@ -15,6 +16,7 @@ def cli() -> None:
 
 cli.add_command(generate_proposal)
 cli.add_command(sync_vault)
+cli.add_command(sync_local)
 
 if __name__ == "__main__":
     cli()

operator_cli/settings.py

@@ -63,7 +63,3 @@
     "ETHEREUM_GOERLI_SUBGRAPH_URL",
     default="https://api.thegraph.com/subgraphs/name/stakewise/stakewise-goerli",
 )
-
-# The legacy key derivation path will be used and new vault will be populated with 1000 keys.
-# Skip this flag in case you are not migrating from the legacy system.
-MIGRATE_LEGACY = config("MIGRATE_LEGACY", cast=bool, default=False)