Repository files navigation

This repository contains a pure python implementation of both:

1. ML-KEM: The NIST Module-Lattice-Based Key-Encapsulation MechanismStandard followingFIPS 203from the NIST post-quantum cryptography project.
2. CRYSTALS-Kyber: following (at the time of writing) the most recentspecification(v3.02)

Note: This project accompaniesdilithium-py which is apure-python implementation of ML-DSA and CRYSTALS-Dilithium and shares a lot ofthe lower-level code of this implementation.

kyber-py has been written as an educational tool. The goal of this project wasto learn about how Kyber works, and to try and create a clean, well commentedimplementation which people can learn from.

This code is not constant time, or written to be performant. Rather, it waswritten so that the python code closely follows the Kyber specificationspecification andFIPS 203. No cryptographic guarantees are made of this work.

This package is available askyber-py onPyPI:

pip install kyber-py

This work started by simply implementing Kyber for fun, however after NISTpicked Kyber to standardise as ML-KEM, the repository grew and now includes bothimplementations of Kyber and ML-KEM. I assume as this repository ages, the Kyberimplementation will get less useful and the ML-KEM one will be the focus, butfor historical reasons we will include both. If only so that people can studythe differences which NIST introduced during the standardisation of theprotocol.

This implementation currently passes all KAT tests forkyber andml_kem Formore information, see the unit tests intest_kyber.pyandtest_ml_kem.py.

The KAT files were either downloaded or generated:

1. ForML-KEM, the KAT files were download from the GitHub repositoryusnistgov/ACVP-Server/ release 1.1.0.35, and are included inassets/ML-KEM-* directories.
2. ForKyber, the KAT files were generated from the projectsGitHubrepository and are included inassets/PQCLkemKAT_*.rsp

Note: for Kyber v3.02, there is a discrepancy between the specification andreference implementation. To ensure all KATs pass, one has to generate thepublic keybefore the random bytes$z = \mathcal{B}^{32}$ in algorithm 7 ofthespecification(v3.02).

Originally this project was planned to have zero dependencies, however to make this workpass the KATs, we needed a deterministic CSRNG. The reference implementation usesAES256 CTR DRBG. I have implemented this inaes256_ctr_drbg.py.However, I have not implemented AES itself, instead I import this frompycryptodome. If this dependency is too annoying, then please make an issue and we can have a pure-python AES included into the repo.

To install dependencies, runpip -r install requirements.

There are four functions exposed on theML_KEM class which are intended foruse:

• ML_KEM.keygen(): generate a keypair(ek, dk)
• ML_KEM.key_derive(seed): generate a keypair(ek, dk) from the providedseed
• ML_KEM.encaps(ek): generate a key and ciphertext pair(key, ct)
• ML_KEM.decaps(dk, ct): generate the shared keykey

Those, together with theML_KEM_512,ML_KEM_768, andML_KEM_1024objects comprise the kyber-py library stable API.

>>>fromkyber_py.ml_kemimportML_KEM_512>>>ek,dk=ML_KEM_512.keygen()>>>key,ct=ML_KEM_512.encaps(ek)>>>_key=ML_KEM_512.decaps(dk,ct)>>>assertkey==_key

The above example would also work withML_KEM_768 andML_KEM_1024.

All times recorded using a Intel Core i7-9750H CPU and averaged over 1000 runs.

There are three functions exposed on theKyber class which are intended foruse:

• Kyber.keygen(): generate a keypair(pk, sk)
• Kyber.encaps(pk): generate shared key and challenge(key, c)
• Kyber.decaps(sk, c): generate the shared keykey

>>>fromkyber_py.kyberimportKyber512>>>pk,sk=Kyber512.keygen()>>>key,c=Kyber512.encaps(pk)>>>_key=Kyber512.decaps(sk,c)>>>assertkey==_key

The above example would also work withKyber768 andKyber1024.

We expect users to pick one of the three initalised classes which use thedefault parameters of the Kyber specification. The three options areKyber512,Kyber768 andKyber1024. However, by following the values inDEFAULT_PARAMETERS one could tweak these values to look at how Kyber behavesfor different default values.

