Merge pull request #45 from valefar-on-discord/pbkdf2-support

Pbkdf2 support

README.md

@@ -162,6 +162,7 @@ You can use `new-mnemonic --help` to see all arguments. Note that if there are m
 | `--folder` | String. Pointing to `./validator_keys` by default | The folder path for the keystore(s) and deposit(s) |
 | `--chain` | String. `mainnet` by default | The chain setting for the signing domain. |
 | `--execution_address` (or `--eth1_withdrawal_address`) | String. Eth1 address in hexadecimal encoded form | If this field is set and valid, the given Eth1 address will be used to create the withdrawal credentials. Otherwise, it will generate withdrawal credentials with the mnemonic-derived withdrawal public key in [ERC-2334 format](https://eips.ethereum.org/EIPS/eip-2334#eth2-specific-parameters). |
+| `--pbkdf2` | Flag | Will use pbkdf2 key derivation instead of scrypt for generated keystore files as defined in EIP-2335. This can be a good alternative if you intend to work with a large number of keys. |
 
 ###### `existing-mnemonic` Arguments
 
@@ -174,6 +175,7 @@ You can use `existing-mnemonic --help` to see all arguments. Note that if there
 | `--folder` | String. Pointing to `./validator_keys` by default | The folder path for the keystore(s) and deposit(s) |
 | `--chain` | String. `mainnet` by default | The chain setting for the signing domain. |
 | `--execution_address` (or `--eth1_withdrawal_address`) | String. Eth1 address in hexadecimal encoded form | If this field is set and valid, the given Eth1 address will be used to create the withdrawal credentials. Otherwise, it will generate withdrawal credentials with the mnemonic-derived withdrawal public key in [ERC-2334 format](https://eips.ethereum.org/EIPS/eip-2334#eth2-specific-parameters). |
+| `--pbkdf2` | Flag | Will use pbkdf2 key derivation instead of scrypt for generated keystore files as defined in EIP-2335. This can be a good alternative if you intend to work with a large number of keys. |
 
 ###### Successful message
 

staking_deposit/cli/generate_keys.py

@@ -105,6 +105,12 @@ def generate_keys_arguments_decorator(function: Callable[..., Any]) -> Callable[
             param_decls=['--execution_address', '--eth1_withdrawal_address'],
             prompt=lambda: load_text(['arg_execution_address', 'prompt'], func='generate_keys_arguments_decorator'),
         ),
+        jit_option(
+            default=False,
+            is_flag=True,
+            param_decls='--pbkdf2',
+            help=lambda: load_text(['arg_pbkdf2', 'help'], func='generate_keys_arguments_decorator'),
+        ),
     ]
     for decorator in reversed(decorators):
         function = decorator(function)
@@ -115,7 +121,7 @@ def generate_keys_arguments_decorator(function: Callable[..., Any]) -> Callable[
 @click.pass_context
 def generate_keys(ctx: click.Context, validator_start_index: int,
                   num_validators: int, folder: str, chain: str, keystore_password: str,
-                  execution_address: HexAddress, **kwargs: Any) -> None:
+                  execution_address: HexAddress, pbkdf2: bool, **kwargs: Any) -> None:
     mnemonic = ctx.obj['mnemonic']
     mnemonic_password = ctx.obj['mnemonic_password']
     amounts = [MAX_DEPOSIT_AMOUNT] * num_validators
@@ -134,6 +140,7 @@ def generate_keys(ctx: click.Context, validator_start_index: int,
         chain_setting=chain_setting,
         start_index=validator_start_index,
         hex_eth1_withdrawal_address=execution_address,
+        use_pbkdf2=pbkdf2
     )
     keystore_filefolders = credentials.export_keystores(password=keystore_password, folder=folder)
     deposits_file = credentials.export_deposit_data_json(folder=folder)

staking_deposit/credentials.py

@@ -14,6 +14,7 @@
 from staking_deposit.key_handling.key_derivation.path import mnemonic_and_path_to_key
 from staking_deposit.key_handling.keystore import (
     Keystore,
+    Pbkdf2Keystore,
     ScryptKeystore,
 )
 from staking_deposit.settings import DEPOSIT_CLI_VERSION, BaseChainSetting
@@ -49,7 +50,8 @@ class Credential:
     """
     def __init__(self, *, mnemonic: str, mnemonic_password: str,
                  index: int, amount: int, chain_setting: BaseChainSetting,
-                 hex_eth1_withdrawal_address: Optional[HexAddress]):
+                 hex_eth1_withdrawal_address: Optional[HexAddress],
+                 use_pbkdf2: Optional[bool] = False):
         # Set path as EIP-2334 format
         # https://eips.ethereum.org/EIPS/eip-2334
         purpose = '12381'
@@ -65,6 +67,7 @@ def __init__(self, *, mnemonic: str, mnemonic_password: str,
         self.amount = amount
         self.chain_setting = chain_setting
         self.hex_eth1_withdrawal_address = hex_eth1_withdrawal_address
+        self.use_pbkdf2 = use_pbkdf2
 
     @property
     def signing_pk(self) -> bytes:
@@ -149,7 +152,10 @@ def deposit_datum_dict(self) -> Dict[str, bytes]:
 
     def signing_keystore(self, password: str) -> Keystore:
         secret = self.signing_sk.to_bytes(32, 'big')
-        return ScryptKeystore.encrypt(secret=secret, password=password, path=self.signing_key_path)
+        if self.use_pbkdf2:
+            return Pbkdf2Keystore.encrypt(secret=secret, password=password, path=self.signing_key_path)
+        else:
+            return ScryptKeystore.encrypt(secret=secret, password=password, path=self.signing_key_path)
 
     def save_signing_keystore(self, password: str, folder: str) -> str:
         keystore = self.signing_keystore(password)
@@ -240,7 +246,8 @@ def from_mnemonic(cls,
                       amounts: List[int],
                       chain_setting: BaseChainSetting,
                       start_index: int,
-                      hex_eth1_withdrawal_address: Optional[HexAddress]) -> 'CredentialList':
+                      hex_eth1_withdrawal_address: Optional[HexAddress],
+                      use_pbkdf2: Optional[bool] = False) -> 'CredentialList':
         if len(amounts) != num_keys:
             raise ValueError(
                 f"The number of keys ({num_keys}) doesn't equal to the corresponding deposit amounts ({len(amounts)})."
@@ -250,7 +257,8 @@ def from_mnemonic(cls,
                                show_percent=False, show_pos=True) as indices:
             return cls([Credential(mnemonic=mnemonic, mnemonic_password=mnemonic_password,
                                    index=index, amount=amounts[index - start_index], chain_setting=chain_setting,
-                                   hex_eth1_withdrawal_address=hex_eth1_withdrawal_address)
+                                   hex_eth1_withdrawal_address=hex_eth1_withdrawal_address,
+                                   use_pbkdf2=use_pbkdf2)
                         for index in indices])
 
     def export_keystores(self, password: str, folder: str) -> List[str]:

staking_deposit/intl/en/cli/generate_keys.json

@@ -22,6 +22,9 @@
             "prompt": "Please enter the optional 20-byte execution address for the new withdrawal credentials. Note that you CANNOT change it once you have set it on chain.",
             "confirm": "Repeat your execution address for confirmation.",
             "mismatch": "Error: the two entered values do not match. Please type again."
+        },
+        "arg_pbkdf2": {
+            "help": "Uses the pbkdf2 hashing function instead of scrypt for generated keystore files. "
         }
     },
     "generate_keys": {

tests/test_cli/test_existing_mnemonic.py

@@ -168,6 +168,84 @@ def test_existing_mnemonic_eth1_address_withdrawal_bad_checksum() -> None:
     clean_key_folder(my_folder_path)
 
 
+def test_pbkdf2_new_mnemonic() -> None:
+    # Prepare pbkdf2 folder
+    pbkdf2_folder_path = os.path.join(os.getcwd(), 'TESTING_TEMP_FOLDER')
+    clean_key_folder(pbkdf2_folder_path)
+    if not os.path.exists(pbkdf2_folder_path):
+        os.mkdir(pbkdf2_folder_path)
+
+    # Prepare scrypt folder
+    scrypt_folder_path = os.path.join(os.getcwd(), 'TESTING_TEMP_FOLDER_2')
+    clean_key_folder(scrypt_folder_path)
+    if not os.path.exists(scrypt_folder_path):
+        os.mkdir(scrypt_folder_path)
+
+    runner = CliRunner()
+    inputs = [
+        'abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about',
+        '0', '0', '1', 'mainnet', 'MyPassword', 'MyPassword',
+    ]
+    data = '\n'.join(inputs)
+    arguments = [
+        '--language', 'english',
+        'existing-mnemonic',
+        '--eth1_withdrawal_address', '',
+        '--folder', pbkdf2_folder_path,
+        '--pbkdf2',
+    ]
+    result = runner.invoke(cli, arguments, input=data)
+    assert result.exit_code == 0
+
+    arguments = [
+        '--language', 'english',
+        'existing-mnemonic',
+        '--eth1_withdrawal_address', '',
+        '--folder', scrypt_folder_path,
+    ]
+    result = runner.invoke(cli, arguments, input=data)
+    assert result.exit_code == 0
+
+    # Load store generated files
+    validator_keys_folder_path = os.path.join(pbkdf2_folder_path, DEFAULT_VALIDATOR_KEYS_FOLDER_NAME)
+    _, _, key_files = next(os.walk(validator_keys_folder_path))
+
+    deposit_file = [key_file for key_file in key_files if key_file.startswith('deposit_data')][0]
+    with open(validator_keys_folder_path + '/' + deposit_file, 'r', encoding='utf-8') as f:
+        pbkdf2_deposit_dict = json.load(f)[0]
+
+    keystore_file = [key_file for key_file in key_files if key_file.startswith('keystore-m_')][0]
+    with open(validator_keys_folder_path + '/' + keystore_file, 'r', encoding='utf-8') as f:
+        pbkdf2_keystore_dict = json.load(f)
+
+    validator_keys_folder_path = os.path.join(scrypt_folder_path, DEFAULT_VALIDATOR_KEYS_FOLDER_NAME)
+    _, _, key_files = next(os.walk(validator_keys_folder_path))
+
+    deposit_file = [key_file for key_file in key_files if key_file.startswith('deposit_data')][0]
+    with open(validator_keys_folder_path + '/' + deposit_file, 'r', encoding='utf-8') as f:
+        scrypt_deposit_dict = json.load(f)[0]
+
+    keystore_file = [key_file for key_file in key_files if key_file.startswith('keystore-m_')][0]
+    with open(validator_keys_folder_path + '/' + keystore_file, 'r', encoding='utf-8') as f:
+        scrypt_keystore_dict = json.load(f)
+
+    # Verify deposit files
+    assert pbkdf2_deposit_dict['withdrawal_credentials'] == scrypt_deposit_dict['withdrawal_credentials']
+    assert pbkdf2_deposit_dict['pubkey'] == scrypt_deposit_dict['pubkey']
+    assert pbkdf2_deposit_dict['signature'] == scrypt_deposit_dict['signature']
+    assert pbkdf2_deposit_dict['deposit_message_root'] == scrypt_deposit_dict['deposit_message_root']
+    assert pbkdf2_deposit_dict['deposit_data_root'] == scrypt_deposit_dict['deposit_data_root']
+
+    # Verify keystore files
+    assert pbkdf2_keystore_dict['crypto']['kdf']['function'] == 'pbkdf2'
+    assert scrypt_keystore_dict['crypto']['kdf']['function'] == 'scrypt'
+    assert pbkdf2_keystore_dict['pubkey'] == scrypt_keystore_dict['pubkey']
+
+    # Clean up
+    clean_key_folder(pbkdf2_folder_path)
+    clean_key_folder(scrypt_folder_path)
+
+
 @pytest.mark.asyncio
 async def test_script() -> None:
     my_folder_path = os.path.join(os.getcwd(), 'TESTING_TEMP_FOLDER')

tests/test_cli/test_exit_transaction_keystore.py

@@ -44,8 +44,6 @@ def test_exit_transaction_keystore() -> None:
     keystore_filepath = credential.save_signing_keystore(keystore_password, exit_transaction_folder_path)
 
     runner = CliRunner()
-    inputs = []
-    data = '\n'.join(inputs)
     arguments = [
         '--language', 'english',
         '--non_interactive',
@@ -57,7 +55,7 @@ def test_exit_transaction_keystore() -> None:
         '--validator_index', '1',
         '--epoch', '1234',
     ]
-    result = runner.invoke(cli, arguments, input=data)
+    result = runner.invoke(cli, arguments)
 
     assert result.exit_code == 0
 
@@ -83,6 +81,108 @@ def test_exit_transaction_keystore() -> None:
     clean_exit_transaction_folder(my_folder_path)
 
 
+def test_exit_transaction_with_pbkdf2() -> None:
+    # Prepare folder
+    pbkdf2_folder_path = os.path.join(os.getcwd(), 'TESTING_TEMP_FOLDER')
+    pbkdf2_exit_transaction_folder_path = os.path.join(pbkdf2_folder_path, DEFAULT_EXIT_TRANSACTION_FOLDER_NAME)
+    clean_exit_transaction_folder(pbkdf2_folder_path)
+    scrypt_folder_path = os.path.join(os.getcwd(), 'TESTING_TEMP_FOLDER2')
+    scrypt_exit_transaction_folder_path = os.path.join(scrypt_folder_path, DEFAULT_EXIT_TRANSACTION_FOLDER_NAME)
+    clean_exit_transaction_folder(pbkdf2_folder_path)
+    if not os.path.exists(pbkdf2_folder_path):
+        os.mkdir(pbkdf2_folder_path)
+    if not os.path.exists(scrypt_folder_path):
+        os.mkdir(scrypt_folder_path)
+    if not os.path.exists(pbkdf2_exit_transaction_folder_path):
+        os.mkdir(pbkdf2_exit_transaction_folder_path)
+    if not os.path.exists(scrypt_exit_transaction_folder_path):
+        os.mkdir(scrypt_exit_transaction_folder_path)
+
+    # Shared parameters
+    chain = 'mainnet'
+    keystore_password = 'solo-stakers'
+
+    # Prepare credential
+    pbkdf2_credential = Credential(
+        mnemonic='aban aban aban aban aban aban aban aban aban aban aban abou',
+        mnemonic_password='',
+        index=0,
+        amount=0,
+        chain_setting=get_chain_setting(chain),
+        hex_eth1_withdrawal_address=None,
+        use_pbkdf2=True,
+    )
+    scrypt_credential = Credential(
+        mnemonic='aban aban aban aban aban aban aban aban aban aban aban abou',
+        mnemonic_password='',
+        index=0,
+        amount=0,
+        chain_setting=get_chain_setting(chain),
+        hex_eth1_withdrawal_address=None,
+        use_pbkdf2=False,
+    )
+
+    # Save keystore file
+    pbkdf2_keystore_filepath = pbkdf2_credential.save_signing_keystore(
+        keystore_password,
+        pbkdf2_exit_transaction_folder_path,
+    )
+    scrypt_keystore_filepath = scrypt_credential.save_signing_keystore(
+        keystore_password,
+        scrypt_exit_transaction_folder_path,
+    )
+
+    runner = CliRunner()
+    arguments = [
+        '--language', 'english',
+        '--non_interactive',
+        'exit-transaction-keystore',
+        '--output_folder', pbkdf2_folder_path,
+        '--chain', chain,
+        '--keystore', pbkdf2_keystore_filepath,
+        '--keystore_password', keystore_password,
+        '--validator_index', '1',
+        '--epoch', '1234',
+    ]
+    result = runner.invoke(cli, arguments)
+    assert result.exit_code == 0
+
+    arguments = [
+        '--language', 'english',
+        '--non_interactive',
+        'exit-transaction-keystore',
+        '--output_folder', scrypt_folder_path,
+        '--chain', chain,
+        '--keystore', scrypt_keystore_filepath,
+        '--keystore_password', keystore_password,
+        '--validator_index', '1',
+        '--epoch', '1234',
+    ]
+    result = runner.invoke(cli, arguments)
+    assert result.exit_code == 0
+
+    # Check files
+    _, _, exit_transaction_files = next(os.walk(pbkdf2_exit_transaction_folder_path))
+    pbkdf2_exit_transaction_file = [f for f in exit_transaction_files if 'signed_exit' in f]
+    assert len(set(pbkdf2_exit_transaction_file)) == 1
+    pbkdf2_json_data = read_json_file(pbkdf2_exit_transaction_folder_path, pbkdf2_exit_transaction_file[0])
+
+    _, _, exit_transaction_files = next(os.walk(scrypt_exit_transaction_folder_path))
+    scrypt_exit_transaction_file = [f for f in exit_transaction_files if 'signed_exit' in f]
+    assert len(set(scrypt_exit_transaction_file)) == 1
+    scrypt_json_data = read_json_file(scrypt_exit_transaction_folder_path, scrypt_exit_transaction_file[0])
+
+    assert pbkdf2_json_data['message']['epoch'] == scrypt_json_data['message']['epoch']
+    assert pbkdf2_json_data['message']['validator_index'] == scrypt_json_data['message']['validator_index']
+    assert pbkdf2_json_data['signature'] == scrypt_json_data['signature']
+
+    verify_file_permission(os, folder_path=pbkdf2_exit_transaction_folder_path, files=pbkdf2_exit_transaction_file)
+    verify_file_permission(os, folder_path=scrypt_exit_transaction_folder_path, files=scrypt_exit_transaction_file)
+
+    clean_exit_transaction_folder(pbkdf2_folder_path)
+    clean_exit_transaction_folder(scrypt_folder_path)
+
+
 def test_invalid_keystore_path() -> None:
     my_folder_path = os.path.join(os.getcwd(), 'TESTING_TEMP_FOLDER')
     clean_exit_transaction_folder(my_folder_path)