Building Post-Quantum Cryptography from Scratch: CRYSTALS-Kyber in Python (Part 1/n)

We will explore how to implement a post-quantum cryptography algorithm(s) CRYSTALS. We start the journey by introducing some needed tools and in the following posts implement the actual algorithm. If you are interested in the dirty details how to turn math to code, this is for you.

[This article was written in another platform first and I had some painful time to get it rendering even remotely correct here. I hope it is readable enough.]

Introduction

CRYSTALS ("Cryptographic Suite for Algebraic Lattices") is post-quantum public key encryption scheme, meaning it is expected to be secure even at the era of quantum computing where many current PKE-variants fail.

CRYSTALS consists of two cryptographic primitives: Kyber and Dilithium. Kyber is key exchange method, i.e. asymmetric encryption providing secure channel to change secrets, and Dilithium is a cryptographic signing method. We will explore the mathematics behind these algorithms by coding them in Python as we go. You can find the code from my GitHub repo.

The reader is assumed certain maturity in mathematics and basic understanding of Python. We don't prove anything, instead the focus is to introduce and build needed machinery that we will use later.
The mathematical part of this presentation follows closely the way Alfred Menezes presents it in his excellent video series on the topic.

Modulo Algebra

When we say two integers $a$ and $b$ are congruent modulo $q,$ we mean $b-a$ is a integer multiple of $q.$ In this case we write

$$a\equiv b\;(\bmod\; q).$$

With $r = a\;(\bmod\; q)$ we mean $r$ is the remainder of integer $a$ divided $q.$ This implies $0\leq r \lt q.$

$\mathbb{Z}_q = {0,1,2,...,q-1}$ is the ring of integers modulo $q.$ In this ring addition and multiplication are performed modulo $q.$

We implement integers in $\mathbb{Z}_q$ in class Zq. Notice the Python modulo-operation % is implemented in a way that is fully compatible with our needs because it can handle negative values correctly. The instantiation can be done with integer value or with an instance of Zq.

class Zq:

    def __init__(self, q, value):
# Sanity checks cleaned
        if isinstance(value,int):
            # Python modulo operation is one of the few correctly implemented
            self.value = value % q
        else:
            self.value = value.get_value() % q
        self.q = q

The class Zq has addition, subtraction, multiplication, str and repr operations implemented. This makes our life a lot easier because we can make arithmetics directly and debug when needed.

To get a feeling how things work, consider the ring $\mathbb{Z}_4={0,1,2,3}.$ For example we have

$$2-9=-7=1\;(\bmod\; 4)$$

$$3+2=5=1\;(\bmod\; 4)$$

and for multiplication

$$3*4=12=0\;(\bmod\; 4)$$

$$3*5=15=3\;(\bmod\; 4).$$

if __name__ == "__main__":
    two = Zq(4,2)
    three = Zq(4,3)
    four = Zq(4,4)
    five = Zq(4,5)
    nine = Zq(4,9)
    print(f"2-9={two-nine}")
    print(f"3+2={three+two}")
    print(f"3*4={three*four}")
    print(f"3*5={three*five}")

(.venv) ..> python Zq.py
2-9=1
3+2=1
3*4=0
3*5=3

Polynomial Rings

Let $q$ be a prime. We define $\mathbb{Z}_q[x]$ to be the set of polynomials of $x$ with all coefficients in the ring $\mathbb{Z}_q.$ This means all coefficient arithmetic is performed in the ring $\mathbb{Z}_q.$

We implement polynomials in the ring with a class ZqPolynomial. Here coefficients is a list of integers, and the length of the list defines $n.$

class ZqPolynomial:

    def __init__(self, q, coefficients):
# Sanity checks cleaned
        self.q = q
        self.coefficients = coefficients
        for i in range(len(coefficients)):
            self.coefficients[i] = Zq(self.q, self.coefficients[i])
        # n from (x^n+1)
        self.n = len(coefficients)

For example, let $q=5,$

$$f(x)=3+2x+4x^2+3x^3$$

and

$$h(x)=2+x+4x^2+4x^3+x^4.$$

We get

$$f(x)+h(x)=3x+3x^2+2x^3+x^4,$$

$$f(x)-h(x)=1+x+4x^3+4x^4$$

