CRYSTALS-Kyber Python Implementation

This repository contains a pure python implementation of CRYSTALS-Kyber following (at the time of writing) the most recent specification (v3.02)

Disclaimer

⚠️ Under no circumstances should this be used for a cryptographic application. ⚠️

I have written kyber-py as a way to learn about the way Kyber works, and to try and create a clean, well commented implementation which people can learn from.

This code is not constant time, or written to be performant. Rather, it was written so that reading though Algorithms 1-9 in the specification closely matches the code which is seen in kyber.py.

KATs

This implementation currently passes all KAT tests from the reference implementation. For more information, see the unit tests in test_kyber.py.

Note: there is a discrepancy between the specification and reference implementation. To ensure all KATs pass, I have to generate the public key before the random bytes $z = \mathcal{B}^{32}$ in algorithm 7 of the specification (v3.02).

Dependencies

Originally this was planned to have zero dependencies, however to make this work pass the KATs, I needed a deterministic CSRNG. The reference implementation uses AES256 CTR DRBG. I have implemented this in aes256_ctr_drbg.py. However, I have not implemented AES itself, instead I import this from pycryptodome.

To install dependencies, run pip -r install requirements.

If you're happy to use system randomness (os.urandom) then you don't need this dependency.

Using kyber-py

There are three functions exposed on the Kyber class which are intended for use:

• Kyber.keygen(): generate a keypair (pk, sk)
• Kyber.enc(pk): generate a challenge and a shared key (c, K)
• Kyber.dec(sk, c): generate the shared key K

To use Kyber() it must be initialised with a dictionary of the protocol parameters. An example can be seen in DEFAULT_PARAMETERS.

Additionally, the class has been initialised with these default parameters, so you can simply import the NIST level you want to play with:

Example

>>> from kyber import Kyber512
>>> pk, sk = Kyber512.keygen()
>>> c, key = Kyber512.enc(pk)
>>> _key = Kyber512.dec(c, sk)
>>> assert key == _key

The above example would also work with Kyber768 and Kyber1024.

Benchmarks

TODO: Better benchmarks? Although this was never about speed haha.

For now, here are some approximate benchmarks:

1000 Iterations	Kyber512	Kyber768	Kyber1024
KeyGen()	6.868s	10.820s	16.172s
Enc()	10.677s	16.094s	22.341s
Dec()	16.822s	25.979s	33.524s

All times recorded using a Intel Core i7-9750H CPU.

Future Plans

• Add documentation on NTT transform for polynomials
• Add documentation for working with DRBG and setting the seed

Include Dilithium

Using polynomials.py and modules.py this work could be extended to have a pure python implementation of CRYSTALS-Dilithium too.

I suppose then this repo should be called crystals-py but I wont get ahead of myself.

Discussion of Implementation

Kyber

TODO:

Add some more information about how working with Kyber works with this
library...

Polynomials

The file polynomials.py contains the classes PolynomialRing and Polynomial. This implements the univariate polynomial ring

$$ R_q = \mathbb{F}_q[X] /(X^n + 1) $$

The implementation is inspired by SageMath and you can create the ring $R_{11} = \mathbb{F}_{11}[X] /(X^8 + 1)$ in the following way:

Example

>>> R = PolynomialRing(11, 8)
>>> x = R.gen()
>>> f = 3*x**3 + 4*x**7
>>> g = R.random_element(); g
5 + x^2 + 5*x^3 + 4*x^4 + x^5 + 3*x^6 + 8*x^7
>>> f*g
8 + 9*x + 10*x^3 + 7*x^4 + 2*x^5 + 5*x^6 + 10*x^7
>>> f + f
6*x^3 + 8*x^7
>>> g - g
0

We additionally include functions for PolynomialRing and Polynomial to move from bytes to polynomials (and back again).

• PolynomialRing
  • parse(bytes) takes $3n$ bytes and produces a random polynomial in $R_q$
  • decode(bytes, l) takes $\ell n$ bits and produces a polynomial in $R_q$
  • cbd(beta, eta) takes $\eta \cdot n / 4$ bytes and produces a polynomial in $R_q$ with coefficents taken from a centered binomial distribution
• Polynomial
  • self.encode(l) takes the polynomial and returns a length $\ell n / 8$ bytearray

Example

>>> R = PolynomialRing(11, 8)
>>> f = R.random_element()
>>> # If we do not specify `l` then it is computed for us (minimal value)
>>> f_bytes = f.encode()
>>> f_bytes.hex()
'06258910'
>>> R.decode(f_bytes) == f
True
>>> # We can also set `l` ourselves
>>> f_bytes = f.encode(l=10)
>>> f_bytes.hex()
'00180201408024010000'
>>> R.decode(f_bytes, l=10) == f
True

