diff --git a/eth2deposit/cli/generate_keys.py b/eth2deposit/cli/generate_keys.py index 5b94ab70..e60cd5a0 100644 --- a/eth2deposit/cli/generate_keys.py +++ b/eth2deposit/cli/generate_keys.py @@ -4,6 +4,7 @@ Any, Callable, ) +from eth_utils import decode_hex from eth2deposit.credentials import ( CredentialList, @@ -14,6 +15,7 @@ validate_password_strength, ) from eth2deposit.utils.constants import ( + BLS_WITHDRAWAL_PREFIX, MAX_DEPOSIT_AMOUNT, DEFAULT_VALIDATOR_KEYS_FOLDER_NAME, ) @@ -57,6 +59,15 @@ def validate_password(cts: click.Context, param: Any, password: str) -> str: return password +def validate_withdrawal_credentials(withdrawal_credentials: str) -> None: + try: + decoded_withdrawal_credentials = decode_hex(withdrawal_credentials) + except Exception: + raise ValueError("Wrong withdrawal_credentials value.") + if len(decoded_withdrawal_credentials) != 32 or decoded_withdrawal_credentials[:1] != BLS_WITHDRAWAL_PREFIX: + raise ValidationError("Wrong withdrawal_credentials format.") + + def generate_keys_arguments_decorator(function: Callable[..., Any]) -> Callable[..., Any]: ''' This is a decorator that, when applied to a parent-command, implements the @@ -91,6 +102,13 @@ def generate_keys_arguments_decorator(function: Callable[..., Any]) -> Callable[ 'to ask you for your mnemonic as otherwise it will appear in your shell history.)'), prompt='Type the password that secures your validator keystore(s)', ), + click.option( + '--withdrawal_credentials', + default='', + help=("The pre-defined withdrawal_credentials in hex string. If it's unset, the CLI will use the " + "withdrawal key created by the mnemonic to generate the withdrawal_credentials"), + type=str, + ), ] for decorator in reversed(decorators): function = decorator(function) @@ -100,7 +118,14 @@ def generate_keys_arguments_decorator(function: Callable[..., Any]) -> Callable[ @click.command() @click.pass_context def generate_keys(ctx: click.Context, validator_start_index: int, - num_validators: int, folder: str, chain: str, keystore_password: str, **kwargs: Any) -> None: + num_validators: int, folder: str, chain: str, + keystore_password: str, withdrawal_credentials: str, **kwargs: Any) -> None: + if withdrawal_credentials == "": + assigned_withdrawal_credentials = None + else: + validate_withdrawal_credentials(withdrawal_credentials) + assigned_withdrawal_credentials = decode_hex(withdrawal_credentials) + mnemonic = ctx.obj['mnemonic'] mnemonic_password = ctx.obj['mnemonic_password'] amounts = [MAX_DEPOSIT_AMOUNT] * num_validators @@ -120,8 +145,14 @@ def generate_keys(ctx: click.Context, validator_start_index: int, start_index=validator_start_index, ) keystore_filefolders = credential_list.export_keystores(password=keystore_password, folder=folder) - deposits_file = credential_list.export_deposit_data_json(folder=folder) - withdrawal_credentials_list = tuple([c.withdrawal_credentials for c in credential_list.credentials]) + deposits_file = credential_list.export_deposit_data_json( + folder=folder, + assigned_withdrawal_credentials=assigned_withdrawal_credentials, + ) + if assigned_withdrawal_credentials is None: + withdrawal_credentials_list = tuple([c.withdrawal_credentials for c in credential_list.credentials]) + else: + withdrawal_credentials_list = (assigned_withdrawal_credentials,) * len(credential_list.credentials) if not credential_list.verify_keystores(keystore_filefolders=keystore_filefolders, password=keystore_password): raise ValidationError("Failed to verify the keystores.") if not verify_deposit_data_json(deposits_file, withdrawal_credentials_list): diff --git a/eth2deposit/credentials.py b/eth2deposit/credentials.py index d9a4a0b1..8edd8255 100644 --- a/eth2deposit/credentials.py +++ b/eth2deposit/credentials.py @@ -2,7 +2,7 @@ import click import time import json -from typing import Dict, List +from typing import Dict, List, Optional from py_ecc.bls import G2ProofOfPossession as bls from eth2deposit.exceptions import ValidationError @@ -72,23 +72,26 @@ def deposit_message(self) -> DepositMessage: amount=self.amount, ) - @property - def signed_deposit(self) -> DepositData: + def generate_signed_deposit(self, assigned_withdrawal_credentials: Optional[bytes]=None) -> DepositData: domain = compute_deposit_domain(fork_version=self.fork_version) - signing_root = compute_signing_root(self.deposit_message, domain) + deposit_message = self.deposit_message + if assigned_withdrawal_credentials is not None: + deposit_message = deposit_message.copy( + withdrawal_credentials=assigned_withdrawal_credentials + ) + signing_root = compute_signing_root(deposit_message, domain) signed_deposit = DepositData( - **self.deposit_message.as_dict(), + **deposit_message.as_dict(), signature=bls.Sign(self.signing_sk, signing_root) ) return signed_deposit - @property - def deposit_datum_dict(self) -> Dict[str, bytes]: + def generate_deposit_datum_dict(self, assigned_withdrawal_credentials: Optional[bytes]=None) -> Dict[str, bytes]: """ Return a single deposit datum for 1 validator including all the information needed to verify and process the deposit. """ - signed_deposit_datum = self.signed_deposit + signed_deposit_datum = self.generate_signed_deposit(assigned_withdrawal_credentials) datum_dict = signed_deposit_datum.as_dict() datum_dict.update({'deposit_message_root': self.deposit_message.hash_tree_root}) datum_dict.update({'deposit_data_root': signed_deposit_datum.hash_tree_root}) @@ -144,10 +147,10 @@ def export_keystores(self, password: str, folder: str) -> List[str]: show_percent=False, show_pos=True) as credentials: return [credential.save_signing_keystore(password=password, folder=folder) for credential in credentials] - def export_deposit_data_json(self, folder: str) -> str: + def export_deposit_data_json(self, folder: str, assigned_withdrawal_credentials: Optional[bytes]=None) -> str: with click.progressbar(self.credentials, label='Creating your depositdata:\t', show_percent=False, show_pos=True) as credentials: - deposit_data = [cred.deposit_datum_dict for cred in credentials] + deposit_data = [cred.generate_deposit_datum_dict(assigned_withdrawal_credentials) for cred in credentials] filefolder = os.path.join(folder, 'deposit_data-%i.json' % time.time()) with open(filefolder, 'w') as f: json.dump(deposit_data, f, default=lambda x: x.hex()) diff --git a/tests/test_cli/test_generate_keys.py b/tests/test_cli/test_generate_keys.py new file mode 100644 index 00000000..0ebfa11a --- /dev/null +++ b/tests/test_cli/test_generate_keys.py @@ -0,0 +1,21 @@ +import pytest + +from eth2deposit.cli.generate_keys import validate_withdrawal_credentials +from eth2deposit.exceptions import ValidationError + + +@pytest.mark.parametrize( + 'withdrawal_credentials, is_valid', + [ + ('0x2222', False), + ('0x0011111111111111111111111111111111111111111111111111111111111111', True), + ('0011111111111111111111111111111111111111111111111111111111111111', True), + ('001111111111111111111111111111111111111111111111111111111111111122', False), + ] +) +def test_validate_withdrawal_credentials(withdrawal_credentials, is_valid) -> None: + if is_valid: + validate_withdrawal_credentials(withdrawal_credentials) + else: + with pytest.raises((ValidationError, ValueError)): + validate_withdrawal_credentials(withdrawal_credentials)