diff --git a/tests/test_dem.py b/tests/test_dem.py new file mode 100644 index 0000000..81be81a --- /dev/null +++ b/tests/test_dem.py @@ -0,0 +1,60 @@ +import pytest +import os + +from umbral.dem import UmbralDEM, DEM_KEYSIZE, DEM_NONCE_SIZE +from cryptography.exceptions import InvalidTag + + +def test_encrypt_decrypt(): + key = os.urandom(32) + + dem = UmbralDEM(key) + + plaintext = b'attack at dawn' + + ciphertext0 = dem.encrypt(plaintext) + ciphertext1 = dem.encrypt(plaintext) + + assert ciphertext0 != plaintext + assert ciphertext1 != plaintext + + # Ciphertext should be different even with same plaintext. + assert ciphertext0 != ciphertext1 + + # Nonce should be different + assert ciphertext0[:DEM_NONCE_SIZE] != ciphertext1[:DEM_NONCE_SIZE] + + cleartext0 = dem.decrypt(ciphertext0) + cleartext1 = dem.decrypt(ciphertext1) + + assert cleartext0 == plaintext + assert cleartext1 == plaintext + + +def test_encrypt_decrypt_associated_data(): + key = os.urandom(32) + aad = b'launch code 0000' + + dem = UmbralDEM(key) + + plaintext = b'attack at dawn' + + ciphertext0 = dem.encrypt(plaintext, authenticated_data=aad) + ciphertext1 = dem.encrypt(plaintext, authenticated_data=aad) + + assert ciphertext0 != plaintext + assert ciphertext1 != plaintext + + assert ciphertext0 != ciphertext1 + + assert ciphertext0[:DEM_NONCE_SIZE] != ciphertext1[:DEM_NONCE_SIZE] + + cleartext0 = dem.decrypt(ciphertext0, authenticated_data=aad) + cleartext1 = dem.decrypt(ciphertext1, authenticated_data=aad) + + assert cleartext0 == plaintext + assert cleartext1 == plaintext + + # Attempt decryption with invalid associated data + with pytest.raises(InvalidTag) as err_info: + cleartext2 = dem.decrypt(ciphertext0, authenticated_data=b'wrong data') diff --git a/umbral/dem.py b/umbral/dem.py index 335f6c5..b1e7432 100644 --- a/umbral/dem.py +++ b/umbral/dem.py @@ -1,29 +1,39 @@ -from nacl.secret import SecretBox +import os +from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305 + + +DEM_KEYSIZE = 32 +DEM_NONCE_SIZE = 12 class UmbralDEM(object): def __init__(self, symm_key: bytes): """ Initializes an UmbralDEM object. Requires a key to perform - Salsa20-Poly1305. + ChaCha20-Poly1305. """ - if len(symm_key) != SecretBox.KEY_SIZE: + if len(symm_key) != DEM_KEYSIZE: raise ValueError( - "Invalid key size, must be {} bytes".format(SecretBox.KEY_SIZE) + "Invalid key size, must be {} bytes".format(DEM_KEYSIZE) ) - self.cipher = SecretBox(symm_key) + self.cipher = ChaCha20Poly1305(symm_key) - def encrypt(self, data: bytes): + def encrypt(self, data: bytes, authenticated_data: bytes=None): """ - Encrypts data using NaCl's Salsa20-Poly1305 secret box symmetric cipher. + Encrypts data using ChaCha20-Poly1305 with optional authenticated data. """ - enc_data = self.cipher.encrypt(data) - return enc_data + nonce = os.urandom(DEM_NONCE_SIZE) + enc_data = self.cipher.encrypt(nonce, data, authenticated_data) + # Ciphertext will be a 12 byte nonce, the ciphertext, and a 16 byte tag. + return nonce + enc_data - def decrypt(self, enc_data: bytes): + def decrypt(self, enc_data: bytes, authenticated_data: bytes=None): """ - Decrypts data using NaCl's Salsa20-Poly1305 secret box symmetric cipher. + Decrypts data using ChaCha20-Poly1305 and validates the provided + authenticated data. """ - plaintext = self.cipher.decrypt(enc_data) + nonce = enc_data[:DEM_NONCE_SIZE] + ciphertext = enc_data[DEM_NONCE_SIZE:] + plaintext = self.cipher.decrypt(nonce, ciphertext, authenticated_data) return plaintext