$$f(x)h(x)=1+2x+2x^2+x^6+3x^7.$$

We can do this with our code as follows (we use extra zeroes in coefficients to prevent the polynomial modulo operation).

if __name__ == "__main__":
    f = ZqPolynomial(5, [3, 2, 4, 3, 0,0,0,0])
    h = ZqPolynomial(5, [2, 1, 4, 4, 1,0,0,0])
    print(f"f(x)+h(x)={f+h}")
    print(f"f(x)-h(x)={f-h}")
    print(f"f(x)h(x)={f*h}")

> python zq_polynomial.py
f(x)+h(x)=3x+3x^2+2x^3+x^4
f(x)-h(x)=1+x+4x^3+4x^4
f(x)h(x)=1+2x+2x^2+x^6+3x^7

Let now $q$ be a prime and $n$ a positive integer. The quotient ring (often called just "polynomial ring") $R_q = \mathbb{Z}_q[x]/(x^n+1)$ consists of polynomials in $\mathbb{Z}_q$ of degree less than $n.$ In ring $R_q$ the multiplication of polynomials is performed modulo the polynomial $x^n+1$ called the reduction polynomial. This means that the product of polynomials $f(x),g(x)\in \mathbb{Z}_q[x]$ is defined as the remainder $r(x)$ of their product when divided by $x^n+1$ in the polynomial ring. Notice that by definition now degree of $r(z)$ is at most $n-1$ and

$$f(x)g(x)= r(x)\in\mathbb{Z}_q[x].$$

One should notice here that remainder is not calculated by the traditional polynomial division algorithm, but with division rules that apply in the polynomial ring. For our purposes it suffices to acknowledge that if the polynomial has degrees $\geq n$ you can apply the rules

$$x^n\equiv -1$$

$$x^{(n+1)}\equiv -x$$

and in general for $k\geq n,$ $x^{(n+k)}\equiv -x^k,$ and then simplify the resulting polynomial normally. To understand why, please visit ring theory and ideals.

Overloading addition and subtraction is straightforward, but multiplication needs special treatment. Here we utilize the fact that Zq has multiplication operation overloaded. In real-life implementations this naive implementation is too slow and NTT-algorithm is used instead. We will return to this later.

    def __mul__(self, other): # Again, sanity checks removed
        n = len(self.coefficients)
        # Initialize result coefficients with zeros
        result = [Zq(self.q, 0) for _ in range(n)]

        # Perform polynomial multiplication
        for i in range(n):
            for j in range(n):
                # Position in the product
                pos = i + j
                # If position exceeds n-1, we need to reduce modulo x^n + 1
                if pos >= n:
                    # When reducing mod x^n + 1, x^n = -1
                    # So x^(n+k) = -x^k
                    pos = pos - n
                    result[pos] -= self.coefficients[i] * other.coefficients[j]
                else:
                    result[pos] += self.coefficients[i] * other.coefficients[j]
        return ZqPolynomial(self.q, result)

For example, consider the ring

$$\mathbb{Z}_5[x]/(x^4+1)$$

and polynomials

$$f(x)=1+3x+4x^2+x^3$$

and

$$g(x)=x+3x^2+3x^3.$$

To get the reminder of $f(x)g(x)$ when divided $x^4+1$ we first calculate the product

$$f(x)g(x)=x+6x^2+16x^3+22x^4+15x^5+3x^6$$

and use the substitution rule to get

$$f(x)g(x)=-22-14x+3x^2+16x^3$$

and with the modulo operations we arrive at

$$f(x)g(x)=3+x+3x^2+x^3\in\mathbb{Z}_5[x]/(x^4+1).$$

With our code we get directly as follows.

if __name__ == "__main__":
    f = ZqPolynomial(5,[1,3,4,1])
    g = ZqPolynomial(5,[0,1,3,3])
    print(f"f*g={f*g}")

>python zq_polynomial.py
f*g=3+x+3x^2+x^3

Polynomials as Vectors

For a programmer it is rather straightforward to see that the polynomial can be represented as vectors. Consider the polynomial

$$f(x)=a_0+a_1x+a_2x^2+\cdots+a_{n-1}x^{n-1}\in\mathbb{Z}_q[x]/(x^n+1).$$

