# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import pytest from datetime import timedelta import pyarrow as pa try: import pyarrow.parquet as pq import pyarrow.parquet.encryption as pe except ImportError: pq = None pe = None else: from pyarrow.tests.parquet.encryption import (InMemoryKmsClient, MockVersioningKmsClient, verify_file_encrypted, read_external_keys_to_dict, parse_wrapped_key) PARQUET_NAME = 'encrypted_table.in_mem.parquet' FOOTER_KEY = b"0123456789112345" FOOTER_KEY_NAME = "footer_key" COL_KEY = b"1234567890123450" COL_KEY_NAME = "col_key" DIRECT_KEY_128 = b"0123456789abcdef" DIRECT_KEY_192 = b"0123456789abcdef01234567" DIRECT_KEY_256 = b"0123456789abcdef0123456789abcdef" DIRECT_AAD_PREFIX = b"test_aad_prefix" KMS_INSTANCE_ID = "kms-123" KMS_INSTANCE_URL = "https://example.com/kms" # Marks all of the tests in this module # Ignore these with pytest ... -m 'not parquet_encryption' # Ignore these with pytest ... -m 'not parquet' pytestmark = [ pytest.mark.parquet_encryption, pytest.mark.parquet ] @pytest.fixture(scope='module') def data_table(): data_table = pa.Table.from_pydict({ 'a': pa.array([1, 2, 3]), 'b': pa.array(['a', 'b', 'c']), 'c': pa.array(['x', 'y', 'z']) }) return data_table @pytest.fixture(scope='module') def basic_encryption_config(): basic_encryption_config = pe.EncryptionConfiguration( footer_key=FOOTER_KEY_NAME, column_keys={ COL_KEY_NAME: ["a", "b"], }) return basic_encryption_config @pytest.fixture(scope='module') def external_encryption_config(): external_encryption_config = pe.EncryptionConfiguration( footer_key=FOOTER_KEY_NAME, column_keys={ COL_KEY_NAME: ["a", "b"], }, internal_key_material=False) return external_encryption_config def setup_encryption_environment(custom_kms_conf): """ Sets up and returns the KMS connection configuration and crypto factory based on provided KMS configuration parameters. """ kms_connection_config = pe.KmsConnectionConfig(custom_kms_conf=custom_kms_conf) def kms_factory(kms_connection_configuration): return InMemoryKmsClient(kms_connection_configuration) # Create our CryptoFactory crypto_factory = pe.CryptoFactory(kms_factory) return kms_connection_config, crypto_factory def write_encrypted_file(path, data_table, footer_key_name, col_key_name, footer_key, col_key, encryption_config): """ Writes an encrypted parquet file based on the provided parameters. """ # Setup the custom KMS configuration with provided keys custom_kms_conf = { footer_key_name: footer_key.decode("UTF-8"), col_key_name: col_key.decode("UTF-8"), } # Setup encryption environment kms_connection_config, crypto_factory = setup_encryption_environment( custom_kms_conf) # Write the encrypted parquet file write_encrypted_parquet(path, data_table, encryption_config, kms_connection_config, crypto_factory) return kms_connection_config, crypto_factory def test_encrypted_parquet_write_read(tempdir, data_table): """Write an encrypted parquet, verify it's encrypted, and then read it.""" path = tempdir / PARQUET_NAME # Encrypt the footer with the footer key, # encrypt column `a` and column `b` with another key, # keep `c` plaintext encryption_config = pe.EncryptionConfiguration( footer_key=FOOTER_KEY_NAME, column_keys={ COL_KEY_NAME: ["a", "b"], }, encryption_algorithm="AES_GCM_V1", cache_lifetime=timedelta(minutes=5.0), data_key_length_bits=256) assert encryption_config.uniform_encryption is False kms_connection_config, crypto_factory = write_encrypted_file( path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, COL_KEY, encryption_config) verify_file_encrypted(path) # Read with decryption properties decryption_config = pe.DecryptionConfiguration( cache_lifetime=timedelta(minutes=5.0)) result_table = read_encrypted_parquet( path, decryption_config, kms_connection_config, crypto_factory) assert data_table.equals(result_table) def test_uniform_encrypted_parquet_write_read(tempdir, data_table): """Write an encrypted parquet, verify it's encrypted, and then read it.""" path = tempdir / PARQUET_NAME # Encrypt the footer and all columns with the footer key, encryption_config = pe.EncryptionConfiguration( footer_key=FOOTER_KEY_NAME, uniform_encryption=True, encryption_algorithm="AES_GCM_V1", cache_lifetime=timedelta(minutes=5.0), data_key_length_bits=256) assert encryption_config.uniform_encryption is True kms_connection_config, crypto_factory = write_encrypted_file( path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, b"", encryption_config) verify_file_encrypted(path) # Read with decryption properties decryption_config = pe.DecryptionConfiguration( cache_lifetime=timedelta(minutes=5.0)) result_table = read_encrypted_parquet( path, decryption_config, kms_connection_config, crypto_factory) assert data_table.equals(result_table) def write_encrypted_parquet(path, table, encryption_config, kms_connection_config, crypto_factory): if encryption_config.internal_key_material: file_encryption_properties = crypto_factory.file_encryption_properties( kms_connection_config, encryption_config) else: file_encryption_properties = crypto_factory.file_encryption_properties( kms_connection_config, encryption_config, path) assert file_encryption_properties is not None with pq.ParquetWriter( path, table.schema, encryption_properties=file_encryption_properties) as writer: writer.write_table(table) def read_encrypted_parquet(path, decryption_config, kms_connection_config, crypto_factory, internal_key_material=True): if internal_key_material: file_decryption_properties = crypto_factory.file_decryption_properties( kms_connection_config, decryption_config) else: file_decryption_properties = crypto_factory.file_decryption_properties( kms_connection_config, decryption_config, path) assert file_decryption_properties is not None meta = pq.read_metadata( path, decryption_properties=file_decryption_properties) assert meta.num_columns == 3 schema = pq.read_schema( path, decryption_properties=file_decryption_properties) assert len(schema.names) == 3 result = pq.ParquetFile( path, decryption_properties=file_decryption_properties) return result.read(use_threads=True) def test_encrypted_parquet_write_read_wrong_key(tempdir, data_table): """Write an encrypted parquet, verify it's encrypted, and then read it using wrong keys.""" path = tempdir / PARQUET_NAME # Encrypt the footer with the footer key, # encrypt column `a` and column `b` with another key, # keep `c` plaintext encryption_config = pe.EncryptionConfiguration( footer_key=FOOTER_KEY_NAME, column_keys={ COL_KEY_NAME: ["a", "b"], }, encryption_algorithm="AES_GCM_V1", cache_lifetime=timedelta(minutes=5.0), data_key_length_bits=256) write_encrypted_file(path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, COL_KEY, encryption_config) verify_file_encrypted(path) wrong_kms_connection_config, wrong_crypto_factory = setup_encryption_environment({ FOOTER_KEY_NAME: COL_KEY.decode("UTF-8"), # Intentionally wrong COL_KEY_NAME: FOOTER_KEY.decode("UTF-8"), # Intentionally wrong }) decryption_config = pe.DecryptionConfiguration( cache_lifetime=timedelta(minutes=5.0)) with pytest.raises(ValueError, match=r"Incorrect master key used"): read_encrypted_parquet( path, decryption_config, wrong_kms_connection_config, wrong_crypto_factory) def test_encrypted_parquet_read_no_decryption_config(tempdir, data_table): """Write an encrypted parquet, verify it's encrypted, but then try to read it without decryption properties.""" test_encrypted_parquet_write_read(tempdir, data_table) # Read without decryption properties with pytest.raises(IOError, match=r"no decryption"): pq.ParquetFile(tempdir / PARQUET_NAME).read() def test_encrypted_parquet_read_metadata_no_decryption_config( tempdir, data_table): """Write an encrypted parquet, verify it's encrypted, but then try to read its metadata without decryption properties.""" test_encrypted_parquet_write_read(tempdir, data_table) # Read metadata without decryption properties with pytest.raises(IOError, match=r"no decryption"): pq.read_metadata(tempdir / PARQUET_NAME) def test_encrypted_parquet_read_schema_no_decryption_config( tempdir, data_table): """Write an encrypted parquet, verify it's encrypted, but then try to read its schema without decryption properties.""" test_encrypted_parquet_write_read(tempdir, data_table) with pytest.raises(IOError, match=r"no decryption"): pq.read_schema(tempdir / PARQUET_NAME) def test_encrypted_parquet_write_no_col_key(tempdir, data_table): """Write an encrypted parquet, but give only footer key, without column key.""" path = tempdir / 'encrypted_table_no_col_key.in_mem.parquet' # Encrypt the footer with the footer key encryption_config = pe.EncryptionConfiguration( footer_key=FOOTER_KEY_NAME) with pytest.raises(OSError, match="Either column_keys or uniform_encryption " "must be set"): # Write with encryption properties write_encrypted_file(path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, b"", encryption_config) def test_encrypted_parquet_write_col_key_and_uniform_encryption(tempdir, data_table): """Write an encrypted parquet, but give only footer key, without column key.""" path = tempdir / 'encrypted_table_col_key_and_uniform_encryption.in_mem.parquet' # Encrypt the footer with the footer key encryption_config = pe.EncryptionConfiguration( footer_key=FOOTER_KEY_NAME, column_keys={ COL_KEY_NAME: ["a", "b"], }, uniform_encryption=True) with pytest.raises(OSError, match=r"Cannot set both column_keys and uniform_encryption"): # Write with encryption properties write_encrypted_file(path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, b"", encryption_config) def test_encrypted_parquet_write_kms_error(tempdir, data_table, basic_encryption_config): """Write an encrypted parquet, but raise KeyError in KmsClient.""" path = tempdir / 'encrypted_table_kms_error.in_mem.parquet' encryption_config = basic_encryption_config # Empty master_keys_map kms_connection_config = pe.KmsConnectionConfig() def kms_factory(kms_connection_configuration): # Empty master keys map will cause KeyError to be raised # on wrap/unwrap calls return InMemoryKmsClient(kms_connection_configuration) crypto_factory = pe.CryptoFactory(kms_factory) with pytest.raises(KeyError, match="footer_key"): # Write with encryption properties write_encrypted_parquet(path, data_table, encryption_config, kms_connection_config, crypto_factory) def test_encrypted_parquet_write_kms_specific_error(tempdir, data_table, basic_encryption_config): """Write an encrypted parquet, but raise KeyError in KmsClient.""" path = tempdir / 'encrypted_table_kms_error.in_mem.parquet' encryption_config = basic_encryption_config # Empty master_keys_map kms_connection_config = pe.KmsConnectionConfig() class ThrowingKmsClient(pe.KmsClient): """A KmsClient implementation that throws exception in wrap/unwrap calls """ def __init__(self, config): """Create an InMemoryKmsClient instance.""" pe.KmsClient.__init__(self) self.config = config def wrap_key(self, key_bytes, master_key_identifier): raise ValueError("Cannot Wrap Key") def unwrap_key(self, wrapped_key, master_key_identifier): raise ValueError("Cannot Unwrap Key") def kms_factory(kms_connection_configuration): # Exception thrown in wrap/unwrap calls return ThrowingKmsClient(kms_connection_configuration) crypto_factory = pe.CryptoFactory(kms_factory) with pytest.raises(ValueError, match="Cannot Wrap Key"): # Write with encryption properties write_encrypted_parquet(path, data_table, encryption_config, kms_connection_config, crypto_factory) def test_encrypted_parquet_write_kms_factory_error(tempdir, data_table, basic_encryption_config): """Write an encrypted parquet, but raise ValueError in kms_factory.""" path = tempdir / 'encrypted_table_kms_factory_error.in_mem.parquet' encryption_config = basic_encryption_config # Empty master_keys_map kms_connection_config = pe.KmsConnectionConfig() def kms_factory(kms_connection_configuration): raise ValueError('Cannot create KmsClient') crypto_factory = pe.CryptoFactory(kms_factory) with pytest.raises(ValueError, match="Cannot create KmsClient"): # Write with encryption properties write_encrypted_parquet(path, data_table, encryption_config, kms_connection_config, crypto_factory) def test_encrypted_parquet_write_kms_factory_type_error( tempdir, data_table, basic_encryption_config): """Write an encrypted parquet, but use wrong KMS client type that doesn't implement KmsClient.""" path = tempdir / 'encrypted_table_kms_factory_error.in_mem.parquet' encryption_config = basic_encryption_config # Empty master_keys_map kms_connection_config = pe.KmsConnectionConfig() class WrongTypeKmsClient(): """This is not an implementation of KmsClient. """ def __init__(self, config): self.master_keys_map = config.custom_kms_conf def wrap_key(self, key_bytes, master_key_identifier): return None def unwrap_key(self, wrapped_key, master_key_identifier): return None def kms_factory(kms_connection_configuration): return WrongTypeKmsClient(kms_connection_configuration) crypto_factory = pe.CryptoFactory(kms_factory) with pytest.raises(TypeError): # Write with encryption properties write_encrypted_parquet(path, data_table, encryption_config, kms_connection_config, crypto_factory) def test_encrypted_parquet_encryption_configuration(): def validate_encryption_configuration(encryption_config): assert FOOTER_KEY_NAME == encryption_config.footer_key assert ["a", "b"] == encryption_config.column_keys[COL_KEY_NAME] assert "AES_GCM_CTR_V1" == encryption_config.encryption_algorithm assert encryption_config.plaintext_footer assert not encryption_config.double_wrapping assert timedelta(minutes=10.0) == encryption_config.cache_lifetime assert not encryption_config.internal_key_material assert 192 == encryption_config.data_key_length_bits encryption_config = pe.EncryptionConfiguration( footer_key=FOOTER_KEY_NAME, column_keys={COL_KEY_NAME: ["a", "b"], }, encryption_algorithm="AES_GCM_CTR_V1", plaintext_footer=True, double_wrapping=False, cache_lifetime=timedelta(minutes=10.0), internal_key_material=False, data_key_length_bits=192, ) validate_encryption_configuration(encryption_config) encryption_config_1 = pe.EncryptionConfiguration( footer_key=FOOTER_KEY_NAME) encryption_config_1.column_keys = {COL_KEY_NAME: ["a", "b"], } encryption_config_1.encryption_algorithm = "AES_GCM_CTR_V1" encryption_config_1.plaintext_footer = True encryption_config_1.double_wrapping = False encryption_config_1.cache_lifetime = timedelta(minutes=10.0) encryption_config_1.internal_key_material = False encryption_config_1.data_key_length_bits = 192 validate_encryption_configuration(encryption_config_1) def test_encrypted_parquet_decryption_configuration(): # Test defaults default_config = pe.DecryptionConfiguration() assert timedelta(minutes=10.0) == default_config.cache_lifetime assert default_config.read_kms_url is False # Test init parameters decryption_config = pe.DecryptionConfiguration( cache_lifetime=timedelta(minutes=5.0), read_kms_url=True) assert timedelta(minutes=5.0) == decryption_config.cache_lifetime assert decryption_config.read_kms_url is True # Test setters decryption_config_1 = pe.DecryptionConfiguration() decryption_config_1.cache_lifetime = timedelta(minutes=5.0) decryption_config_1.read_kms_url = True assert timedelta(minutes=5.0) == decryption_config_1.cache_lifetime assert decryption_config_1.read_kms_url is True # Can pass integer number of seconds as cache lifetime decryption_config_2 = pe.DecryptionConfiguration(cache_lifetime=300) assert timedelta(minutes=5.0) == decryption_config_2.cache_lifetime def test_encrypted_parquet_kms_configuration(): def validate_kms_connection_config(kms_connection_config): assert "Instance1" == kms_connection_config.kms_instance_id assert "URL1" == kms_connection_config.kms_instance_url assert "MyToken" == kms_connection_config.key_access_token assert ({"key1": "key_material_1", "key2": "key_material_2"} == kms_connection_config.custom_kms_conf) kms_connection_config = pe.KmsConnectionConfig( kms_instance_id="Instance1", kms_instance_url="URL1", key_access_token="MyToken", custom_kms_conf={ "key1": "key_material_1", "key2": "key_material_2", }) validate_kms_connection_config(kms_connection_config) kms_connection_config_1 = pe.KmsConnectionConfig() kms_connection_config_1.kms_instance_id = "Instance1" kms_connection_config_1.kms_instance_url = "URL1" kms_connection_config_1.key_access_token = "MyToken" kms_connection_config_1.custom_kms_conf = { "key1": "key_material_1", "key2": "key_material_2", } validate_kms_connection_config(kms_connection_config_1) def test_encrypted_parquet_write_read_plain_footer_single_wrapping( tempdir, data_table): """ Write an encrypted parquet, with plaintext footer and with single wrapping, verify it's encrypted, and then read plaintext columns. Runs once with a flat schema and once where the encrypted column `b` is itself a nested (struct) field. """ path = tempdir / PARQUET_NAME for nested in [False, True]: if nested: table = pa.Table.from_pydict({ 'a': pa.array([1, 2, 3]), 'b': pa.array( [{'x': 1, 'y': 2}, {'x': 3, 'y': 4}, {'x': 5, 'y': 6}], type=pa.struct([('x', pa.int32()), ('y', pa.int32())])), 'c': pa.array(['x', 'y', 'z']) }) else: table = data_table # Encrypt the footer with the footer key, # encrypt column `a` and column `b` with another key, keep `c` plaintext encryption_config = pe.EncryptionConfiguration( footer_key=FOOTER_KEY_NAME, column_keys={ COL_KEY_NAME: ["a", "b"], }, plaintext_footer=True, double_wrapping=False) kms_connection_config = pe.KmsConnectionConfig( custom_kms_conf={ FOOTER_KEY_NAME: FOOTER_KEY.decode("UTF-8"), COL_KEY_NAME: COL_KEY.decode("UTF-8"), } ) def kms_factory(kms_connection_configuration): return InMemoryKmsClient(kms_connection_configuration) crypto_factory = pe.CryptoFactory(kms_factory) # Write with encryption properties write_encrypted_parquet(path, table, encryption_config, kms_connection_config, crypto_factory) # Read the plaintext column without decryption properties with pq.ParquetFile(path) as result: result_table = result.read(columns='c', use_threads=False) assert table.num_rows == result_table.num_rows assert table.select(['c']).equals(result_table) def test_encrypted_parquet_write_read_external(tempdir, data_table, external_encryption_config): """Write an encrypted parquet file with external key material, verify it's encrypted, then read both the table and external store. """ path = tempdir / PARQUET_NAME kms_connection_config, crypto_factory = write_encrypted_file( path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, COL_KEY, external_encryption_config) verify_file_encrypted(path) decryption_config = pe.DecryptionConfiguration() result_table = read_encrypted_parquet( path, decryption_config, kms_connection_config, crypto_factory, internal_key_material=False) store = pa._parquet_encryption.FileSystemKeyMaterialStore.for_file(path) assert len(key_ids := store.get_key_id_set()) == ( len(external_encryption_config.column_keys[COL_KEY_NAME]) + 1) assert all([store.get_key_material(k) is not None for k in key_ids]) assert data_table.equals(result_table) @pytest.mark.parametrize( ("double_wrap_initial", "double_wrap_rotated"), [ pytest.param(True, True, id="double wrapping"), pytest.param(False, True, id="single to double wrapped"), pytest.param(True, False, id="double to singe wrapped"), pytest.param(False, False, id="single wrapping")]) def test_external_key_material_rotation( reusable_tempdir, data_table, double_wrap_initial, double_wrap_rotated): """Tests CryptoFactory.rotate_master_keys Note: The CryptoFactory.rotate_master_keys() double_wrapping keword arg may be either True (the default) or False regardless of whether EncryptionConfig.double_wrapping was set to true (also the default) when the external key material store was written. This means double wrapping may be set one way initially and then applied or removed during rotation. """ path = reusable_tempdir / PARQUET_NAME encryption_config = pe.EncryptionConfiguration( footer_key=FOOTER_KEY_NAME, column_keys={COL_KEY_NAME: ["a", "b"]}, internal_key_material=False, double_wrapping=double_wrap_initial) # initial master key version - see MockVersioningKmsClient docstring kms_connection_config = pe.KmsConnectionConfig(key_access_token="1") def kms_factory(kms_connection_configuration): return MockVersioningKmsClient(kms_connection_configuration) crypto_factory = pe.CryptoFactory(kms_factory) write_encrypted_parquet( path, data_table, encryption_config, kms_connection_config, crypto_factory) before_keys = read_external_keys_to_dict(path) # "rotate" kms master key kms_connection_config.refresh_key_access_token("2") crypto_factory.rotate_master_keys( kms_connection_config, path, double_wrapping=double_wrap_rotated) after_keys = read_external_keys_to_dict(path) verify_file_encrypted(path) table_read_after_rotation = read_encrypted_parquet( path, pe.DecryptionConfiguration(), kms_connection_config, crypto_factory, internal_key_material=False) assert FOOTER_KEY_NAME in before_keys assert COL_KEY_NAME in before_keys assert FOOTER_KEY_NAME in after_keys assert COL_KEY_NAME in after_keys def check_rotated_external_keys(master_key_id: str) -> None: before_key_mat = before_keys[master_key_id] if double_wrap_initial: before_key_wrapped = before_key_mat.wrapped_kek else: before_key_wrapped = before_key_mat.wrapped_dek _, before_ver, _ = parse_wrapped_key(before_key_wrapped) after_key_mat = after_keys[master_key_id] if double_wrap_rotated: after_key_wrapped = after_key_mat.wrapped_kek else: after_key_wrapped = after_key_mat.wrapped_dek _, after_ver, _ = parse_wrapped_key(after_key_wrapped) # CryptoFactory rewrapped keys if after version is later than before assert before_ver < after_ver check_rotated_external_keys(FOOTER_KEY_NAME) check_rotated_external_keys(COL_KEY_NAME) assert data_table.equals(table_read_after_rotation) def recording_kms_factory(created_configs, client_class=None): """Create a KMS client factory that appends the KMS instance ID and URL of each connection configuration it is given to created_configs""" if client_class is None: client_class = InMemoryKmsClient def kms_factory(kms_connection_configuration): created_configs.append( (kms_connection_configuration.kms_instance_id, kms_connection_configuration.kms_instance_url)) return client_class(kms_connection_configuration) return kms_factory @pytest.mark.parametrize("read_kms_url", [False, True]) def test_read_kms_url_from_file( tempdir, data_table, basic_encryption_config, read_kms_url): """Read a file written with KMS connection properties configured, using a KmsConnectionConfig that doesn't specify them""" path = tempdir / PARQUET_NAME custom_kms_conf = { FOOTER_KEY_NAME: FOOTER_KEY.decode("UTF-8"), COL_KEY_NAME: COL_KEY.decode("UTF-8"), } write_config = pe.KmsConnectionConfig( kms_instance_id=KMS_INSTANCE_ID, kms_instance_url=KMS_INSTANCE_URL, custom_kms_conf=custom_kms_conf) write_crypto_factory = pe.CryptoFactory(InMemoryKmsClient) write_encrypted_parquet(path, data_table, basic_encryption_config, write_config, write_crypto_factory) verify_file_encrypted(path) # Leave the KMS instance ID and URL unset when reading read_config = pe.KmsConnectionConfig(custom_kms_conf=custom_kms_conf) created_configs = [] read_crypto_factory = pe.CryptoFactory( recording_kms_factory(created_configs)) decryption_config = pe.DecryptionConfiguration( read_kms_url=read_kms_url) result_table = read_encrypted_parquet( path, decryption_config, read_config, read_crypto_factory) assert data_table.equals(result_table) if read_kms_url: # The URL is read from the file key material assert created_configs == [(KMS_INSTANCE_ID, KMS_INSTANCE_URL)] else: # The URL in the key material is ignored and the default provided # instead. assert created_configs == [(KMS_INSTANCE_ID, "DEFAULT")] @pytest.mark.parametrize("read_kms_url", [False, True]) def test_key_rotation_reads_kms_url_from_file(reusable_tempdir, data_table, read_kms_url): """Rotate the keys of a file written with KMS connection properties configured, using a KmsConnectionConfig that doesn't specify them""" path = reusable_tempdir / PARQUET_NAME encryption_config = pe.EncryptionConfiguration( footer_key=FOOTER_KEY_NAME, column_keys={COL_KEY_NAME: ["a", "b"]}, internal_key_material=False) # Write initial encrypted file with external key material write_config = pe.KmsConnectionConfig( kms_instance_id=KMS_INSTANCE_ID, kms_instance_url=KMS_INSTANCE_URL, key_access_token="1") write_crypto_factory = pe.CryptoFactory(MockVersioningKmsClient) write_encrypted_parquet(path, data_table, encryption_config, write_config, write_crypto_factory) # Rotate keys without specifying the KMS instance ID and URL rotation_config = pe.KmsConnectionConfig(key_access_token="2") created_configs = [] rotation_crypto_factory = pe.CryptoFactory( recording_kms_factory(created_configs, MockVersioningKmsClient)) rotation_crypto_factory.rotate_master_keys( rotation_config, path, read_kms_url=read_kms_url) if read_kms_url: # The empty config provided is used to wrap new keys, # and the config from the file was used to unwrap the original keys. assert created_configs == [("", ""), (KMS_INSTANCE_ID, KMS_INSTANCE_URL)] else: # The default URL is used when unwrapping the original keys. assert created_configs == [("", ""), (KMS_INSTANCE_ID, "DEFAULT")] # New key material always uses the provided KMS connection configuration rotated_keys = read_external_keys_to_dict(path) footer_key_material = rotated_keys[FOOTER_KEY_NAME] assert footer_key_material.kms_instance_id == "DEFAULT" assert footer_key_material.kms_instance_url == "DEFAULT" table_read_after_rotation = read_encrypted_parquet( path, pe.DecryptionConfiguration(), rotation_config, rotation_crypto_factory, internal_key_material=False) assert data_table.equals(table_read_after_rotation) def test_encrypted_parquet_loop(tempdir, data_table, basic_encryption_config): """Write an encrypted parquet, verify it's encrypted, and then read it multithreaded in a loop.""" path = tempdir / PARQUET_NAME # Encrypt the footer with the footer key, # encrypt column `a` and column `b` with another key, # keep `c` plaintext, defined in basic_encryption_config kms_connection_config, crypto_factory = write_encrypted_file( path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, COL_KEY, basic_encryption_config) verify_file_encrypted(path) decryption_config = pe.DecryptionConfiguration( cache_lifetime=timedelta(minutes=5.0)) for i in range(50): # Read with decryption properties file_decryption_properties = crypto_factory.file_decryption_properties( kms_connection_config, decryption_config) assert file_decryption_properties is not None result = pq.ParquetFile( path, decryption_properties=file_decryption_properties) result_table = result.read(use_threads=True) assert data_table.equals(result_table) def test_read_with_deleted_crypto_factory(tempdir, data_table, basic_encryption_config): """ Test that decryption properties can be used if the crypto factory is no longer alive """ path = tempdir / PARQUET_NAME kms_connection_config, crypto_factory = write_encrypted_file( path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, COL_KEY, basic_encryption_config) verify_file_encrypted(path) # Create decryption properties and delete the crypto factory that created # the properties afterwards. decryption_config = pe.DecryptionConfiguration( cache_lifetime=timedelta(minutes=5.0)) file_decryption_properties = crypto_factory.file_decryption_properties( kms_connection_config, decryption_config) del crypto_factory result = pq.ParquetFile( path, decryption_properties=file_decryption_properties) result_table = result.read(use_threads=True) assert data_table.equals(result_table) def test_encrypted_parquet_read_table(tempdir, data_table, basic_encryption_config): """Write an encrypted parquet then read it back using read_table.""" path = tempdir / PARQUET_NAME # Write the encrypted parquet file using the utility function kms_connection_config, crypto_factory = write_encrypted_file( path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, COL_KEY, basic_encryption_config) decryption_config = pe.DecryptionConfiguration( cache_lifetime=timedelta(minutes=5.0)) file_decryption_properties = crypto_factory.file_decryption_properties( kms_connection_config, decryption_config) # Read the encrypted parquet file using read_table result_table = pq.read_table(path, decryption_properties=file_decryption_properties) # Assert that the read table matches the original data assert data_table.equals(result_table) # Read the encrypted parquet folder using read_table result_table = pq.read_table( tempdir, decryption_properties=file_decryption_properties) assert data_table.equals(result_table) class TestDirectKeyEncryption: """Tests for create_encryption_properties / create_decryption_properties.""" @pytest.mark.parametrize("key", [ DIRECT_KEY_128, DIRECT_KEY_192, DIRECT_KEY_256, ], ids=["aes128", "aes192", "aes256"]) def test_roundtrip_key_sizes(self, tempdir, data_table, key): path = tempdir / f"direct_{len(key) * 8}.parquet" enc_props = pe.create_encryption_properties(footer_key=key) pq.write_table(data_table, path, encryption_properties=enc_props) dec_props = pe.create_decryption_properties(footer_key=key) result = pq.read_table(path, decryption_properties=dec_props) assert data_table.equals(result) def test_roundtrip_with_aad_prefix(self, tempdir, data_table): path = tempdir / "direct_aad.parquet" enc_props = pe.create_encryption_properties( footer_key=DIRECT_KEY_128, aad_prefix=DIRECT_AAD_PREFIX, ) pq.write_table(data_table, path, encryption_properties=enc_props) dec_props = pe.create_decryption_properties( footer_key=DIRECT_KEY_128, aad_prefix=DIRECT_AAD_PREFIX, ) result = pq.read_table(path, decryption_properties=dec_props) assert data_table.equals(result) def test_roundtrip_aad_prefix_not_stored(self, tempdir, data_table): """When store_aad_prefix=False, reader must supply aad_prefix.""" path = tempdir / "direct_aad_not_stored.parquet" enc_props = pe.create_encryption_properties( footer_key=DIRECT_KEY_128, aad_prefix=DIRECT_AAD_PREFIX, store_aad_prefix=False, ) pq.write_table(data_table, path, encryption_properties=enc_props) # Reading without aad_prefix should fail dec_props_no_aad = pe.create_decryption_properties( footer_key=DIRECT_KEY_128, ) with pytest.raises(IOError, match="AAD"): pq.read_table(path, decryption_properties=dec_props_no_aad) # Reading with correct aad_prefix should succeed dec_props = pe.create_decryption_properties( footer_key=DIRECT_KEY_128, aad_prefix=DIRECT_AAD_PREFIX, ) result = pq.read_table(path, decryption_properties=dec_props) assert data_table.equals(result) def test_wrong_aad_prefix_fails(self, tempdir, data_table): path = tempdir / "direct_wrong_aad.parquet" enc_props = pe.create_encryption_properties( footer_key=DIRECT_KEY_128, aad_prefix=DIRECT_AAD_PREFIX, ) pq.write_table(data_table, path, encryption_properties=enc_props) dec_props = pe.create_decryption_properties( footer_key=DIRECT_KEY_128, aad_prefix=b"wrong_prefix", ) with pytest.raises(IOError, match="AAD"): pq.read_table(path, decryption_properties=dec_props) def test_encrypted_file_has_pare_magic(self, tempdir, data_table): path = tempdir / "direct_magic.parquet" enc_props = pe.create_encryption_properties( footer_key=DIRECT_KEY_128) pq.write_table(data_table, path, encryption_properties=enc_props) with open(path, "rb") as f: magic = f.read(4) assert magic == b"PARE" def test_plaintext_footer(self, tempdir, data_table): path = tempdir / "direct_plaintext_footer.parquet" enc_props = pe.create_encryption_properties( footer_key=DIRECT_KEY_128, plaintext_footer=True, ) pq.write_table(data_table, path, encryption_properties=enc_props) dec_props = pe.create_decryption_properties( footer_key=DIRECT_KEY_128) result = pq.read_table(path, decryption_properties=dec_props) assert data_table.equals(result) def test_aes_gcm_ctr_v1_algorithm(self, tempdir, data_table): path = tempdir / "direct_ctr.parquet" enc_props = pe.create_encryption_properties( footer_key=DIRECT_KEY_128, encryption_algorithm="AES_GCM_CTR_V1", ) pq.write_table(data_table, path, encryption_properties=enc_props) dec_props = pe.create_decryption_properties( footer_key=DIRECT_KEY_128) result = pq.read_table(path, decryption_properties=dec_props) assert data_table.equals(result) def test_wrong_key_fails(self, tempdir, data_table): path = tempdir / "direct_wrong_key.parquet" enc_props = pe.create_encryption_properties( footer_key=DIRECT_KEY_128) pq.write_table(data_table, path, encryption_properties=enc_props) wrong_key = b"fedcba9876543210" dec_props = pe.create_decryption_properties(footer_key=wrong_key) with pytest.raises(IOError, match="decrypt"): pq.read_table(path, decryption_properties=dec_props) def test_reading_without_decryption_fails(self, tempdir, data_table): path = tempdir / "direct_no_decrypt.parquet" enc_props = pe.create_encryption_properties( footer_key=DIRECT_KEY_128) pq.write_table(data_table, path, encryption_properties=enc_props) with pytest.raises(IOError, match="encrypted metadata"): pq.read_table(path) def test_allow_plaintext_files(self, tempdir, data_table): """Plaintext file reads should work when allow_plaintext_files=True.""" path = tempdir / "plaintext.parquet" pq.write_table(data_table, path) dec_props = pe.create_decryption_properties( footer_key=DIRECT_KEY_128, allow_plaintext_files=True, ) result = pq.read_table(path, decryption_properties=dec_props) assert data_table.equals(result) def test_plaintext_file_rejected_by_default(self, tempdir, data_table): """Default allow_plaintext_files=False should reject plaintext files.""" path = tempdir / "plaintext_rejected.parquet" pq.write_table(data_table, path) dec_props = pe.create_decryption_properties( footer_key=DIRECT_KEY_128) with pytest.raises(IOError, match="plaintext"): pq.read_table(path, decryption_properties=dec_props) def test_check_footer_integrity_false(self, tempdir, data_table): """check_footer_integrity=False should still allow decryption.""" path = tempdir / "direct_no_footer_check.parquet" enc_props = pe.create_encryption_properties( footer_key=DIRECT_KEY_128) pq.write_table(data_table, path, encryption_properties=enc_props) dec_props = pe.create_decryption_properties( footer_key=DIRECT_KEY_128, check_footer_integrity=False, ) result = pq.read_table(path, decryption_properties=dec_props) assert data_table.equals(result) def test_plaintext_footer_has_par1_magic(self, tempdir, data_table): """plaintext_footer=True should produce PAR1 magic, not PARE.""" path = tempdir / "direct_plaintext_magic.parquet" enc_props = pe.create_encryption_properties( footer_key=DIRECT_KEY_128, plaintext_footer=True, ) pq.write_table(data_table, path, encryption_properties=enc_props) with open(path, "rb") as f: magic = f.read(4) assert magic == b"PAR1" def test_invalid_key_length_raises(self): with pytest.raises(ValueError, match="16, 24, or 32 bytes"): pe.create_encryption_properties(footer_key=b"short") with pytest.raises(ValueError, match="16, 24, or 32 bytes"): pe.create_encryption_properties(footer_key=b"") with pytest.raises(ValueError, match="16, 24, or 32 bytes"): pe.create_decryption_properties(footer_key=b"short") def test_invalid_algorithm_raises(self): with pytest.raises(ValueError, match="Invalid cipher name"): pe.create_encryption_properties( footer_key=DIRECT_KEY_128, encryption_algorithm="INVALID", ) def test_footer_key_rejects_non_bytes(self): with pytest.raises(TypeError, match="footer_key must be bytes"): pe.create_encryption_properties(footer_key="0123456789abcdef") with pytest.raises(TypeError, match="footer_key must be bytes"): pe.create_decryption_properties(footer_key="0123456789abcdef") with pytest.raises(TypeError, match="footer_key must be bytes"): pe.create_encryption_properties(footer_key=None) def test_aad_prefix_rejects_str(self, tempdir, data_table): with pytest.raises(TypeError, match="aad_prefix must be bytes"): pe.create_encryption_properties( footer_key=DIRECT_KEY_128, aad_prefix="not_bytes", )