Initial commit

README.md

+# LARKG from FrodoKEM and Kyber
+
+We implement LARKG using [Microsoft's FrodoKEM Python implementation](https://github.com/microsoft/PQCrypto-LWEKE/blob/master/python3/frodokem.py) and a [Kyber implmentation in pure Python](https://github.com/jack4818/kyber-py).

frodo_larkg.py

+from frodokem import FrodoKEM
+import numpy as np
+import struct, bitstring
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.hmac import HMAC
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+import timeit
+
+def hkdf(ikm, info, length=64):
+    hkdf = HKDF(
+        algorithm=hashes.SHA512(),
+        length=length,
+        salt=None,
+        info=info,
+        backend=default_backend(),
+    )
+
+    return hkdf.derive(ikm)
+
+def hmac(key, data):
+    hmac = HMAC(key, hashes.SHA512(), default_backend())
+    hmac.update(data)
+    return hmac.finalize()
+
+class ARKG_KEM(FrodoKEM):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.print_intermediate_values = False
+        self.seedA = b' '*self.len_seedA_bytes
+        self.A = self.gen(self.seedA)
+
+    def keygen(self):
+        """Generate a public key / secret key pair (FrodoKEM specification, 
+        Algorithm 12)"""
+        # 1. Choose uniformly random seeds s || seedSE || z
+        s_seedSE_z = self.randombytes(self.len_s_bytes + self.len_seedSE_bytes + self.len_z_bytes)
+        self.__print_intermediate_value("randomness", s_seedSE_z)
+        s = bytes(s_seedSE_z[0:self.len_s_bytes])
+        seedSE = bytes(s_seedSE_z[self.len_s_bytes : self.len_s_bytes + self.len_seedSE_bytes])
+        z = bytes(s_seedSE_z[self.len_s_bytes + self.len_seedSE_bytes : self.len_s_bytes + self.len_seedSE_bytes + self.len_z_bytes])
+        # 2. Generate pseudorandom seed seedA = SHAKE(z, len_seedA) (length in bits)
+        #seedA = self.shake(z, self.len_seedA_bytes)
+        #self.__print_intermediate_value("seedA", seedA)
+        # 3. A = Frodo.Gen(seedA)
+        #A = self.gen(seedA)
+        # self.__print_intermediate_value("A", A)
+        # 4. r = SHAKE(0x5F || seedSE, 2*n*nbar*len_chi) (length in bits), parsed as 2*n*nbar len_chi-bit integers in little-endian byte order
+        rbytes = self.shake(bytes(b'\x5f') + seedSE, 2 * self.n * self.nbar * self.len_chi_bytes)
+        r = [struct.unpack_from('<H', rbytes, 2*i)[0] for i in range(2 * self.n * self.nbar)]
+        self.__print_intermediate_value("r", r)
+        # 5. S^T = Frodo.SampleMatrix(r[0 .. n*nbar-1], nbar, n)
+        Stransposed = self.sample_matrix(r[0 : self.n * self.nbar], self.nbar, self.n)
+        self.__print_intermediate_value("S^T", Stransposed)
+        S = self.__matrix_transpose(Stransposed)
+        # 6. E = Frodo.SampleMatrix(r[n*nbar .. 2*n*nbar-1], n, nbar)
+        E = self.sample_matrix(r[self.n * self.nbar : 2 * self.n * self.nbar], self.n, self.nbar)
+        # self.__print_intermediate_value("E", E)
+        # 7. B = A S + E
+        #print'E',E)
+        #print'Froebenius Norm of E', np.linalg.norm(E), np.linalg.norm(E, np.inf))
+        B = self.__matrix_add(self.__matrix_mul(self.A,S), E)
+        self.__print_intermediate_value("B", B)
+        #print'B',B)
+        ##print'A',A)
+        # 8. b = Pack(B)
+        b = self.pack(B)
+        self.__print_intermediate_value("b", b)
+        # 9. pkh = SHAKE(seedA || b, len_pkh) (length in bits)
+        pkh = self.shake(self.seedA + b, self.len_pkh_bytes)
+        self.__print_intermediate_value("pkh", pkh)
+        # 10. pk = seedA || b, sk = (s || seedA || b, S^T, pkh)
+        pk = self.seedA + b
+        assert len(pk) == self.len_pk_bytes
+        sk = bitstring.BitArray()
+        sk.append(s + self.seedA + b)
+        for i in range(self.nbar):
+            for j in range(self.n):
+                sk.append(bitstring.BitArray(intle = Stransposed[i][j], length = 16))
+        sk.append(pkh)
+        sk = sk.bytes
+        assert len(sk) == self.len_sk_bytes
+        return (pk, sk, S, B)
+
+
+    def kem_keygen(self):
+        """Generate a public key / secret key pair (FrodoKEM specification, 
+        Algorithm 12)"""
+        # 1. Choose uniformly random seeds s || seedSE || z
+        s_seedSE_z = self.randombytes(self.len_s_bytes + self.len_seedSE_bytes + self.len_z_bytes)
+        self.__print_intermediate_value("randomness", s_seedSE_z)
+        s = bytes(s_seedSE_z[0:self.len_s_bytes])
+        seedSE = bytes(s_seedSE_z[self.len_s_bytes : self.len_s_bytes + self.len_seedSE_bytes])
+        #z = bytes(s_seedSE_z[self.len_s_bytes + self.len_seedSE_bytes : self.len_s_bytes + self.len_seedSE_bytes + self.len_z_bytes])
+        # 2. Generate pseudorandom seed seedA = SHAKE(z, len_seedA) (length in bits)
+        #seedA = self.shake(z, self.len_seedA_bytes)
+        #self.__print_intermediate_value("seedA", seedA)
+        # 3. A = Frodo.Gen(seedA)
+        A = self.A
+        # self.__print_intermediate_value("A", A)
+        # 4. r = SHAKE(0x5F || seedSE, 2*n*nbar*len_chi) (length in bits), parsed as 2*n*nbar len_chi-bit integers in little-endian byte order
+        rbytes = self.shake(bytes(b'\x5f') + seedSE, 2 * self.n * self.nbar * self.len_chi_bytes)
+        r = [struct.unpack_from('<H', rbytes, 2*i)[0] for i in range(2 * self.n * self.nbar)]
+        self.__print_intermediate_value("r", r)
+        # 5. S^T = Frodo.SampleMatrix(r[0 .. n*nbar-1], nbar, n)
+        Stransposed = self.sample_matrix(r[0 : self.n * self.nbar], self.nbar, self.n)
+        self.__print_intermediate_value("S^T", Stransposed)
+        S = self.__matrix_transpose(Stransposed)
+        # 6. E = Frodo.SampleMatrix(r[n*nbar .. 2*n*nbar-1], n, nbar)
+        E = self.sample_matrix(r[self.n * self.nbar : 2 * self.n * self.nbar], self.n, self.nbar)
+        # self.__print_intermediate_value("E", E)
+        # 7. B = A S + E
+        #print'E',E)
+        #print'Froebenius Norm of E', np.linalg.norm(E), np.linalg.norm(E, np.inf))
+        B = self.__matrix_add(self.__matrix_mul(A,S), E)
+        self.__print_intermediate_value("B", B)
+        #print'B',B)
+        #print'A',A)
+        # 8. b = Pack(B)
+        b = self.pack(B)
+        self.__print_intermediate_value("b", b)
+        # 9. pkh = SHAKE(seedA || b, len_pkh) (length in bits)
+        pkh = self.shake(self.seedA + b, self.len_pkh_bytes)
+        self.__print_intermediate_value("pkh", pkh)
+        # 10. pk = seedA || b, sk = (s || seedA || b, S^T, pkh)
+        pk = self.seedA + b
+        assert len(pk) == self.len_pk_bytes
+        sk = bitstring.BitArray()
+        sk.append(s + self.seedA + b)
+        for i in range(self.nbar):
+            for j in range(self.n):
+                sk.append(bitstring.BitArray(intle = Stransposed[i][j], length = 16))
+        sk.append(pkh)
+        sk = sk.bytes
+        assert len(sk) == self.len_sk_bytes
+        return (pk, sk, S, B)
+        
+    def kem_encaps(self, pk):
+        """Encapsulate against a public key to create a ciphertext and shared secret 
+        (FrodoKEM specification, Algorithm 13)"""
+        # Parse pk = seedA || b
+        assert len(pk) == self.len_seedA_bytes + self.D * self.n * self.nbar / 8, "Incorrect public key length"
+        seedA = self.seedA #pk[0 : self.len_seedA_bytes]
+        b = pk[self.len_seedA_bytes:]
+        # 1. Choose a uniformly random key mu in {0,1}^len_mu (length in bits)
+        mu = self.randombytes(self.len_mu_bytes)
+        self.__print_intermediate_value("mu", mu)
+        # 2. pkh = SHAKE(pk, len_pkh)
+        pkh = self.shake(pk, self.len_pkh_bytes)
+        self.__print_intermediate_value("pkh", pkh)
+        # 3. seedSE || k = SHAKE(pkh || mu, len_seedSE + len_k) (length in bits)
+        seedSE_k = self.shake(pkh + mu, self.len_seedSE_bytes + self.len_k_bytes)
+        seedSE = seedSE_k[0:self.len_seedSE_bytes]
+        self.__print_intermediate_value("seedSE", seedSE)
+        k = seedSE_k[self.len_seedSE_bytes:self.len_seedSE_bytes + self.len_k_bytes]
+        self.__print_intermediate_value("k", k)
+        # 4. r = SHAKE(0x96 || seedSE, 2*mbar*n + mbar*nbar*len_chi) (length in bits)
+        rbytes = self.shake(bytes(b'\x96') + seedSE, (2 * self.mbar * self.n + self.mbar * self.mbar) * self.len_chi_bytes)
+        r = [struct.unpack_from('<H', rbytes, 2*i)[0] for i in range(2 * self.mbar * self.n + self.mbar * self.nbar)]
+        self.__print_intermediate_value("r", r)
+        # 5. S' = Frodo.SampleMatrix(r[0 .. mbar*n-1], mbar, n)
+        Sprime = self.sample_matrix(r[0 : self.mbar * self.n], self.mbar, self.n)
+        self.__print_intermediate_value("S'", Sprime)
+        # 6. E' = Frodo.SampleMatrix(r[mbar*n .. 2*mbar*n-1], mbar, n)
+        Eprime = self.sample_matrix(r[self.mbar * self.n : 2 * self.mbar * self.n], self.mbar, self.n)
+        self.__print_intermediate_value("E'", Eprime)
+        # 7. A = Frodo.Gen(seedA)
+        A = self.A #self.gen(seedA)
+        # 8. B' = S' A + E'
+        Bprime = self.__matrix_add(self.__matrix_mul(Sprime, A), Eprime)
+        self.__print_intermediate_value("B'", Bprime)
+        # 9. c1 = Frodo.Pack(B')
+        c1 = self.pack(Bprime)
+        self.__print_intermediate_value("c1", c1)
+        # 10. E'' = Frodo.SampleMatrix(r[2*mbar*n .. 2*mbar*n + mbar*nbar-1], mbar, n)
+        Eprimeprime = self.sample_matrix(r[2 * self.mbar * self.n : 2 * self.mbar * self.n + self.mbar * self.nbar], self.mbar, self.nbar)
+        self.__print_intermediate_value("E''", Eprimeprime)
+        # 11. B = Frodo.Unpack(b, n, nbar)
+        B = self.unpack(b, self.n, self.nbar)
+        self.__print_intermediate_value("B", B)
+        # 12. V = S' B + E''
+        V = self.__matrix_add(self.__matrix_mul(Sprime, B), Eprimeprime)
+        self.__print_intermediate_value("V", V)
+        # 13. C = V + Frodo.Encode(mu)
+        self.__print_intermediate_value("mu_encoded", self.encode(mu))
+        C = self.__matrix_add(V, self.encode(mu))
+        self.__print_intermediate_value("C", C)
+        # 14. c2 = Frodo.Pack(C)
+        c2 = self.pack(C)
+        self.__print_intermediate_value("c2", c2)
+        # 15. ss = SHAKE(c1 || c2 || k, len_ss)
+        ss = self.shake(c1 + c2 + k, self.len_ss_bytes)
+        ct = c1 + c2
+        assert len(ct) == self.len_ct_bytes
+        assert len(ss) == self.len_ss_bytes
+        return (ct, k)
+    
+    
+    def kem_decaps(self, sk, ct):
+        """Decapsulate a ciphertext using a secret key to obtain a shared secret 
+        (FrodoKEM specification, Algorithm 14)"""
+        # Parse ct = c1 || c2
+        assert len(ct) == self.len_ct_bytes, "Incorrect ciphertext length"
+        offset = 0; length = int(self.mbar * self.n * self.D / 8)
+        c1 = ct[offset:offset+length]
+        self.__print_intermediate_value("c1", c1)
+        offset += length; length = int(self.mbar * self.nbar * self.D / 8)
+        c2 = ct[offset:offset+length]
+        self.__print_intermediate_value("c2", c2)
+        # Parse sk = (s || seedA || b, S^T, pkh)
+        assert len(sk) == self.len_sk_bytes
+        offset = 0; length = self.len_s_bytes
+        s = sk[offset:offset+length]
+        self.__print_intermediate_value("s", s)
+        offset += length; length = self.len_seedA_bytes
+        seedA = self.seedA #sk[offset:offset+length]
+        self.__print_intermediate_value("seedA", seedA)
+        offset += length; length = int(self.D * self.n * self.nbar / 8)
+        b = sk[offset:offset+length]
+        self.__print_intermediate_value("b", b)
+        offset += length; length = int(self.n * self.nbar * 16 / 8)
+        Sbytes = bitstring.ConstBitStream(sk[offset:offset+length])
+        Stransposed = [[0 for j in range(self.n)] for i in range(self.nbar)]
+        for i in range(self.nbar):
+            for j in range(self.n):
+                Stransposed[i][j] = Sbytes.read('intle:16')
+        self.__print_intermediate_value("S^T", Stransposed)
+        S = self.__matrix_transpose(Stransposed)
+        offset += length; length = self.len_pkh_bytes
+        pkh = sk[offset:offset+length]
+        self.__print_intermediate_value("pkh", pkh)
+        # 1. B' = Frodo.Unpack(c1, mbar, n)
+        Bprime = self.unpack(c1, self.mbar, self.n)
+        self.__print_intermediate_value("B'", Bprime)
+        # 2. C = Frodo.Unpack(c2, mbar, nbar)
+        C = self.unpack(c2, self.mbar, self.nbar)
+        self.__print_intermediate_value("C", C)
+        # 3. M = C - B' S
+        BprimeS = self.__matrix_mul(Bprime, S)
+        self.__print_intermediate_value("B'S", BprimeS)
+        M = self.__matrix_sub(C, BprimeS)
+        self.__print_intermediate_value("M", M)
+        # 4. mu' = Frodo.Decode(M)
+        muprime = self.decode(M)
+        self.__print_intermediate_value("mu'", muprime)
+        # 5. Parse pk = seedA || b
+        # (done above)
+        # 6. seedSE' || k' = SHAKE(pkh || mu', len_seedSE + len_k) (length in bits)
+        seedSEprime_kprime = self.shake(pkh + muprime, self.len_seedSE_bytes + self.len_k_bytes)
+        seedSEprime = seedSEprime_kprime[0:self.len_seedSE_bytes]
+        self.__print_intermediate_value("seedSE'", seedSEprime)
+        kprime = seedSEprime_kprime[self.len_seedSE_bytes:self.len_seedSE_bytes + self.len_k_bytes]
+        self.__print_intermediate_value("k'", kprime)
+        # 7. r = SHAKE(0x96 || seedSE', 2*mbar*n + mbar*nbar*len_chi) (length in bits)
+        rbytes = self.shake(bytes(b'\x96') + seedSEprime, (2 * self.mbar * self.n + self.mbar * self.mbar) * self.len_chi_bytes)
+        r = [struct.unpack_from('<H', rbytes, 2*i)[0] for i in range(2 * self.mbar * self.n + self.mbar * self.nbar)]
+        self.__print_intermediate_value("r", r)
+        # 8. S' = Frodo.SampleMatrix(r[0 .. mbar*n-1], mbar, n)
+        Sprime = self.sample_matrix(r[0 : self.mbar * self.n], self.mbar, self.n)
+        self.__print_intermediate_value("S'", Sprime)
+        # 9. E' = Frodo.SampleMatrix(r[mbar*n .. 2*mbar*n-1], mbar, n)
+        Eprime = self.sample_matrix(r[self.mbar * self.n : 2 * self.mbar * self.n], self.mbar, self.n)
+        self.__print_intermediate_value("E'", Eprime)
+        # 10. A = Frodo.Gen(seedA)
+        A = self.A #self.gen(seedA)
+        # 11. B'' = S' A + E'
+        Bprimeprime = self.__matrix_add(self.__matrix_mul(Sprime, A), Eprime)
+        self.__print_intermediate_value("B''", Bprimeprime)
+        # 12. E'' = Frodo.SampleMatrix(r[2*mbar*n .. 2*mbar*n + mbar*nbar-1], mbar, n)
+        Eprimeprime = self.sample_matrix(r[2 * self.mbar * self.n : 2 * self.mbar * self.n + self.mbar * self.nbar], self.mbar, self.nbar)
+        self.__print_intermediate_value("E''", Eprimeprime)
+        # 13. B = Frodo.Unpack(b, n, nbar)
+        B = self.unpack(b, self.n, self.nbar)
+        self.__print_intermediate_value("B", B)
+        # 14. V = S' B + E''
+        V = self.__matrix_add(self.__matrix_mul(Sprime, B), Eprimeprime)
+        self.__print_intermediate_value("V", V)
+        # 15. C' = V + Frodo.Encode(muprime)
+        Cprime = self.__matrix_add(V, self.encode(muprime))
+        self.__print_intermediate_value("C'", Cprime)
+        # 16. (in constant time) kbar = kprime if (B' || C == B'' || C') else kbar = s
+        # Needs to avoid branching on secret data as per:
+        #     Qian Guo, Thomas Johansson, Alexander Nilsson. A key-recovery timing attack on post-quantum 
+        #     primitives using the Fujisaki-Okamoto transformation and its application on FrodoKEM. In CRYPTO 2020.
+        #use_kprime = self.__ctverify(Bprime + C, Bprimeprime + Cprime)
+        #kbar = self.__ctselect(kprime, s, use_kprime)
+        # 17. ss = SHAKE(c1 || c2 || kbar, len_ss) (length in bits)
+        ss = self.shake(c1 + c2 + kprime, self.len_ss_bytes)
+        assert len(ss) == self.len_ss_bytes
+        return kprime
+
+
+    def mat_add(self, A, B):
+        return self._FrodoKEM__matrix_add(A, B)
+    
+    def __matrix_add(self, A, B):
+        return self._FrodoKEM__matrix_add(A, B)
+    
+    def mat_mul(self, A, B):
+        return self._FrodoKEM__matrix_mul(A, B)
+    
+    def __matrix_mul(self, A, B):
+        return self._FrodoKEM__matrix_mul(A, B)
+    
+    def mat_mul(self, A, B):
+        return self._FrodoKEM__matrix_mul(A, B)
+        
+    def __print_intermediate_value(self, name, value):
+        return self._FrodoKEM__print_intermediate_value(name, value)
+    
+    def matrix_sub(self, a, b):
+        return self._FrodoKEM__matrix_sub(a,b)
+    
+    def __matrix_sub(self, a, b):
+        return self._FrodoKEM__matrix_sub(a,b)
+    
+    def __ctverify(self, a, b):
+        return self._FrodoKEM__ctverify(a,b)
+    
+    def __ctselect(self, a, b, c):
+        return self._FrodoKEM__ctselect(a,b,c)
+    
+    def __matrix_transpose(self, e):
+        return self._FrodoKEM__matrix_transpose(e)
+    
+    def sample(self, r):
+        """Sample from the error distribution using noise r (a two-byte array 
+        encoding a 16-bit integer in little-endian byte order) (FrodoKEM 
+        specification, Algorithm 5)"""
+        # 1. t = sum_{i=1}^{len_x - 1} r_i * 2^{i-1}
+        t = r >> 1
+        # 2. e = 0
+        e = 0
+        # 3. for z = 0; z < s; z += 1
+        for z in range(len(self.T_chi) - 1):
+            # 4. if t > T_chi(z)
+            if t > self.T_chi[z]:
+                # 5. e = e + 1
+                e += 1
+        # 6. e = (-1)^{r_0} * e
+        r0 = r % 2
+        e = ((-1) ** r0) * e
+        return e
+
+    def sample_matrix(self, r, n1, n2):
+        """Sample an n1 x n2 matrix from the error distribution using noise r 
+        (FrodoKEM specification, Algorithm 6)"""
+        E = [[None for j in range(n2)] for i in range(n1)]
+        # 1. for i = 0; i < n1; i += 1
+        for i in range(n1):
+            # 2. for j = 0; j < n2; j += 1
+            for j in range(n2):
+                # 3. E[i][j] = Frodo.Sample(r^{i*n2+j}, T_chi)
+                E[i][j] = self.sample(r[i * n2 + j])
+        return E
+
+
+def concat(in_array):
+    out = []
+    for i in range(0, len(in_array), 2):
+        out.append(int.from_bytes(bytearray([in_array[i], in_array[i+1]]), 'little'))
+    return out
+
+
+def to_signed(input):
+    for i in range(arkg_kem.n):
+        for j in range(arkg_kem.nbar):
+            if input[i][j] > np.abs(input[i][j] - 65536):
+                input[i][j] = input[i][j] - 65536
+    return input
+
+
+def derive_pk(S):
+    c, K = arkg_kem.kem_encaps(S)
+    kmac = hkdf(K, b'1')
+    kcred = concat([b for b in hkdf(K, b'2', length=arkg_kem.n*arkg_kem.nbar*2)])
+    kcred = arkg_kem.sample_matrix(kcred,arkg_kem.n, arkg_kem.nbar)
+
+    mu = hmac(kmac, c)
+    P = arkg_kem.mat_add(arkg_kem.unpack(S, arkg_kem.n, arkg_kem.nbar), arkg_kem.mat_mul(arkg_kem.A, kcred))
+    
+    p = arkg_kem.seedA + arkg_kem.pack(P)
+    cred = c, mu, p
+    return p, cred
+
+
+def derive_sk(sk, sm, cred):
+    c, mu, pkp = cred
+    K = arkg_kem.kem_decaps(sk, c)
+
+    kmac = hkdf(K, b'1')
+    kcred = concat([b for b in hkdf(K, b'2', length=arkg_kem.n*arkg_kem.nbar*2)])   
+    kcred = arkg_kem.sample_matrix(kcred,arkg_kem.n, arkg_kem.nbar)
+
+    sk = arkg_kem.unpack(sk, arkg_kem.nbar, arkg_kem.n)
+
+    p = to_signed(arkg_kem.mat_add(sm, kcred))
+    s = arkg_kem.pack(p)
+
+    Stransposed = arkg_kem._ARKG_KEM__matrix_transpose(p)
+    pkh = arkg_kem.shake(b'blah', arkg_kem.len_pkh_bytes) # match hash of pkp
+    skn = bitstring.BitArray()
+    skn.append(b' '* 24 + pkp)
+    for i in range(arkg_kem.nbar):
+        for j in range(arkg_kem.n):
+            skn.append(bitstring.BitArray(intle = Stransposed[i][j], length = 16))
+    skn.append(pkh)
+    skn = skn.bytes
+
+    return  skn
+
+
+def check(pkp, skp, P, B, kcred):
+    return 'NI'
+
+

kyber_larkg.py

+from kyber_skem import KyberSKEM, Kyber, DEFAULT_PARAMETERS
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import hashes
+from collections import namedtuple
+from random import uniform
+from operator import mul
+from functools import reduce
+import time
+
+
+PublicParams = namedtuple('PublicParams', 'skem kdf_1 kdf_2 rsp rsp_dom M')
+Cred = namedtuple('Cred', 'B_prime c mu')
+
+
+def _hkdf(msg):
+    kdf = HKDF(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=None,
+        info=None,
+        backend=default_backend())
+    return kdf.derive(msg)
+
+
+''' Invert Montgomery and flatten. '''
+def _from_montgomery_to_flat(pp, poly):
+    def inverse(elem):
+        return elem * pp.skem.R.ntt_helper.mont_r_inv % pp.skem.q
+
+    return [v - pp.skem.q if (v := inverse(e)) >= pp.skem.q - 2*pp.skem.eta_1
+            else v for row in poly for col in row for e in col]
+
+
+def setup(skem, rsp, rsp_dom, repetition_rate):
+    kdf_1 = lambda K: skem._generate_error_vector(K, skem.eta_1, 0)[0]
+    return PublicParams(skem, kdf_1, _hkdf, rsp, rsp_dom, repetition_rate)
+
+
+def keygen(pp):
+    return pp.skem.keygen_dec()
+
+
+def derive_pk(pp, pk):
+    B = pk
+    B_poly = pp.skem.M.decode(B, pp.skem.k, 1, l=12, is_ntt=True)
+
+    B_prime, S_prime = pp.skem.keygen_enc()                              # Ln 1
+    c, k_seed = pp.skem.encaps(S_prime, B)                               # Ln 2
+
+    K = pp.kdf_1(k_seed)                                                 # Ln 3
+    K.to_ntt()
+
+    mu = pp.kdf_2(k_seed)                                                # Ln 4
+
+    E_prime = pp.skem._generate_error_vector(pp.skem.random_bytes(32),   # Ln 5
+                                             pp.skem.eta_1, 0)[0]
+    E_prime.to_ntt()
+
+
+    P = (pp.skem.A @ K).to_montgomery() + E_prime + B_poly               # Ln 6
+    P.reduce_coefficents()
+    return P.encode(l=12), Cred(B_prime, c, mu)                          # Ln 7
+
+
+def derive_sk(pp, sk, cred):
+    B_prime, c, mu = cred
+    S = sk
+    S_poly = pp.skem.M.decode(S, pp.skem.k, 1, l=12, is_ntt=True)
+
+    S_prime_prime = False                                                # Ln 1
+    k_seed = pp.skem.decaps(S, c, B_prime)                               # Ln 2
+
+    K = pp.kdf_1(k_seed)                                                 # Ln 3
+    K.to_ntt()
+
+    mu_star = pp.kdf_2(k_seed)                                           # Ln 4
+
+    if mu_star == mu:                                                    # Ln 5
+        S_prime_prime = S_poly + K                                       # Ln 6
+        S_prime_prime.reduce_coefficents()
+        S_prime_prime_encoded = S_prime_prime.encode(l=12)
+    else:
+        return False  # Abort early as chi_b_S_prime_prime otherwise undefined
+
+    u = uniform(0, 1)  # NOT CRYPTOGRAPHICALLY SECURE                    # Ln 7
+
+    S_poly.from_ntt()
+    S_flat = _from_montgomery_to_flat(pp, S_poly)
+
+    S_prime_prime.from_ntt()
+    S_prime_prime_flat = _from_montgomery_to_flat(pp, S_prime_prime)
+
+    chi_a_S = [pp.rsp[abs(elem)]/pp.rsp_dom for elem in S_flat]
+    chi_b_S_prime_prime = [pp.rsp[abs(elem - S_flat[i])]/pp.rsp_dom for i, elem
+                           in enumerate(S_prime_prime_flat)]
+
+    reject_value = reduce(mul, map(lambda pair: pair[1] / pair[0],
+                                   zip(chi_a_S, chi_b_S_prime_prime)), 1)
+
+    if u < reject_value * pp.M:
+        return False                                                     # Ln 8
+
+    return S_prime_prime_encoded                                         # Ln 9
+
+
+def test(pp, skp, pkp):
+    if not skp:
+        return False
+
+    pkp = pkp + pp.skem.rho
+    m = b' '*32
+    c = pp.skem._cpapke_enc(pkp, m, pp.skem.random_bytes(32))
+    mm = pp.skem._cpapke_dec(skp, c)
+    return m == mm
+
+
+if __name__ == '__main__':
+    pp = setup(KyberSKEM(DEFAULT_PARAMETERS['kyber_1024']),
+               rsp=[6, 4, 1, 0, 0], rsp_dom=16,
+               repetition_rate=162_755)
+    computed_in = []
+    for i in range(1000):
+        count = 0
+        skp = False
+        while not skp:
+            pk, sk = keygen(pp)
+            pkp, cred = derive_pk(pp, pk)
+            skp = derive_sk(pp, sk, cred)
+            count += 1
+        computed_in.append(count)
+        print(i)
+    print(sum(computed_in))
+        #assert test(pp, skp, pkp)
+    print(rejects)

kyber_skem.py

+from kyber import Kyber, DEFAULT_PARAMETERS
+
+class KyberSKEM(Kyber):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.keygen_dec = self.keygen
+
+        # Fix public parameters rho and A.
+        self.rho, _ = self._g(self.random_bytes(32))
+        self.A = self._generate_matrix_from_seed(self.rho, is_ntt=True)
+        self.At = self._generate_matrix_from_seed(self.rho, transpose=True,
+                                                  is_ntt=True)
+
+
+    ''' Kyber CPA PKE KeyGen with fixed A. '''
+    def _cpapke_keygen(self):
+        N, _, sigma = 0, *self._g(self.random_bytes(32))
+
+        # Generate the error vector s ∈ R^k
+        s, N = self._generate_error_vector(sigma, self.eta_1, N)
+        s.to_ntt()
+
+        # Generate the error vector e ∈ R^k
+        e, N = self._generate_error_vector(sigma, self.eta_1, N)
+        e.to_ntt()
+
+        # Construct the public key
+        t = (self.A @ s).to_montgomery() + e
+
+        # Reduce vectors mod^+ q
+        t.reduce_coefficents()
+        s.reduce_coefficents()
+
+        # Encode elements to bytes and return
+        pk = t.encode(l=12) + self.rho
+        sk = s.encode(l=12)
+        return pk, sk
+
+
+    ''' Returns encapsulation keypair. '''
+    def keygen_enc(self):
+        N, coins = 0, self.random_bytes(32)
+
+        # Generate the error vector r ∈ R^k
+        skp, N = self._generate_error_vector(coins, self.eta_1, N)
+        skp.to_ntt()
+
+        # Generate the error vector e1 ∈ R^k
+        e1, N = self._generate_error_vector(coins, self.eta_2, N)
+
+        # Module/Polynomial arithmatic
+        u = (self.At @ skp).from_ntt() + e1
+
+        u.reduce_coefficents()
+        skp.reduce_coefficents()
+
+        # Ciphertext to bytes
+        pkp = u.encode(l=12)
+
+        return pkp, skp
+
+
+    ''' Encapsulate for pk using skp. '''
+    def encaps(self, skp, pk):
+        N, K = 0, self._h(self.random_bytes(32))
+        tt = self.M.decode(pk, 1, self.k, l=12, is_ntt=True)
+
+        # Encode message as polynomial
+        m_poly = self.R.decode(K, l=1).decompress(1)
+
+        coins = self.random_bytes(32)
+
+        # Generate the error polynomial e2 ∈ R
+        input_bytes = self._prf(coins,  bytes([N]), 64*self.eta_2)
+        e2 = self.R.cbd(input_bytes, self.eta_2)
+
+        # Module/Polynomial arithmatic 
+        v = (tt @ skp)[0][0].from_ntt()
+        v = v + e2 + m_poly
+
+        # Ciphertext to bytes
+        c = v.compress(self.dv).encode(l=self.dv)
+
+        return c, K
+
+
+    ''' Decapsulate ciphertext C with sk and pkp. '''
+    def decaps(self, sk, c, pkp):
+        u = self.M.decode(pkp, self.k, 1, l=12)
+        u.to_ntt()
+
+        v = self.R.decode(c, l=self.dv).decompress(self.dv)
+
+        # s_transpose (already in NTT form)
+        st = self.M.decode(sk, 1, self.k, l=12, is_ntt=True)
+
+        # Recover message as polynomial
+        m = (st @ u)[0][0].from_ntt()
+        m = v - m
+
+        # Return message as bytes
+        m = m.compress(1).encode(l=1)
+        return m
+
+
+if __name__ == '__main__':
+    skem = KyberSKEM(DEFAULT_PARAMETERS['kyber_1024'])
+    pk, sk = skem.keygen_dec()
+    pkp, skp = skem.keygen_enc()
+    c, K = skem.encaps(skp, pk)
+    kk = skem.decaps(sk, c, pkp)
+    assert K == kk