The obvious way to write that as a vector is $f=[a_0,a_1,...,a_{n-1}]^T$ (convention is to use column vectors). The polynomial addition is now component-wise addition modulo $q,$ and subtraction is component-wise subtraction modulo $q.$ Multiplication is polynomial multiplication as shown earlier and the resulting polynomial is stored in a vector.

We used this implicitly earlier when defining ZqPolynomial.

We extend this notation. Let $k>0$ be a positive integer. Module $R^k_q$ consists of length $k$ vectors of polynomials of $R_q$ . Addition and subtraction is again component-wise. It follows that the resulting vectors in both cases is in $R_q^k.$

We will later use $R_q^{k\times l}$ to represent $k\times l$ -matrices with polynomial entries. This extension is similar to that from vectors to matrices in linear algebra.

The module $R_q^k$ includes $k\times 1$ -matrices with elements from the polynomial ring. The easiest way to handle these is to define class PolyMatrix that holds ZqPolynomials in a list of lists.

The multiplication in $R_q^k$ is defined as inner product of two vectors in $R_q^k.$ This means that the polynomials, that are elements of the vector in $R_q^k,$ are multiplied in $R_q$ and added together. The result is in $R_q.$

For example, let

$$A=\begin{pmatrix}4+x+x^2 \newline x+3x^2\end{pmatrix}\in R_5^2$$

and

$$B=\begin{pmatrix}1+2x+3x^2 \newline 4x\end{pmatrix}\in R_5^2.$$

We get directly

$$A+B=\begin{pmatrix} 3x \newline 3x^2\end{pmatrix}\in R_5^2$$

$$A-B=\begin{pmatrix} 3+4x+4x^2 \newline 2x+3x^2\end{pmatrix}\in R_5^2$$

and in Python as follows.

if __name__ == "__main__":
    A = PolyMatrix(2,1,5,3)
    A[0,0] = ZqPolynomial(5,[4,1,2])
    A[1,0] = ZqPolynomial(5,[0,1,3])

    B = PolyMatrix(2,1,5,3)
    B[0,0] = ZqPolynomial(5,[1,2,3])
    B[1,0] = ZqPolynomial(5,[0,4,0])

    print(f"A+B={A+B}")
    print(f"A-B={A-B}")

(.venv)...> python poly_matrix.py
A+B=[3x]
[3x^2]
A-B=[3+4x+4x^2]
[2x+3x^2]

Using the $A$ and $B$ defined earlier, we get

$$A^TB=(4+x+x^2,x+3x^2)\begin{bmatrix}1+2x+3x^2\newline 4x\end{bmatrix}$$

$$=(4+x+x^2)(1+2x+3x^2)+(x+3x^2)(1+2x+3x^2)$$

$$=3x.$$

And in Python.

if __name__ == "__main__":
    A = PolyMatrix(2,1,5,3)
    A[0,0] = ZqPolynomial(5,[4,1,2])
    A[1,0] = ZqPolynomial(5,[0,1,3])

    B = PolyMatrix(2,1,5,3)
    B[0,0] = ZqPolynomial(5,[1,2,3])
    B[1,0] = ZqPolynomial(5,[0,4,0])
    print(f"A.T@B={A.T@B}")

(.venv)..>python poly_matrix.py
A.T@B=[3x]

To enable matrix multiplication, we implemented the matmul operator.

    def __matmul__(self, other):  # enables @ operator
        result = PolyMatrix(self.rows, other.cols, self.q, self.n)

        for i in range(self.rows):
            for j in range(other.cols):
                # Sum of products for this position
                for k in range(self.cols):
                    result[i, j] += self[i, k] * other[k, j]

        return result

We can use the bracket-notation with PolyMatrix because we defined getitem and setitem methods.

Size of Polynomials

Next we need notion of size that will become useful later. First let us define symmetric mod.

Let $q\in\mathbb{Z}$ be odd and $r\in\mathbb{Z}_q$ . We define

$$r\bmod_s q = \begin{cases} r & \mathrm{if\ } r\leq (q-1)/2 \newline r-q & \mathrm{if } r > (q-1)/2. \end{cases}$$

This immediately gives $-(q-1)/2\leq r\bmod_s q\leq (q-1)/2.$ Here $\bmod_s$ is symmetric modulo operation.

Let $q=11.$ We now have by definition $-5\leq r\bmod_s 11\leq 5.$

