From 26ad9c25a5abd18da96c554825caa370337cd0ac Mon Sep 17 00:00:00 2001
From: Rubidium <rubidium@openttd.org>
Date: Tue, 26 Mar 2024 19:54:56 +0100
Subject: [PATCH] Add Python class to access incremental aead interface

---
 c_monocypher.pyx        | 52 +++++++++++++++++++++++++++++++++++++++++
 test/test_monocypher.py | 25 ++++++++++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/c_monocypher.pyx b/c_monocypher.pyx
index 77c52f1..1f3558b 100644
--- a/c_monocypher.pyx
+++ b/c_monocypher.pyx
@@ -203,6 +203,58 @@ def unlock(key, nonce, mac, message, associated_data=None):
     return plain_text
 
 
+cdef class IncrementalAuthenticatedEncryption:
+    cdef crypto_aead_ctx _ctx
+
+    """Instantiate the incremental authenticated encryption handler.
+
+    :param key: The 32-byte shared session key.
+    :param nonce: The 24-byte number, used only once with any given session key.
+    """
+    def __init__(self, key, nonce):
+        if len(key) != 32:
+            raise ValueError(f'Invalid key length {len(key)} != 32')
+
+        if len(nonce) != 24:
+            raise ValueError(f'Invalid nonce length {len(key)} != 24')
+
+        crypto_aead_init_x(&self._ctx, key, nonce)
+
+    def lock(self, message, associated_data=None):
+        """Perform authenticated encryption.
+
+        :param message: The secret message to encrypt.
+        :param associated_data: The additional data to authenticate which
+            is NOT encrypted.
+        :return: the tuple of (MAC, ciphertext).  MAC is the 16-byte message
+            authentication code.  ciphertext is the encrypted message.
+        """
+        mac = bytes(16)
+        crypto_text = bytes(len(message))
+        associated_data = b'' if associated_data is None else associated_data
+        crypto_aead_write(&self._ctx, crypto_text, mac, associated_data, len(associated_data), message, len(message))
+        return mac, crypto_text
+
+    def unlock(self, mac, message, associated_data=None):
+        """Perform authenticated decryption.
+
+        :param mac: The 16-byte message authentication code produced by :func:`lock`.
+        :param message: The ciphertext encrypted message to decrypt produced by :func:`lock`.
+        :param associated_data: The additional data to authenticate which
+            is NOT encrypted.
+        :return: The secret message or None on authentication failure.
+        """
+        if len(mac) != 16:
+            raise ValueError(f'Invalid mac length {len(mac)} != 16')
+
+        plain_text = bytearray(len(message))
+        associated_data = b'' if associated_data is None else associated_data
+        rv = crypto_aead_read(&self._ctx, plain_text, mac, associated_data, len(associated_data), message, len(message))
+        if 0 != rv:
+            return None
+        return plain_text
+
+
 def chacha20(key, nonce, message):
     """Encrypt/Decrypt a message with ChaCha20.
 
diff --git a/test/test_monocypher.py b/test/test_monocypher.py
index 2741649..d3ea5a7 100644
--- a/test/test_monocypher.py
+++ b/test/test_monocypher.py
@@ -75,6 +75,31 @@ def test_symmetric_aead(self):
             msg2 = monocypher.unlock(key, nonce, mac, c, associated_data=aead)
             self.assertEqual(msg, msg2)
 
+    def test_IncrementalAuthenticatedEncryption(self):
+        random = np.random.RandomState(seed=1)
+        for i in range(10):
+            message_length = random.randint(1, 4096)
+            aead_length = random.randint(1, 128)
+            key = bytes(random.randint(0, 256, 32, dtype=np.uint8))
+            nonce = bytes(random.randint(0, 256, 24, dtype=np.uint8))
+            aead = bytes(random.randint(0, 256, aead_length, dtype=np.uint8))
+            msg = bytes(random.randint(0, 256, message_length, dtype=np.uint8))
+
+            sender = monocypher.IncrementalAuthenticatedEncryption(key, nonce)
+            receiver = monocypher.IncrementalAuthenticatedEncryption(key, nonce)
+
+            mac, c = sender.lock(msg, associated_data=aead)
+            msg2 = receiver.unlock(mac, c, associated_data=aead)
+            self.assertEqual(msg, msg2)
+
+            # A second encryption works, but creates a different cipher text
+            mac2, c2 = sender.lock(msg, associated_data=aead)
+            self.assertNotEqual(c, c2)
+            self.assertNotEqual(mac, mac2)
+
+            msg2 = receiver.unlock(mac2, c2, associated_data=aead)
+            self.assertEqual(msg, msg2)
+
     def test_sign(self):
         random = np.random.RandomState(seed=1)
         for i in range(10):