implemented ChaCha-20

*implemented and added tests for ChaCha-20 Poly1305

Author: pvprajwal

pydatastructs/strings/algorithms.py

@@ -334,3 +334,101 @@ def sha256_encrypt(text) -> str:
             h = [(x + y) & 0xFFFFFFFF for x, y in zip(h, [a, b, c, d, e, f, g, h0])]
 
         return ''.join(f'{value:08x}' for value in h)
+
+    @staticmethod
+    def yield_chacha20_xor_stream(key, iv, position=0):
+        """Generate the xor stream with the ChaCha20 cipher."""
+        if not isinstance(position, int):
+            raise TypeError
+        if position & ~0xFFFFFFFF:
+            raise ValueError("Position is not uint32.")
+        if not isinstance(key, bytes):
+            raise TypeError
+        if not isinstance(iv, bytes):
+            raise TypeError
+        if len(key) != 32:
+            raise ValueError
+        if len(iv) != 8:
+            raise ValueError
+
+        def rotate(v, c):
+            return ((v << c) & 0xFFFFFFFF) | v >> (32 - c)
+
+        def quarter_round(x, a, b, c, d):
+            x[a] = (x[a] + x[b]) & 0xFFFFFFFF
+            x[d] = rotate(x[d] ^ x[a], 16)
+            x[c] = (x[c] + x[d]) & 0xFFFFFFFF
+            x[b] = rotate(x[b] ^ x[c], 12)
+            x[a] = (x[a] + x[b]) & 0xFFFFFFFF
+            x[d] = rotate(x[d] ^ x[a], 8)
+            x[c] = (x[c] + x[d]) & 0xFFFFFFFF
+            x[b] = rotate(x[b] ^ x[c], 7)
+
+        ctx = [0] * 16
+        ctx[:4] = (1634760805, 857760878, 2036477234, 1797285236)
+        ctx[4:12] = struct.unpack("<8L", key)
+        ctx[12] = ctx[13] = position
+        ctx[14:16] = struct.unpack("<LL", iv)
+        while 1:
+            x = list(ctx)
+            for i in range(10):
+                quarter_round(x, 0, 4, 8, 12)
+                quarter_round(x, 1, 5, 9, 13)
+                quarter_round(x, 2, 6, 10, 14)
+                quarter_round(x, 3, 7, 11, 15)
+                quarter_round(x, 0, 5, 10, 15)
+                quarter_round(x, 1, 6, 11, 12)
+                quarter_round(x, 2, 7, 8, 13)
+                quarter_round(x, 3, 4, 9, 14)
+            for c in struct.pack(
+                "<16L", *((x[i] + ctx[i]) & 0xFFFFFFFF for i in range(16))
+            ):
+                yield c
+            ctx[12] = (ctx[12] + 1) & 0xFFFFFFFF
+            if ctx[12] == 0:
+                ctx[13] = (ctx[13] + 1) & 0xFFFFFFFF
+
+    @staticmethod
+    def chacha20(data:bytes, key:bytes, nonce=None, position=0):
+        """
+        Encrypt (or decrypt) data using the ChaCha20 stream cipher.
+
+        Parameters
+        ==========
+        data : bytes
+            Input data to encrypt/decrypt
+        key : bytes
+            32-byte encryption key. Shorter keys will be padded by repetition.
+        nonce : bytes, optional
+            8-byte nonce
+
+        Returns
+        =======
+        bytes
+            Processed output
+
+        Notes
+        =====
+        Chacha20 is symmetric - same operation encrypts and decrypts.
+
+        References
+        ==========
+
+        .. [1] https://en.wikipedia.org/wiki/ChaCha20-Poly1305
+        .. [2] https://github.com/pts/chacha20/blob/master/chacha20_python3.py
+
+        """
+
+        if nonce is None:
+            nonce = b"\0" * 8
+
+        if not key:
+            raise ValueError("Key is empty.")
+        if len(key) < 32:
+            key = (key * (32 // len(key) + 1))[:32]
+        if len(key) > 32:
+            key = key[:32]
+
+        return bytes(
+            a ^ b for a, b in zip(data, Crypto.yield_chacha20_xor_stream(key, nonce, position))
+        )

pydatastructs/strings/tests/test_algorithms.py

@@ -97,3 +97,17 @@ def _test_sha256_encrypt():
 
     for text, expected in zip(test_cases, expected_hashes):
         assert Crypto.sha256_encrypt(text) == expected
+
+def _test_chacha20():
+    import binascii
+    cnvt = lambda x: binascii.unhexlify(bytes(x, 'ascii'))
+    test_cases = [
+        ('76b8e0ada0f13d90405d6ae55386bd28bdd219b8a08ded1aa836efcc8b770dc7da41597c5157488d7724e03fb8d84a376a43b8f41518a11cc387b669', '0000000000000000000000000000000000000000000000000000000000000000', '0000000000000000'),
+        ('4540f05a9f1fb296d7736e7b208e3c96eb4fe1834688d2604f450952ed432d41bbe2a0b6ea7566d2a5d1e7e20d42af2c53d792b1c43fea817e9ad275', '0000000000000000000000000000000000000000000000000000000000000001', '0000000000000000'),
+        ('de9cba7bf3d69ef5e786dc63973f653a0b49e015adbff7134fcb7df137821031e85a050278a7084527214f73efc7fa5b5277062eb7a0433e445f41e3', '0000000000000000000000000000000000000000000000000000000000000000', '0000000000000001'),
+        ('ef3fdfd6c61578fbf5cf35bd3dd33b8009631634d21e42ac33960bd138e50d32111e4caf237ee53ca8ad6426194a88545ddc497a0b466e7d6bbdb004', '0000000000000000000000000000000000000000000000000000000000000000', '0100000000000000'),
+        ('f798a189f195e66982105ffb640bb7757f579da31602fc93ec01ac56f85ac3c134a4547b733b46413042c9440049176905d3be59ea1c53f15916155c2be8241a38008b9a26bc35941e2444177c8ade6689de95264986d95889fb60e84629c9bd9a5acb1cc118be563eb9b3a4a472f82e09a7e778492b562ef7130e88dfe031c79db9d4f7c7a899151b9a475032b63fc385245fe054e3dd5a97a5f576fe064025d3ce042c566ab2c507b138db853e3d6959660996546cc9c4a6eafdc777c040d70eaf46f76dad3979e5c5360c3317166a1c894c94a371876a94df7628fe4eaaf2ccb27d5aaae0ad7ad0f9d4b6ad3b54098746d4524d38407a6deb', '000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f', '0001020304050607')
+    ]
+
+    for i, (ciphertext, key, iv) in enumerate(map(lambda t: tuple(map(cnvt, t)), test_cases)):
+        assert Crypto.chacha20_encrypt(b'\0' * len(ciphertext), key, iv) == ciphertext