Add --_withdrawal_credentials flag

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 == "":
+        assgined_withdrawal_credentials = None
+    else:
+        validate_withdrawal_credentials(withdrawal_credentials)
+        assgined_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,
+        assgined_withdrawal_credentials=assgined_withdrawal_credentials,
+    )
+    if assgined_withdrawal_credentials is None:
+        withdrawal_credentials_list = tuple([c.withdrawal_credentials for c in credential_list.credentials])
+    else:
+        withdrawal_credentials_list = (assgined_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):

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, assgined_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 assgined_withdrawal_credentials is not None:
+            deposit_message = deposit_message.copy(
+                withdrawal_credentials=assgined_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, assgined_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(assgined_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, assgined_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(assgined_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())

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)