For example

$$\begin{align*}4\bmod_s 11 &=4 \newline 8\bmod_s 11 &= -3.\end{align*}$$

The definition of symmetric modulo is slightly different for even $q.$ Let $q\in\mathbb{Z}$ be even and $r\in\mathbb{Z}_q.$

Then

$$r\bmod_s q = \begin{cases} r & \mathrm{if\ } r\leq q/2 \newline r-q & \mathrm{if\ } r > q/2.\end{cases}$$

and we get $-q/2< r\bmod_s q\leq q/2$ .

In the code we implement the symmetric modulo in Zq and use that from ZqPolynomial and PolyMatrix.

    def to_symmetric(self):
        """Convert value to symmetric representation in [-(q-1)/2, q/2]"""
        if self.q % 2 == 0:  # even q
            if self.value >= self.q // 2:
                self.value = self.value - self.q
        else:  # odd q
            if self.value > self.q // 2:
                self.value = self.value - self.q
        return self

Let $q\in\mathbb{Z}.$ We define

$$|r|{\infty} = |r\bmod_s q|.$$

For example, let

$q=11.$

Then

$$\begin{align*} |5|{\infty} &= 5 \newline |8|{\infty} &=|-3|= 3\newline |10|{\infty} &= |-1| = 1.\end{align*}$$

You can think this as "clock-algebra", the further the integer is from noon, the bigger its norm.

We have the following immediate corollaries

$$0\leq |r|{\infty}\leq (q-1)/2$$

if q is odd and

$$0\leq |r|{\infty}\leq q/2$$

if q is even.

For polynomial ring elements the size of a polynomial is defined with the maximum operation. Let

$$f(x)=a_0+a_1x+a_2x^2+a_3x^3+\cdots +a_{n-1}x^{n-1}.$$

We define

$$|f(x)|{\infty} = \max_i |a_i|{\infty}.$$

For example, let $q=11,$ $n=4$ and

$$f(x)=8+x+5x^2+3x^3\in\mathbb{Z}q/(x^n+1).$$

Then $|f(x)|{\infty}=5$ because

$$\begin{align*} |a_0|{\infty} &= |8\bmod_s 11| &=|-3| &= 3\newline |a_1|{\infty} &= |1\bmod_s 11| &=|1| &= 1\newline |a_2|{\infty} &= |5\bmod_s 11| &=|5| &= 5\newline |a_3|{\infty} &= |3\bmod_s 11| &=|3| &= 3\end{align*}$$

and clearly $\max_i |a_i|_{\infty} =5.$

This definition can be generalized to elements of module $R_q^k.$ Let

$$a=[a_1,a_2,...,a_k]^T\in R_q^k.$$

We define

$$|a|{\infty}=\max_i|a_i|{\infty}.$$

Small Polynomials

We say a polynomial $f$ is small if $|f|_{\infty}$ is small. Notice that this means that all coefficients of $f$ need to be small due the way the norm is defined. What "small" means is defined per context.

Let $\eta$ be a positive integer less than $q/2.$ We define

$$S_\eta = \lbrace f\in R_q\,|\,|f|{\infty}\leq \eta\rbrace$$

to be the set of polynomials in $R_q$ where each polynomials each coefficient is of size at most $\eta.$ We use $S\eta$ to define a set of "small" polynomials.

For example consider polynomial

$$f(x)=10+9x+9x^2+2x^3\in R_{11}.$$

Now

$$|f(x)|_{\infty}=2$$

and hence $f(x)\in S_2.$

Observe that $S_1$ is the set of polynomials in $R_q$ with all coefficients in the set $\lbrace -1,0,1\rbrace$ (when reduced $\bmod_s q$ ).

Let $f(x)\in S_{\eta_1}$ and $g(x)\in S_{\eta_2}.$ Without proof we state that $f(x)g(x)\in S_{n\eta_{1}\eta_{2}}.$ This generalizes to vectors or polynomials. Let $a\in S^{k}{\eta{1}}$ and $b\in S^{k}{\eta_2}.$ Then we have $ab^T\in S{kn\eta_{1}\eta_{2}}.$

Next steps

In the next article we utilize the presented machinery to implement basic CRYSTALS-Kyber public key encryption. Stay tuned.