Merge pull request #33 from tuxxy/aead-dem

Implement cryptography.io's AEAD for ChaCha20-Poly1305 in the UmbralDEM

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')

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