NOTE: it is relatively easy to change the parameters$k$,$\eta_1$,$\eta_2$$d_u$ and$d_v$ from the Kyber specification. However, if you wish to change thepolynomial ring itself, then you will lose access to the NTT transforms whichcurrently only support$q = 3329$ and$n = 256$.

All times recorded using a Intel Core i7-9750H CPU and averaged over 1000 runs.

• https://kyber-py.readthedocs.io/en/latest/

There are two main things to worry about when implementing Kyber/ML-KEM. Thefirst thing to consider is the mathematics, which requires performing linearalgebra in a module with elements in the ring$R_q = \mathbb{F}_q[X] /(X^n + 1)$and the second is the sampling, compression and decompression, which links tothe cryptographic assurance of the protocol.

For those who don't know, a module is a generalisation of a vector space, whereelements of a matrix are not selected from a field (such as the rationals, orelement of a finite field$\mathbb{F}_{p^k}$), but rather in a ring (we do notrequire each element in a ring to have a multiplicative inverse). The ring in question for Kyber/ML-KEM is a polynomial ring where polynomials have coefficients in$\mathbb{F}_{q}$ with$q = 3329$ and the polynomial ring has a modulus$X^n + 1$ with$n = 256$ (and so every element of the polynomial ring has at most 256 coefficients).

To help with experimenting with these polynomial rings themselves, the filepolynomials_generic.py has an implementation of the univariate polynomial ring

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

where the user can select any$q, n$. For example, you can create thering$R_{11} = \mathbb{F}_{11}[X] /(X^8 + 1)$ in the following way:

>>>fromkyber_py.polynomials.polynomials_genericimportPolynomialRing>>>R=PolynomialRing(11,8)>>>x=R.gen()>>>f=3*x**3+4*x**7>>>g=R.random_element();g5+x^2+5*x^3+4*x^4+x^5+3*x^6+8*x^7>>>f*g8+9*x+10*x^3+7*x^4+2*x^5+5*x^6+10*x^7>>>f+f6*x^3+8*x^7>>>g-g0

We hope that this allows for some hands-on experience at working with thesepolynomials before starting to play with the whole of Kyber/ML-KEM.

For the "Kyber-specific" functions, needed to implement the protocol itself, wehave made a child classPolynomialRingKyber(PolynomialRing) which has thefollowing additional methods:

• PolynomialRingKyber
  • ntt_sample(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
• PolynomialKyber
  • encode(l) takes the polynomial and returns a length$\ell n / 8$ bytearray
  • to_ntt() converts the polynomial into the NTT domain for efficientpolynomial multiplication and returns an element of typePolynomialKyberNTT
• PolynomialKyberNTT
  • from_ntt() converts the polynomial back from the NTT domain and returns anelement of typePolynomialKyber

This class fixes$q = 3329$ and$n = 256$

Lastly, we define aself.compress(d) andself.decompress(d) method forpolynomials following page 2 of thespecification

$$ \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 functionscompress anddecompress are defined for the coefficients of apolynomial and a polynomial is (de)compressed by acting the function on everycoefficient. Similarly, an element of a module is (de)compressed by acting thefunction on every polynomial.

Note: compression is lossy! We do not get the same polynomial back bycomputingf.compress(d).decompress(d). They are howeverclose. See thespecification for more information.

Building onpolynomials_generic.py we also include a filemodules_generic.py which has all ofthe functions needed to perform linear algebra given a ring.

Note thatMatrix allows elements of the module to be of size$m \times n$ butfor Kyber, we only need vectors of length$k$ and square matrices of size$k\times k$.

As an example of the operations we can perform with outModule lets revisitthe ring from the previous 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 matrices 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 coefficients>>>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]

On top of this class, we have the classesModuleKyber(Module) andMatrixKyber(Matrix) which have helper functions which (for example) encodeevery element of a matrix, or convert every element to or from the NTT domain.These are simple functions which call the respectivePolynomialKyber methodsfor every element.