Lastly, we define a self.compress(d) and self.decompress(d) method for polynomials following page 2 of the specification

$$ \textsf{compress}_q(x, d) = \lceil (2^d / q) \cdot x \rfloor \textrm{mod}^+ 2^d, $$

$$ \textsf{decompress}_q(x, d) = \lceil (q / 2^d) \cdot x \rfloor. $$

The functions compress and decompress are defined for the coefficients of a polynomial and a polynomial is (de)compressed by acting the function on every coefficient. Similarly, an element of a module is (de)compressed by acting the function on every polynomial.

Example

>>> R = PolynomialRing(11, 8)
>>> f = R.random_element()
>>> f
9 + 3*x + 5*x^2 + 2*x^3 + 9*x^4 + 10*x^5 + 6*x^6 + x^7
>>> f.compress(1)
x + x^2 + x^6
>>> f.decompress(1)
6*x + 6*x^2 + 6*x^6

Note: compression is lossy! We do not get the same polynomial back by computing f.compress(d).decompress(d). They are however close. See the specification for more information.

Number Theoretic Transform

TODO:

This is now handled by `NTTHelper` which is passed to `PolynomialRing`
and has functions which are accessed by `Polynomial`.

Talk about what is available, and how they are used.

Modules

The file modules.py contains the classes Module and Matrix. A module is a generalisation of a vector space, where the field of scalars is replaced with a ring. In the case of Kyber, we need the module with the ring $R_q$ as described above.

Matrix allows elements of the module to be of size $m \times n$ but for Kyber, we only need vectors of length $k$ and square matricies of size $k \times k$.

As an example of the operations we can perform with out Module lets revisit the ring from the previous example:

Example

>>> R = PolynomialRing(11, 8)
>>> x = R.gen()
>>>
>>> M = Module(R)
>>> # We create a matrix by feeding the coefficients to M
>>> A = M([[x + 3*x**2, 4 + 3*x**7], [3*x**3 + 9*x**7, x**4]])
>>> A
[    x + 3*x^2, 4 + 3*x^7]
[3*x^3 + 9*x^7,       x^4]
>>> # We can add and subtract matricies of the same size
>>> A + A
[  2*x + 6*x^2, 8 + 6*x^7]
[6*x^3 + 7*x^7,     2*x^4]
>>> A - A
[0, 0]
[0, 0]
>>> # A vector can be constructed by a list of coefficents
>>> v = M([3*x**5, x])
>>> v
[3*x^5, x]
>>> # We can compute the transpose
>>> v.transpose()
[3*x^5]
[    x]
>>> v + v
[6*x^5, 2*x]
>>> # We can also compute the transpose in place
>>> v.transpose_self()
[3*x^5]
[    x]
>>> v + v
[6*x^5]
[  2*x]
>>> # Matrix multiplication follows python standards and is denoted by @
>>> A @ v
[8 + 4*x + 3*x^6 + 9*x^7]
[        2 + 6*x^4 + x^5]

We also carry through Matrix.encode() and Module.decode(bytes, n_rows, n_cols) which simply use the above functions defined for polynomials and run for each element.

Example

We can see how encoding / decoding a vector works in the following example. Note that we can swap the rows/columns to decode bytes into the transpose when working with a vector.

>>> R = PolynomialRing(11, 8)
>>> M = Module(R)
>>> v = M([R.random_element() for _ in range(2)])
>>> v_bytes = v.encode()
>>> v_bytes.hex()
'd'
>>> M.decode(v_bytes, 1, 2) == v
True
>>> v_bytes = v.encode(l=10)
>>> v_bytes.hex()
'a014020100103004000040240a03009030080200'
>>> M.decode(v_bytes, 1, 2, l=10) == v
True
>>> M.decode(v_bytes, 2, 1, l=10) == v.transpose()
True
>>> # We can also compress and decompress elements of the module
>>> v
[5 + 10*x + 4*x^2 + 2*x^3 + 8*x^4 + 3*x^5 + 2*x^6, 2 + 9*x + 5*x^2 + 3*x^3 + 9*x^4 + 3*x^5 + x^6 + x^7]
>>> v.compress(1)
[1 + x^2 + x^4 + x^5, x^2 + x^3 + x^5]
>>> v.decompress(1)
[6 + 6*x^2 + 6*x^4 + 6*x^5, 6*x^2 + 6*x^3 + 6*x^5]

Baby Kyber

A great resource for learning Kyber is available at Approachable Cryptography.

We include code corresponding to their example in baby_kyber.py.