From b4fd8c29ffcc7939ff6c3495bce5c94dd85d65d3 Mon Sep 17 00:00:00 2001
From: Leo Vasanko <git@vasanko.com>
Date: Sat, 28 Oct 2023 18:42:34 +0000
Subject: [PATCH] Python module with Numpy BitGenerator to integrate with
 np.random, and other Python use.

---
 randquik/cha.py    | 71 ++++++++++++++++++++++++++++++++++++++++++++++
 randquik/nprand.py | 31 ++++++++++++++++++++
 2 files changed, 102 insertions(+)
 create mode 100644 randquik/cha.py
 create mode 100644 randquik/nprand.py

diff --git a/randquik/cha.py b/randquik/cha.py
new file mode 100644
index 0000000..f2caad8
--- /dev/null
+++ b/randquik/cha.py
@@ -0,0 +1,71 @@
+from pathlib import Path
+from typing import Any
+
+import cffi
+
+src = Path(__file__).parent.parent / "src"
+
+if not src.is_dir():
+    raise RuntimeError("Unable to find RandQuik C sources in {src}")
+
+ffi = cffi.FFI()
+ffi.cdef(
+    """
+    typedef struct cha_ctx { uint32_t input[16]; } cha_ctx;
+
+    int cha_generate(uint8_t* out, uint64_t outlen, const uint8_t key[32], const uint8_t iv[16]);
+
+    void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv);
+    void cha_wipe(cha_ctx* ctx);
+    int cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen);
+    """
+)
+lib = ffi.dlopen("../build/librandquik-chacha20.so")
+
+
+def _processKeys(key, iv):
+    if len(key) != 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"
+        )
+    return ffi.from_buffer(key), ffi.from_buffer(iv)
+
+
+def _processBuffer(out):
+    if not out:
+        raise ValueError("Output buffer of non-zero size is required")
+    try:
+        outlen = out.nbytes
+    except AttributeError:
+        out = memoryview(out)
+        outlen = out.nbytes
+    if getattr(out, "readonly", None):
+        raise ValueError("The output buffer must be writable, not e.g. `bytes`")
+    return ffi.from_buffer(out), outlen
+
+
+class Cha:
+    def __init__(self, key: bytes | Any, iv: bytes | Any):
+        """Construct a generator that holds its internal state, moving forward on each call."""
+        key, iv = _processKeys(key, iv)
+        self.ctx = ffi.new("cha_ctx*")
+        lib.cha_init(self.ctx, key, iv)
+
+    def __del__(self):
+        lib.cha_wipe(self.ctx)
+
+    def __call__(self, out: bytearray | Any):
+        """Fill the parameter with random bytes"""
+        out, outlen = _processBuffer(out)
+        lib.cha_update(self.ctx, out, outlen)
+        return out
+
+
+def generate(out: bytearray | Any, key: bytes | Any, iv: bytes | Any):
+    """Setup a generator, fill the out buffer and dispose the generator"""
+    key, iv =_processKeys(key, iv)
+    out, outlen = _processBuffer(out)
+    lib.cha_generate(out, outlen, key, iv)
+    return out
diff --git a/randquik/nprand.py b/randquik/nprand.py
new file mode 100644
index 0000000..fd16ac2
--- /dev/null
+++ b/randquik/nprand.py
@@ -0,0 +1,31 @@
+import secrets
+
+import cha
+import numpy as np
+import sys
+
+class ChaRandom(np.random.BitGenerator):
+    def __init__(self, seed=None):
+        super().__init__(seed)
+        sys.stderr.write("Construct\n")
+        if seed is None:
+            key = secrets.token_bytes(32)
+        else:
+            key = (
+                np.random.SeedSequence(seed, pool_size=8)
+                .generate_state(4, dtype=np.uint64)
+                .tobytes()
+            )
+        self._generator = cha.Cha(key, bytes(8) + b"NumpyGen")
+
+    def random_raw(self, size=None):
+        sys.stderr.write(f"Random raw {size=}\n")
+        if size is None:
+            return int.from_bytes(self._generator(bytearray(8)), "little")
+        ret = np.empty(size, np.uint64)
+        self._generator(ret.data)
+        return ret
+
+    def spawn(self, n):
+        sys.stderr.write(f"Spawn {n=}\n")
+        raise NotImplementedError