Refactor recover command

src/commands/recover.py

@@ -2,13 +2,15 @@
 from pathlib import Path
 
 import click
-from eth_typing import HexAddress
+from eth_typing import BlockNumber, HexAddress, HexStr
 from eth_utils import add_0x_prefix
-from sw_utils.consensus import EXITED_STATUSES
+from sw_utils.consensus import EXITED_STATUSES, ValidatorStatus
 
 from src.common.clients import consensus_client, execution_client
-from src.common.contracts import vault_contract
+from src.common.contracts import v2_pool_contract, vault_contract
+from src.common.contrib import greenify
 from src.common.credentials import CredentialManager
+from src.common.execution import SECONDS_PER_MONTH
 from src.common.password import generate_password, get_or_create_password_file
 from src.common.utils import log_verbose
 from src.common.validators import validate_eth_address, validate_mnemonic
@@ -83,8 +85,10 @@ def recover(
         vault=vault,
         data_dir=Path(data_dir),
     )
-    if config.vault_dir.exists():
-        raise click.ClickException(f'Vault directory {config.vault_dir} already exists.')
+    if not config.vault_dir.exists():
+        # create vault dir if it does not exist
+        config.vault_dir.mkdir(parents=True)
+        click.secho(f'Vault directory {config.vault_dir} created.', bold=True, fg='green')
 
     keystores_dir = config.vault_dir / 'keystores'
     password_file = keystores_dir / 'password.txt'
@@ -122,106 +126,134 @@ async def main(
     validators = await _fetch_registered_validators()
     if not validators:
         raise click.ClickException('No registered validators')
-    click.secho(f'Found {len(validators)} validators, recovering...')
+
+    total_validators = len(validators)
+    click.confirm(
+        f'Vault has {total_validators} registered validator(s), '
+        f'recover active keystores from provided mnemonic?',
+        default=True,
+        abort=True,
+    )
+
+    if keystores_dir.exists():
+        click.confirm(
+            f'Clean up existing {keystores_dir} keystores directory?',
+            default=True,
+            abort=True,
+        )
+        for file in keystores_dir.glob('*'):
+            file.unlink()
 
     mnemonic_next_index = await _generate_keystores(
-        mnemonic,
-        keystores_dir,
-        password_file,
-        validators,
-        per_keystore_password,
+        mnemonic=mnemonic,
+        keystores_dir=keystores_dir,
+        password_file=password_file,
+        validator_statuses=validators,
+        per_keystore_password=per_keystore_password,
     )
 
     config.save(settings.network, mnemonic, mnemonic_next_index)
     click.secho(
-        f'Keystores for vault {settings.vault} successfully recovered to {keystores_dir}',
-        bold=True,
-        fg='green',
+        f'Successfully recovered {greenify(mnemonic_next_index)} '
+        f'keystores for vault {greenify(settings.vault)}',
     )
 
 
 # pylint: disable-next=too-many-locals
-async def _fetch_registered_validators() -> dict[str, str]:
+async def _fetch_registered_validators() -> dict[HexStr, ValidatorStatus | None]:
     """Fetch registered validators."""
-    block = await execution_client.eth.get_block('latest')
-    current_block = block['number']
-    keeper_genesis_block = settings.network_config.KEEPER_GENESIS_BLOCK
-
-    page_size = 50_000
-    public_keys = []
-
-    for cursor in range(keeper_genesis_block, current_block, page_size):
-        page_start = cursor
-        page_end = min(cursor + page_size - 1, current_block)
+    click.secho('Fetching registered validators...', bold=True)
+    current_block = await execution_client.eth.get_block_number()
+    public_keys = await vault_contract.get_registered_validators_public_keys(
+        from_block=settings.network_config.KEEPER_GENESIS_BLOCK,
+        to_block=current_block,
+    )
 
-        events = await vault_contract.events.ValidatorRegistered.get_logs(
-            fromBlock=page_start, toBlock=page_end
+    if vault_contract.contract_address == settings.network_config.GENESIS_VAULT_CONTRACT_ADDRESS:
+        # fetch registered validators from v2 pool contract
+        # new validators won't be registered after upgrade to the v3,
+        # no need to check up to the latest block
+        blocks_per_month = int(SECONDS_PER_MONTH // settings.network_config.SECONDS_PER_BLOCK)
+        to_block = BlockNumber(
+            min(
+                settings.network_config.KEEPER_GENESIS_BLOCK + blocks_per_month,
+                current_block,
+            )
         )
-        for event in events:
-            hex_key = event['args']['publicKey'].hex()
-            public_keys.append(add_0x_prefix(hex_key))
-
-    validators_status = {}
+        public_keys.extend(
+            await v2_pool_contract.get_registered_validators_public_keys(
+                from_block=settings.network_config.V2_POOL_GENESIS_BLOCK, to_block=to_block
+            )
+        )
+    click.secho(f'Fetched {len(public_keys)} registered validators', bold=True)
 
+    click.secho('Fetching validators statuses...', bold=True)
+    validator_statuses: dict[HexStr, ValidatorStatus | None] = {
+        public_key: None for public_key in public_keys
+    }
     for i in range(0, len(public_keys), settings.validators_fetch_chunk_size):
         validators = await consensus_client.get_validators_by_ids(
             public_keys[i : i + settings.validators_fetch_chunk_size]
         )
         for beacon_validator in validators['data']:
-            validators_status[beacon_validator['validator']['pubkey']] = beacon_validator['status']
+            public_key = add_0x_prefix(beacon_validator['validator']['pubkey'])
+            validator_statuses[public_key] = ValidatorStatus(beacon_validator['status'])
+    click.secho('Fetched statuses for registered validators', bold=True)
 
-    return validators_status
+    return validator_statuses
 
 
 # pylint: disable-next=too-many-arguments,too-many-locals
 async def _generate_keystores(
     mnemonic: str,
     keystores_dir: Path,
     password_file: Path,
-    validators_status: dict[str, str],
+    validator_statuses: dict[HexStr, ValidatorStatus | None],
     per_keystore_password: bool,
 ):
-    keystores_dir.mkdir(parents=True, exist_ok=True)
-    exited_statuses = [x.value for x in EXITED_STATUSES]
-
-    total_validators = len(validators_status)
-
     index = 0
     failed_attempts = 0
 
-    with click.progressbar(  # type: ignore
-        length=total_validators, label='Generating keystores'
-    ) as progress_bar:
-        while total_validators > 0:
-            credential = CredentialManager.generate_credential(
-                network=settings.network,
-                vault=settings.vault,
-                mnemonic=mnemonic,
-                index=index,
-            )
+    validators_count = len(validator_statuses)
+    # stop once failed 1000 times
+    while failed_attempts != 1000:
+        # generate credential
+        credential = CredentialManager.generate_credential(
+            network=settings.network,
+            vault=settings.vault,
+            mnemonic=mnemonic,
+            index=index,
+        )
+        public_key = add_0x_prefix(credential.public_key)
+        # increase index for next iteration
+        index += 1
+
+        # check whether public key is registered
+        if public_key not in validator_statuses:
+            failed_attempts += 1
+            continue
+
+        # get validator status
+        validator_status = validator_statuses.pop(public_key)
+        validators_count -= 1
+
+        # update progress, reset failed attempts
+        failed_attempts = 0
+
+        # skip if validator is already exited
+        if validator_status in EXITED_STATUSES:
+            continue
+
+        # generate password and save keystore
+        password = (
+            generate_password()
+            if per_keystore_password
+            else get_or_create_password_file(str(password_file))
+        )
+        credential.save_signing_keystore(password, str(keystores_dir), per_keystore_password)
+        click.secho(
+            f'Keystore for validator {greenify(public_key)} successfully '
+            f'recovered from mnemonic index {greenify(index - 1)}',
+        )
 
-            public_key = add_0x_prefix(credential.public_key)
-            if public_key in validators_status:
-                validator_status = validators_status.pop(public_key)
-                total_validators -= 1
-                progress_bar.update(1)
-
-                if validator_status not in exited_statuses:
-                    password = (
-                        generate_password()
-                        if per_keystore_password
-                        else get_or_create_password_file(str(password_file))
-                    )
-                    credential.save_signing_keystore(
-                        password, str(keystores_dir), per_keystore_password
-                    )
-                failed_attempts = 0
-            else:
-                failed_attempts += 1
-
-            if failed_attempts > 100:
-                raise click.ClickException('Keystores not found, check mnemonic')
-
-            index += 1
-
-    return index
+    return index - failed_attempts