import tempfile
from pathlib import Path
import pytest
from self_encryption import (
PyDataMap,
encrypt,
encrypt_from_file,
decrypt,
decrypt_from_storage,
streaming_decrypt_from_storage,
)
def test_direct_encryption_decryption():
data = b"Hello, World!" * 1000
data_map, chunks = encrypt(data)
decrypted = decrypt(data_map, chunks)
assert data == decrypted
def test_file_encryption_decryption():
with tempfile.TemporaryDirectory() as temp_dir:
input_path = Path(temp_dir) / "input.dat"
data = b"x" * 1_000_000 input_path.write_bytes(data)
chunk_dir = Path(temp_dir) / "chunks"
chunk_dir.mkdir()
result = encrypt_from_file(str(input_path), str(chunk_dir))
data_map, chunk_names = result
def get_chunk(chunk_name: str) -> bytes:
chunk_path = Path(chunk_dir) / chunk_name
return chunk_path.read_bytes()
output_path = Path(temp_dir) / "output.dat"
decrypt_from_storage(data_map, str(output_path), get_chunk)
assert input_path.read_bytes() == output_path.read_bytes()
def test_streaming_decryption():
with tempfile.TemporaryDirectory() as temp_dir:
input_path = Path(temp_dir) / "input.dat"
data = b"x" * 1_000_000 input_path.write_bytes(data)
chunk_dir = Path(temp_dir) / "chunks"
chunk_dir.mkdir()
result = encrypt_from_file(str(input_path), str(chunk_dir))
data_map, chunk_names = result
def get_chunks(chunk_names: list) -> list:
return [
(index, (Path(chunk_dir) / name).read_bytes())
for index, name in chunk_names
]
output_path = Path(temp_dir) / "output.dat"
streaming_decrypt_from_storage(data_map, str(output_path), get_chunks)
assert input_path.read_bytes() == output_path.read_bytes()
def test_data_map_json():
chunk_infos = []
data_map = PyDataMap(chunk_infos)
json_str = data_map.to_json()
data_map2 = PyDataMap.from_json(json_str)
assert data_map.to_json() == data_map2.to_json()
if __name__ == "__main__":
pytest.main([__file__])