From 4f42fb344778c1c7232018429383956a819d5283 Mon Sep 17 00:00:00 2001
From: Leo Vasanko <git@vasanko.com>
Date: Sun, 29 Oct 2023 12:17:05 +0000
Subject: [PATCH] Added Python convenience for different nonce sizes.
 Statistical tests of output randomness.

---
 pyproject.toml       |  6 +++---
 randquik/cha.py      | 18 +++++++++++-----
 tests/test_chacha.py | 50 +++++++++++++++++++++++++++++++++++++-------
 3 files changed, 58 insertions(+), 16 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 8c38894..06dfa1f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,7 +3,7 @@ requires = ["hatchling", "hatch-vcs", "wheel", "cffi"]
 build-backend = "hatchling.build"
 
 [project]
-name = "RandQuik"
+name = "randquik"
 version = "0.1.0"
 description = "Extremely fast and cryptographically secure random number generator."
 readme = "README.md"
@@ -16,7 +16,7 @@ classifiers = [
     "Topic :: Security",
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
-dependencies = ["cffi>=1.0.1"]
+dependencies = ["cffi>=1.0.1", "numpy"]
 requires-python = ">=3.10"
 keywords = [
     "random",
@@ -30,7 +30,7 @@ keywords = [
 [project.urls]
 
 [project.optional-dependencies]
-dev = ["pytest", "ruff", "cryptography"]
+dev = ["pytest", "ruff", "cryptography", "scipy"]
 
 [tool.hatchling]
 
diff --git a/randquik/cha.py b/randquik/cha.py
index c368218..430682e 100644
--- a/randquik/cha.py
+++ b/randquik/cha.py
@@ -30,12 +30,20 @@ lib = ffi.dlopen(
 
 
 def _processKeys(key, iv):
-    if len(key) != 32:
+    key = memoryview(key)
+    iv = memoryview(iv)
+    if key.nbytes != 32:
         raise ValueError("key must be 32 bytes")
-    if len(iv) != 16:
-        raise ValueError(
-            "iv must be full 16 bytes, starting with the counter - usually zeroes - followed by nonce"
-        )
+    if iv.nbytes != 16:
+        # Allow original and IETF nonces with zero counter
+        if iv.nbytes == 8:
+            iv = bytes(8) + iv
+        elif iv.nbytes == 12:
+            iv = bytes(4) + iv
+        else:
+            raise ValueError(
+                "iv lenth must be 8 (original nonce), 12 (IETF) or 16 (counter in initial 8 bytes)"
+            )
     return ffi.from_buffer(key), ffi.from_buffer(iv)
 
 
diff --git a/tests/test_chacha.py b/tests/test_chacha.py
index c2903d4..0af824f 100644
--- a/tests/test_chacha.py
+++ b/tests/test_chacha.py
@@ -1,17 +1,50 @@
 from secrets import randbelow, token_bytes
 
+import numpy as np
 from cryptography.hazmat.primitives.ciphers import Cipher
 from cryptography.hazmat.primitives.ciphers.algorithms import ChaCha20
+from scipy.stats import chisquare
 
 from randquik import cha
 
 
+def test_cha_generate_statistical():
+    """Requests in multiple of ChaCha20 block size 64 bytes"""
+    ROUNDS = 100
+    min_p = 1 - 0.99 ** (1 / ROUNDS)  # 99 % confidence over the entire test
+    print(min_p)
+
+    for round in range(ROUNDS):
+        # First round both zero, use 8-byte nonce as round counter
+        key = bytes(32)
+        nonce = round.to_bytes(8, "little")
+        iv = bytes(8) + nonce  # Cryptography module requires IV (counter, nonce)
+        # Varying sizes to test internal processing that occurs in 64 bit blocks
+        # and with SIMD implementations also 256 or 512 bytes at a time.
+        N = 10000 + randbelow(10000)
+        ct0, ct1 = bytearray(N), bytearray(N)
+        Cipher(ChaCha20(key, iv), None, None).encryptor().update_into(bytes(N), ct0)
+        cha.generate(ct1, key, nonce)
+        assert len(ct0) == len(ct1)
+        assert ct0.hex() == ct1.hex()
+        assert ct1.count(0)
+
+        # Zeroing out 50 bytes fail the test in ~30 rounds
+        # ct1[:50] = bytes(50)
+
+        # Test that all byte values are equivalently common (despite zero inputs)
+        observed = np.bincount(ct1, minlength=256)
+        expected = np.full(256, N / 256)
+        chi2, p_value = chisquare(observed, expected)
+        assert p_value >= min_p, f"{round=} {N=}"
+
+
 def test_cipherstreams_fullblocks():
     """Requests in multiple of ChaCha20 block size 64 bytes"""
     key = token_bytes(32)
-    nonce = token_bytes(16)
-    c0 = Cipher(ChaCha20(key, nonce), None, None).encryptor()
-    c1 = cha.Cha(key, nonce)
+    iv = token_bytes(16)
+    c0 = Cipher(ChaCha20(key, iv), None, None).encryptor()
+    c1 = cha.Cha(key, iv)
 
     for i in range(2048):
         N = 64 * (1 + randbelow(2048))
@@ -24,9 +57,9 @@ def test_cipherstreams_fullblocks():
 def test_cipherstreams_partial_updates():
     """Odd-sized requests that retain leftover buffers"""
     key = token_bytes(32)
-    nonce = token_bytes(16)
-    c0 = Cipher(ChaCha20(key, nonce), None, None).encryptor()
-    c1 = cha.Cha(key, nonce)
+    iv = token_bytes(16)
+    c0 = Cipher(ChaCha20(key, iv), None, None).encryptor()
+    c1 = cha.Cha(key, iv)
 
     for i in range(2048):
         N = 1 + randbelow(2048)
@@ -39,8 +72,9 @@ def test_cipherstreams_partial_updates():
 def test_cipherstreams_32leftover():
     """Test the special case where the second response comes entirely from the leftover buffer"""
     key = token_bytes(32)
-    nonce = token_bytes(16)
-    c0 = Cipher(ChaCha20(key, nonce), None, None).encryptor()
+    nonce = token_bytes(12)  # IETF nonce size
+    iv = bytes(4) + nonce
+    c0 = Cipher(ChaCha20(key, iv), None, None).encryptor()
     c1 = cha.Cha(key, nonce)
 
     for N in [512 - 32, 32]: