From dfff5c67223ef2d3b36fa56919495d411ccc694b Mon Sep 17 00:00:00 2001 From: evgeny-stakewise <123374581+evgeny-stakewise@users.noreply.github.com> Date: Thu, 15 Feb 2024 13:06:39 +0300 Subject: [PATCH] Rework remote db (#293) * Del parent-public-key in remote-db * Rename test_tasks -> test_commands * Fixes --- src/remote_db/commands.py | 16 +- src/remote_db/database.py | 70 +---- src/remote_db/tasks.py | 24 +- src/remote_db/tests/test_commands.py | 344 ++++++++++++++++++++++++ src/remote_db/tests/test_tasks.py | 388 +-------------------------- src/remote_db/typings.py | 1 - 6 files changed, 383 insertions(+), 460 deletions(-) create mode 100644 src/remote_db/tests/test_commands.py diff --git a/src/remote_db/commands.py b/src/remote_db/commands.py index 52dabb08..4fb136d3 100644 --- a/src/remote_db/commands.py +++ b/src/remote_db/commands.py @@ -109,9 +109,7 @@ def cleanup(ctx: Context) -> None: click.echo(f'Successfully removed all the entries for the {greenify(settings.vault)} vault.') -@remote_db_group.command( - help='Generates shares for the local keypairs, updates configs in the remote DB.' -) +@remote_db_group.command(help='Uploads key-pairs to remote DB. Updates configs in the remote DB.') @click.option( '--encrypt-key', envvar='REMOTE_DB_ENCRYPT_KEY', @@ -133,12 +131,19 @@ def cleanup(ctx: Context) -> None: help='Path to the deposit_data.json file. ' 'Default is the file generated with "create-keys" command.', ) +@click.option( + '--pool-size', + help='Number of processes in a pool.', + envvar='POOL_SIZE', + type=int, +) @click.pass_context def upload_keypairs( ctx: Context, encrypt_key: str, execution_endpoints: str, deposit_data_file: str | None, + pool_size: int | None, ) -> None: settings.set( vault=settings.vault, @@ -148,12 +153,11 @@ def upload_keypairs( deposit_data_file=deposit_data_file, verbose=settings.verbose, execution_endpoints=execution_endpoints, + pool_size=pool_size, ) try: asyncio.run(tasks.upload_keypairs(ctx.obj['db_url'], encrypt_key)) - click.echo( - f'Successfully uploaded keypairs and shares for the {greenify(settings.vault)} vault.' - ) + click.echo(f'Successfully uploaded keypairs for the {greenify(settings.vault)} vault.') except Exception as e: log_verbose(e) diff --git a/src/remote_db/database.py b/src/remote_db/database.py index 7978f409..94822d19 100644 --- a/src/remote_db/database.py +++ b/src/remote_db/database.py @@ -50,7 +50,7 @@ def get_first_keypair(self) -> RemoteDatabaseKeyPair | None: with self.db_connection.cursor() as cur: cur.execute( f''' - SELECT parent_public_key, public_key, private_key, nonce + SELECT public_key, private_key, nonce FROM {self.table} WHERE vault = %s ORDER BY public_key @@ -63,32 +63,20 @@ def get_first_keypair(self) -> RemoteDatabaseKeyPair | None: return None return RemoteDatabaseKeyPair( vault=settings.vault, - parent_public_key=row[0], - public_key=row[1], - private_key=row[2], - nonce=row[3], + public_key=row[0], + private_key=row[1], + nonce=row[2], ) - def get_keypairs( - self, has_parent_public_key: bool | None = None - ) -> list[RemoteDatabaseKeyPair]: + def get_keypairs(self) -> list[RemoteDatabaseKeyPair]: """Returns keypairs from the database.""" - where_list = [] params: dict = {} - if has_parent_public_key is True: - where_list.append('parent_public_key IS NOT NULL') - elif has_parent_public_key is False: - where_list.append('parent_public_key IS NULL') - query = f''' - SELECT parent_public_key, public_key, private_key, nonce + SELECT public_key, private_key, nonce FROM {self.table} ''' - if where_list: - query += f'WHERE {" AND ".join(where_list)}\n' - query += 'ORDER BY public_key' with self.db_connection.cursor() as cur: @@ -98,22 +86,21 @@ def get_keypairs( return [ RemoteDatabaseKeyPair( vault=settings.vault, - parent_public_key=row[0], - public_key=row[1], - private_key=row[2], - nonce=row[3], + public_key=row[0], + private_key=row[1], + nonce=row[2], ) for row in res ] - def remove_keypairs(self, in_parent_public_keys: set[HexStr] | None = None) -> None: + def remove_keypairs(self, in_public_keys: set[HexStr] | None = None) -> None: """Removes keypairs from the database.""" where_list = ['vault = %(vault)s'] params: dict = {'vault': settings.vault} - if in_parent_public_keys is not None: - where_list.append('parent_public_key IN %(in_parent_public_keys)s') - params['in_parent_public_keys'] = tuple(in_parent_public_keys) + if in_public_keys is not None: + where_list.append('public_key IN %(in_public_keys)s') + params['in_public_keys'] = tuple(in_public_keys) query = f''' DELETE FROM {self.table} @@ -130,18 +117,16 @@ def upload_keypairs(self, keypairs: list[RemoteDatabaseKeyPair]) -> None: f''' INSERT INTO {self.table} ( vault, - parent_public_key, public_key, private_key, nonce ) - VALUES (%s, %s, %s, %s, %s) + VALUES (%s, %s, %s, %s) ON CONFLICT DO NOTHING ''', [ ( keypair.vault, - keypair.parent_public_key, keypair.public_key, keypair.private_key, keypair.nonce, @@ -157,7 +142,6 @@ def create_table(self) -> None: f''' CREATE TABLE IF NOT EXISTS {self.table} ( vault VARCHAR(42) NOT NULL, - parent_public_key VARCHAR(98), public_key VARCHAR(98) UNIQUE NOT NULL, private_key VARCHAR(66) UNIQUE NOT NULL, nonce VARCHAR(34) UNIQUE NOT NULL @@ -167,7 +151,6 @@ def create_table(self) -> None: class ConfigsCrud: - remote_signer_config_name = 'remote_signer_config.json' deposit_data_name = 'deposit_data.json' def __init__(self, db_connection: Any | None = None, db_url: str | None = None): @@ -187,18 +170,6 @@ def get_configs_count(self) -> int: row = cur.fetchone() return row[0] - def get_remote_signer_config(self) -> dict | None: - """Returns the remote signer config from the database.""" - with self.db_connection.cursor() as cur: - cur.execute( - f'SELECT data FROM {self.table} WHERE vault = %s AND name = %s', - (settings.vault, self.remote_signer_config_name), - ) - row = cur.fetchone() - if row is None: - return None - return json.loads(row[0]) - def get_deposit_data(self) -> list | None: """Returns the deposit data from the database.""" with self.db_connection.cursor() as cur: @@ -211,19 +182,6 @@ def get_deposit_data(self) -> list | None: return None return json.loads(row[0]) - def update_remote_signer_config(self, data: dict) -> None: - """Updates the remote signer config in the database.""" - data_string = json.dumps(data) - with self.db_connection.cursor() as cur: - cur.execute( - f''' - INSERT INTO {self.table} (vault, name, data) - VALUES (%s, %s, %s) - ON CONFLICT (vault, name) DO UPDATE SET data = %s - ''', - (settings.vault, self.remote_signer_config_name, data_string, data_string), - ) - def update_deposit_data(self, deposit_data: list[dict]) -> None: """Updates the deposit data in the database.""" data_string = json.dumps(deposit_data) diff --git a/src/remote_db/tasks.py b/src/remote_db/tasks.py index 04d6938c..efe3d559 100644 --- a/src/remote_db/tasks.py +++ b/src/remote_db/tasks.py @@ -64,7 +64,7 @@ def cleanup(db_url: str) -> None: # pylint: disable=too-many-locals async def upload_keypairs(db_url: str, b64_encrypt_key: str) -> None: - """Generates shares for the local keypairs, updates configs in the remote DB.""" + """Uploads key-pairs to remote DB. Updates configs in the remote DB.""" encryption_key = _check_encryption_key(db_url, b64_encrypt_key) # load and check deposit data file @@ -79,7 +79,7 @@ async def upload_keypairs(db_url: str, b64_encrypt_key: str) -> None: if len(keystore) == 0: raise click.ClickException('Keystore not found.') - click.echo(f'Calculating and encrypting shares for {len(keystore)} keystores...') + click.echo(f'Encrypting {len(keystore)} keystores...') key_records: list[RemoteDatabaseKeyPair] = [] for public_key, private_key in keystore.keys.items(): # pylint: disable=no-member encrypted_priv_key, nonce = _encrypt_private_key(private_key, encryption_key) @@ -149,7 +149,7 @@ def setup_validator( output_dir: Path, ) -> None: """Generate validator configs for Lighthouse, Teku and Prysm clients.""" - keypairs = KeyPairsCrud(db_url=db_url).get_keypairs(has_parent_public_key=False) + keypairs = KeyPairsCrud(db_url=db_url).get_keypairs() if not keypairs: raise click.ClickException('No keypairs found in the remote db.') @@ -234,20 +234,24 @@ def _check_encryption_key(db_url: str, b64_encrypt_key: str) -> bytes: encryption_key = base64.b64decode(b64_encrypt_key) if len(encryption_key) != CIPHER_KEY_LENGTH: raise click.ClickException('Invalid encryption key length.') + except Exception as exc: + raise click.ClickException('Invalid encryption key.') from exc - keypair = KeyPairsCrud(db_url=db_url).get_first_keypair() - if keypair is None: - return encryption_key + keypair = KeyPairsCrud(db_url=db_url).get_first_keypair() + if keypair is None: + return encryption_key + try: decrypted_private_key = _decrypt_private_key( private_key=Web3.to_bytes(hexstr=keypair.private_key), encryption_key=encryption_key, nonce=Web3.to_bytes(hexstr=keypair.nonce), ) - if bls.SkToPk(decrypted_private_key) != Web3.to_bytes(hexstr=keypair.public_key): - raise click.ClickException('Failed to decrypt first private key.') - except Exception as exc: - raise click.ClickException('Invalid encryption key.') from exc + except Exception as e: + raise click.ClickException('Failed to decrypt first private key.') from e + + if bls.SkToPk(decrypted_private_key) != Web3.to_bytes(hexstr=keypair.public_key): + raise click.ClickException('Failed to decrypt first private key.') return encryption_key diff --git a/src/remote_db/tests/test_commands.py b/src/remote_db/tests/test_commands.py new file mode 100644 index 00000000..6945deaf --- /dev/null +++ b/src/remote_db/tests/test_commands.py @@ -0,0 +1,344 @@ +import base64 +from collections import defaultdict +from pathlib import Path +from secrets import randbits +from typing import Generator +from unittest import mock + +import milagro_bls_binding as bls +import pytest +from click.testing import CliRunner +from eth_typing import ChecksumAddress, HexAddress +from py_ecc.bls import G2ProofOfPossession +from web3 import Web3 + +from src.remote_db.commands import remote_db_group +from src.remote_db.database import ConfigsCrud, KeyPairsCrud +from src.remote_db.tasks import _encrypt_private_key +from src.remote_db.typings import RemoteDatabaseKeyPair +from src.validators.typings import BLSPrivkey + + +@pytest.fixture +def _patch_check_db_connection() -> Generator: + with mock.patch('src.remote_db.commands.check_db_connection', return_value=None): + yield + + +@pytest.fixture +def _patch_get_db_connection() -> Generator: + with mock.patch('src.remote_db.tasks.get_db_connection'), mock.patch( + 'src.remote_db.database.get_db_connection' + ): + yield + + +@pytest.fixture +def _patch_check_deposit_data_root() -> Generator: + with mock.patch('src.remote_db.tasks.check_deposit_data_root'): + yield + + +def _get_remote_db_keypairs( + encryption_key: bytes, vault_address: HexAddress, total_validators: int +) -> list[RemoteDatabaseKeyPair]: + keystores = {} + for _ in range(total_validators): + seed = randbits(256).to_bytes(32, 'big') + private_key = BLSPrivkey(G2ProofOfPossession.KeyGen(seed).to_bytes(32, 'big')) + public_key = Web3.to_hex(bls.SkToPk(private_key)) + keystores[public_key] = private_key + + key_records: list[RemoteDatabaseKeyPair] = [] + for public_key, private_key in keystores.items(): # pylint: disable=no-member + encrypted_priv_key, nonce = _encrypt_private_key(private_key, encryption_key) + key_records.append( + RemoteDatabaseKeyPair( + vault=ChecksumAddress(vault_address), + public_key=public_key, + private_key=Web3.to_hex(encrypted_priv_key), + nonce=Web3.to_hex(nonce), + ) + ) + return key_records + + +@pytest.mark.usefixtures('_init_vault', '_patch_check_db_connection', '_patch_get_db_connection') +class TestRemoteDbSetup: + def test_setup_works( + self, + data_dir: Path, + vault_address: HexAddress, + runner: CliRunner, + ): + db_url = 'postgresql://user:password@localhost:5432/dbname' + + args = ['--db-url', db_url, '--vault', vault_address, '--data-dir', str(data_dir), 'setup'] + with mock.patch.object( + KeyPairsCrud, 'get_keypairs_count', return_value=0 + ) as get_keypairs_count_mock, mock.patch.object( + ConfigsCrud, 'get_configs_count', return_value=0 + ) as get_configs_count_mock, mock.patch( + 'src.remote_db.tasks.get_random_bytes', return_value=b'1' + ): + result = runner.invoke(remote_db_group, args) + assert get_keypairs_count_mock.call_count == 1 + assert get_configs_count_mock.call_count == 1 + + output = ( + 'Successfully configured remote database.\n' + 'Encryption key: MQ==\n' + 'NB! You must store your encryption in a secure cold storage!' + ) + assert output.strip() == result.output.strip() + + def test_cleanup_works( + self, + data_dir: Path, + vault_address: HexAddress, + runner: CliRunner, + ): + db_url = 'postgresql://user:password@localhost:5432/dbname' + args = [ + '--db-url', + db_url, + '--vault', + vault_address, + '--data-dir', + str(data_dir), + 'cleanup', + ] + with mock.patch.object( + KeyPairsCrud, 'remove_keypairs' + ) as remove_keypairs_mock, mock.patch.object( + ConfigsCrud, 'remove_configs' + ) as remove_configs_mock: + result = runner.invoke(remote_db_group, args) + assert remove_keypairs_mock.call_count == 1 + assert remove_configs_mock.call_count == 1 + + output = 'Successfully removed all the entries for the ' f'{vault_address} vault.' + assert output.strip() == result.output.strip() + + def test_setup_fails_when_keypairs_not_empty( + self, + data_dir: Path, + vault_address: HexAddress, + runner: CliRunner, + ): + db_url = 'postgresql://user:password@localhost:5432/dbname' + + args = ['--db-url', db_url, '--vault', vault_address, '--data-dir', str(data_dir), 'setup'] + with mock.patch.object( + KeyPairsCrud, 'get_keypairs_count', return_value=1 + ) as get_keypairs_count_mock: + result = runner.invoke(remote_db_group, args) + assert get_keypairs_count_mock.call_count == 1 + + output = ( + 'Error: Error: the remote database is not empty. Please clean up with "clean" ' + 'command first.' + ) + assert output.strip() == result.output.strip() + + def test_setup_fails_when_configs_not_empty( + self, + data_dir: Path, + vault_address: HexAddress, + runner: CliRunner, + ): + db_url = 'postgresql://user:password@localhost:5432/dbname' + + args = ['--db-url', db_url, '--vault', vault_address, '--data-dir', str(data_dir), 'setup'] + with ( + mock.patch.object(KeyPairsCrud, 'get_keypairs_count', return_value=0), + mock.patch.object( + ConfigsCrud, 'get_configs_count', return_value=1 + ) as get_configs_count_mock, + ): + result = runner.invoke(remote_db_group, args) + assert get_configs_count_mock.call_count == 1 + + output = ( + 'Error: Error: the remote database is not empty. Please clean up with "clean" ' + 'command first.' + ) + assert output.strip() == result.output.strip() + + +@pytest.mark.usefixtures( + '_patch_check_db_connection', + '_patch_get_db_connection', + '_patch_check_deposit_data_root', +) +@pytest.mark.usefixtures('_init_vault', '_create_keys') +class TestRemoteDbUploadKeypairs: + def test_upload_keypairs_works( + self, + data_dir: Path, + vault_address: HexAddress, + runner: CliRunner, + execution_endpoints: str, + ): + db_url = 'postgresql://user:password@localhost:5432/dbname' + encryption_key = '43ueY4nqsiajWHTnkdqrc3OWj2W+t0bbdBISJFjZ3Ck=' + + args = [ + '--db-url', + db_url, + '--vault', + vault_address, + '--data-dir', + str(data_dir), + 'upload-keypairs', + '--execution-endpoints', + execution_endpoints, + '--encrypt-key', + encryption_key, + ] + + with mock.patch.object(KeyPairsCrud, 'get_first_keypair', return_value=None): + result = runner.invoke(remote_db_group, args) + output = f'Successfully uploaded keypairs for the {vault_address} vault.' + assert output.strip() in result.output.strip() + + +@pytest.mark.usefixtures( + '_patch_check_db_connection', + '_patch_get_db_connection', +) +@pytest.mark.usefixtures('_init_vault', '_create_keys') +class TestRemoteDbSetupWeb3Signer: + def test_setup_web3signer_works( + self, + data_dir: Path, + vault_address: HexAddress, + runner: CliRunner, + execution_endpoints: str, + ): + db_url = 'postgresql://user:password@localhost:5432/dbname' + encryption_key = '43ueY4nqsiajWHTnkdqrc3OWj2W+t0bbdBISJFjZ3Ck=' + + args = [ + '--db-url', + db_url, + '--vault', + vault_address, + '--data-dir', + str(data_dir), + 'setup-web3signer', + '--output-dir', + './web3signer', + '--encrypt-key', + encryption_key, + ] + keypairs = _get_remote_db_keypairs( + base64.b64decode(encryption_key), vault_address, total_validators=3 + ) + + with runner.isolated_filesystem(), mock.patch.object( + KeyPairsCrud, 'get_first_keypair', return_value=keypairs[0] + ), mock.patch.object(KeyPairsCrud, 'get_keypairs', return_value=keypairs): + result = runner.invoke(remote_db_group, args) + output = 'Successfully retrieved web3signer private keys from the database.\n' + assert output.strip() in result.output.strip() + + +@pytest.mark.usefixtures( + '_patch_check_db_connection', + '_patch_get_db_connection', +) +@pytest.mark.usefixtures('_init_vault', '_create_keys') +class TestRemoteDbSetupValidator: + def test_setup_validator( + self, + data_dir: Path, + vault_address: HexAddress, + runner: CliRunner, + ): + db_url = 'postgresql://user:password@localhost:5432/dbname' + encryption_key = '43ueY4nqsiajWHTnkdqrc3OWj2W+t0bbdBISJFjZ3Ck=' + + args = [ + '--db-url', + db_url, + '--vault', + vault_address, + '--data-dir', + str(data_dir), + 'setup-validator', + '--output-dir', + './validator', + '--validator-index', + '0', + '--total-validators', + '1', + '--web3signer-endpoint', + 'http://localhost:8080', + '--fee-recipient', + vault_address, + ] + keypairs = _get_remote_db_keypairs( + base64.b64decode(encryption_key), vault_address, total_validators=3 + ) + + with runner.isolated_filesystem(), mock.patch.object( + KeyPairsCrud, 'get_first_keypair', return_value=keypairs[0] + ), mock.patch.object(KeyPairsCrud, 'get_keypairs', return_value=keypairs): + result = runner.invoke(remote_db_group, args) + output = ( + 'Generated configs with 3 keys ' + 'for validator with index 0.\n' + 'Validator definitions for Lighthouse saved to ' + 'validator/validator_definitions.yml file.\n' + 'Signer keys for Teku\\Prysm saved to ' + 'validator/signer_keys.yml file.\n' + 'Proposer config for Teku\\Prysm saved to ' + 'validator/proposer_config.json file.\n' + ) + assert output.strip() in result.output.strip() + + +@pytest.mark.usefixtures( + '_patch_check_db_connection', + '_patch_get_db_connection', +) +@pytest.mark.usefixtures('_init_vault', '_create_keys') +class TestRemoteDbSetupOperator: + def test_setup_operator( + self, + data_dir: Path, + vault_address: HexAddress, + runner: CliRunner, + ): + db_url = 'postgresql://user:password@localhost:5432/dbname' + encryption_key = '43ueY4nqsiajWHTnkdqrc3OWj2W+t0bbdBISJFjZ3Ck=' + + args = [ + '--db-url', + db_url, + '--vault', + vault_address, + '--data-dir', + str(data_dir), + 'setup-operator', + '--output-dir', + './operator', + ] + keypairs = _get_remote_db_keypairs( + base64.b64decode(encryption_key), vault_address, total_validators=3 + ) + remote_config: dict[str, list[str]] = defaultdict(list) + for keypair in keypairs: + remote_config['public_key'].append(keypair.public_key) + + with runner.isolated_filesystem(), mock.patch.object( + KeyPairsCrud, 'get_first_keypair', return_value=keypairs[0] + ), mock.patch.object( + KeyPairsCrud, 'get_keypairs', return_value=keypairs + ), mock.patch.object( + ConfigsCrud, 'get_deposit_data', return_value=[] + ): + result = runner.invoke(remote_db_group, args) + output = 'Successfully created operator configuration file.\n' + assert output.strip() in result.output.strip() diff --git a/src/remote_db/tests/test_tasks.py b/src/remote_db/tests/test_tasks.py index e0d888f6..725564c8 100644 --- a/src/remote_db/tests/test_tasks.py +++ b/src/remote_db/tests/test_tasks.py @@ -1,390 +1,4 @@ -import base64 -from collections import defaultdict -from pathlib import Path -from secrets import randbits -from typing import Generator -from unittest import mock - -import milagro_bls_binding as bls -import pytest -from click.testing import CliRunner -from eth_typing import ChecksumAddress, HexAddress -from py_ecc.bls import G2ProofOfPossession -from web3 import Web3 - -from src.common.typings import Oracles -from src.remote_db.commands import remote_db_group -from src.remote_db.database import ConfigsCrud, KeyPairsCrud -from src.remote_db.tasks import _encrypt_private_key, _get_key_indexes -from src.remote_db.typings import RemoteDatabaseKeyPair -from src.validators.signing.key_shares import private_key_to_private_key_shares -from src.validators.typings import BLSPrivkey - - -@pytest.fixture -def _patch_check_db_connection() -> Generator: - with mock.patch('src.remote_db.commands.check_db_connection', return_value=None): - yield - - -@pytest.fixture -def _patch_get_db_connection() -> Generator: - with mock.patch('src.remote_db.tasks.get_db_connection'), mock.patch( - 'src.remote_db.database.get_db_connection' - ): - yield - - -@pytest.fixture -def _patch_check_deposit_data_root() -> Generator: - with mock.patch('src.remote_db.tasks.check_deposit_data_root'): - yield - - -def _get_remote_db_keypairs( - mocked_oracles: Oracles, encryption_key: bytes, vault_address: HexAddress -) -> list[RemoteDatabaseKeyPair]: - oracles = mocked_oracles - total_oracles = len(oracles.public_keys) - keystores = {} - for _ in range(3): - seed = randbits(256).to_bytes(32, 'big') - private_key = BLSPrivkey(G2ProofOfPossession.KeyGen(seed).to_bytes(32, 'big')) - public_key = Web3.to_hex(bls.SkToPk(private_key)) - keystores[public_key] = private_key - - key_records: list[RemoteDatabaseKeyPair] = [] - for public_key, private_key in keystores.items(): # pylint: disable=no-member - encrypted_priv_key, nonce = _encrypt_private_key(private_key, encryption_key) - key_records.append( - RemoteDatabaseKeyPair( - vault=ChecksumAddress(vault_address), - public_key=public_key, - private_key=Web3.to_hex(encrypted_priv_key), - nonce=Web3.to_hex(nonce), - ) - ) - # calculate shares for keystore private key - private_key_shares = private_key_to_private_key_shares( - private_key=private_key, - threshold=oracles.exit_signature_recover_threshold, - total=total_oracles, - ) - - # update remote signer config and shares keystores - for share_private_key in private_key_shares: - share_public_key = Web3.to_hex(bls.SkToPk(share_private_key)) - encrypted_priv_key, nonce = _encrypt_private_key(share_private_key, encryption_key) - key_records.append( - RemoteDatabaseKeyPair( - vault=ChecksumAddress(vault_address), - parent_public_key=public_key, - public_key=share_public_key, - private_key=Web3.to_hex(encrypted_priv_key), - nonce=Web3.to_hex(nonce), - ) - ) - return key_records - - -@pytest.mark.usefixtures('_init_vault', '_patch_check_db_connection', '_patch_get_db_connection') -class TestRemoteDbSetup: - def test_setup_works( - self, - data_dir: Path, - vault_address: HexAddress, - runner: CliRunner, - ): - db_url = 'postgresql://user:password@localhost:5432/dbname' - - args = ['--db-url', db_url, '--vault', vault_address, '--data-dir', str(data_dir), 'setup'] - with mock.patch.object( - KeyPairsCrud, 'get_keypairs_count', return_value=0 - ) as get_keypairs_count_mock, mock.patch.object( - ConfigsCrud, 'get_configs_count', return_value=0 - ) as get_configs_count_mock, mock.patch( - 'src.remote_db.tasks.get_random_bytes', return_value=b'1' - ): - result = runner.invoke(remote_db_group, args) - assert get_keypairs_count_mock.call_count == 1 - assert get_configs_count_mock.call_count == 1 - - output = ( - 'Successfully configured remote database.\n' - 'Encryption key: MQ==\n' - 'NB! You must store your encryption in a secure cold storage!' - ) - assert output.strip() == result.output.strip() - - def test_cleanup_works( - self, - data_dir: Path, - vault_address: HexAddress, - runner: CliRunner, - ): - db_url = 'postgresql://user:password@localhost:5432/dbname' - args = [ - '--db-url', - db_url, - '--vault', - vault_address, - '--data-dir', - str(data_dir), - 'cleanup', - ] - with mock.patch.object( - KeyPairsCrud, 'remove_keypairs' - ) as remove_keypairs_mock, mock.patch.object( - ConfigsCrud, 'remove_configs' - ) as remove_configs_mock: - result = runner.invoke(remote_db_group, args) - assert remove_keypairs_mock.call_count == 1 - assert remove_configs_mock.call_count == 1 - - output = 'Successfully removed all the entries for the ' f'{vault_address} vault.' - assert output.strip() == result.output.strip() - - def test_setup_fails_when_keypairs_not_empty( - self, - data_dir: Path, - vault_address: HexAddress, - runner: CliRunner, - ): - db_url = 'postgresql://user:password@localhost:5432/dbname' - - args = ['--db-url', db_url, '--vault', vault_address, '--data-dir', str(data_dir), 'setup'] - with mock.patch.object( - KeyPairsCrud, 'get_keypairs_count', return_value=1 - ) as get_keypairs_count_mock: - result = runner.invoke(remote_db_group, args) - assert get_keypairs_count_mock.call_count == 1 - - output = ( - 'Error: Error: the remote database is not empty. Please clean up with "clean" ' - 'command first.' - ) - assert output.strip() == result.output.strip() - - def test_setup_fails_when_configs_not_empty( - self, - data_dir: Path, - vault_address: HexAddress, - runner: CliRunner, - ): - db_url = 'postgresql://user:password@localhost:5432/dbname' - - args = ['--db-url', db_url, '--vault', vault_address, '--data-dir', str(data_dir), 'setup'] - with ( - mock.patch.object(KeyPairsCrud, 'get_keypairs_count', return_value=0), - mock.patch.object( - ConfigsCrud, 'get_configs_count', return_value=1 - ) as get_configs_count_mock, - ): - result = runner.invoke(remote_db_group, args) - assert get_configs_count_mock.call_count == 1 - - output = ( - 'Error: Error: the remote database is not empty. Please clean up with "clean" ' - 'command first.' - ) - assert output.strip() == result.output.strip() - - -@pytest.mark.usefixtures( - '_patch_check_db_connection', - '_patch_get_db_connection', - '_patch_check_deposit_data_root', -) -@pytest.mark.usefixtures('_init_vault', '_create_keys') -class TestRemoteDbUploadKeypairs: - def test_upload_keypairs_works( - self, - data_dir: Path, - vault_address: HexAddress, - mocked_oracles: Oracles, - runner: CliRunner, - execution_endpoints: str, - ): - db_url = 'postgresql://user:password@localhost:5432/dbname' - encryption_key = '43ueY4nqsiajWHTnkdqrc3OWj2W+t0bbdBISJFjZ3Ck=' - - args = [ - '--db-url', - db_url, - '--vault', - vault_address, - '--data-dir', - str(data_dir), - 'upload-keypairs', - '--execution-endpoints', - execution_endpoints, - '--encrypt-key', - encryption_key, - ] - - with mock.patch.object( - KeyPairsCrud, 'get_first_keypair', return_value=None - ), mock.patch.object(ConfigsCrud, 'get_remote_signer_config', return_value=None): - result = runner.invoke(remote_db_group, args) - output = f'Successfully uploaded keypairs and shares for the {vault_address} vault.' - assert output.strip() in result.output.strip() - - -@pytest.mark.usefixtures( - '_patch_check_db_connection', - '_patch_get_db_connection', -) -@pytest.mark.usefixtures('_init_vault', '_create_keys') -class TestRemoteDbSetupWeb3Signer: - def test_setup_web3signer_works( - self, - data_dir: Path, - vault_address: HexAddress, - mocked_oracles: Oracles, - runner: CliRunner, - execution_endpoints: str, - ): - db_url = 'postgresql://user:password@localhost:5432/dbname' - encryption_key = '43ueY4nqsiajWHTnkdqrc3OWj2W+t0bbdBISJFjZ3Ck=' - - args = [ - '--db-url', - db_url, - '--vault', - vault_address, - '--data-dir', - str(data_dir), - 'setup-web3signer', - '--output-dir', - './web3signer', - '--encrypt-key', - encryption_key, - ] - keypairs = _get_remote_db_keypairs( - mocked_oracles, base64.b64decode(encryption_key), vault_address - ) - - with runner.isolated_filesystem(), mock.patch.object( - KeyPairsCrud, 'get_first_keypair', return_value=keypairs[0] - ), mock.patch.object( - KeyPairsCrud, 'get_keypairs', return_value=keypairs - ), mock.patch.object( - ConfigsCrud, 'get_remote_signer_config', return_value=None - ): - result = runner.invoke(remote_db_group, args) - output = 'Successfully retrieved web3signer private keys from the database.\n' - assert output.strip() in result.output.strip() - - -@pytest.mark.usefixtures( - '_patch_check_db_connection', - '_patch_get_db_connection', -) -@pytest.mark.usefixtures('_init_vault', '_create_keys') -class TestRemoteDbSetupValidator: - def test_setup_validator( - self, - data_dir: Path, - vault_address: HexAddress, - mocked_oracles: Oracles, - runner: CliRunner, - ): - db_url = 'postgresql://user:password@localhost:5432/dbname' - encryption_key = '43ueY4nqsiajWHTnkdqrc3OWj2W+t0bbdBISJFjZ3Ck=' - - args = [ - '--db-url', - db_url, - '--vault', - vault_address, - '--data-dir', - str(data_dir), - 'setup-validator', - '--output-dir', - './validator', - '--validator-index', - '0', - '--total-validators', - '1', - '--web3signer-endpoint', - 'http://localhost:8080', - '--fee-recipient', - vault_address, - ] - keypairs = _get_remote_db_keypairs( - mocked_oracles, base64.b64decode(encryption_key), vault_address - ) - - with runner.isolated_filesystem(), mock.patch.object( - KeyPairsCrud, 'get_first_keypair', return_value=keypairs[0] - ), mock.patch.object( - KeyPairsCrud, 'get_keypairs', return_value=keypairs - ), mock.patch.object( - ConfigsCrud, 'get_remote_signer_config', return_value=None - ): - result = runner.invoke(remote_db_group, args) - output = ( - 'Generated configs with 12 keys ' - 'for validator with index 0.\n' - 'Validator definitions for Lighthouse saved to ' - 'validator/validator_definitions.yml file.\n' - 'Signer keys for Teku\\Prysm saved to ' - 'validator/signer_keys.yml file.\n' - 'Proposer config for Teku\\Prysm saved to ' - 'validator/proposer_config.json file.\n' - ) - assert output.strip() in result.output.strip() - - -@pytest.mark.usefixtures( - '_patch_check_db_connection', - '_patch_get_db_connection', -) -@pytest.mark.usefixtures('_init_vault', '_create_keys') -class TestRemoteDbSetupOperator: - def test_setup_operator( - self, - data_dir: Path, - vault_address: HexAddress, - mocked_oracles: Oracles, - runner: CliRunner, - ): - db_url = 'postgresql://user:password@localhost:5432/dbname' - encryption_key = '43ueY4nqsiajWHTnkdqrc3OWj2W+t0bbdBISJFjZ3Ck=' - - args = [ - '--db-url', - db_url, - '--vault', - vault_address, - '--data-dir', - str(data_dir), - 'setup-operator', - '--output-dir', - './operator', - ] - keypairs = _get_remote_db_keypairs( - mocked_oracles, base64.b64decode(encryption_key), vault_address - ) - remote_config: dict[str, list[str]] = defaultdict(list) - for keypair in keypairs: - if keypair.parent_public_key is None: - remote_config['public_key'] = [] - else: - remote_config['public_key'].append(keypair.public_key) - - with runner.isolated_filesystem(), mock.patch.object( - KeyPairsCrud, 'get_first_keypair', return_value=keypairs[0] - ), mock.patch.object( - KeyPairsCrud, 'get_keypairs', return_value=keypairs - ), mock.patch.object( - ConfigsCrud, 'get_remote_signer_config', return_value=remote_config - ), mock.patch.object( - ConfigsCrud, 'get_deposit_data', return_value=[] - ): - result = runner.invoke(remote_db_group, args) - output = 'Successfully created operator configuration file.\n' - assert output.strip() in result.output.strip() +from src.remote_db.tasks import _get_key_indexes def test_get_key_indexes(): diff --git a/src/remote_db/typings.py b/src/remote_db/typings.py index 1643d8c0..a87fa8a8 100644 --- a/src/remote_db/typings.py +++ b/src/remote_db/typings.py @@ -9,4 +9,3 @@ class RemoteDatabaseKeyPair: public_key: HexStr private_key: HexStr nonce: HexStr - parent_public_key: HexStr | None = None