{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c00661f3-96d7-4100-86c6-4f9f8aebb2d7",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mikolajlh/pqlattice/blob/main/docs/examples/constructing_and_attacking_NTRU_pkc.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e64c1dd7-561e-44c4-bacd-cbbc1464cbe4",
   "metadata": {},
   "source": [
    "# Constructing and attacking NTRU public key cryptosystem\n",
    "In contrast to cryptosystems such as RSA, Diffie-Hellman or ECC which are based on Group operations, the NTRU cryptosystem is a ring-based cryptosystem, specifically convolution polynomial rings, but it's underlying hard mathematical problem can also be interpreted as **SVP** or **CVP** in a lattice."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1c8afb8-3dec-4c71-a185-d6e8b1d3d99a",
   "metadata": {},
   "source": [
    "# NTRUEncrypt\n",
    "The NTRU public key cryptosystem, parameterized by integers $N$, $p$, $q$ and $d$ such that\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\text{1.}&\\; N \\in \\mathbb{P} \\; \\text{is prime} \\newline\n",
    "\\text{2.}&\\; \\gcd(N,q) = \\gcd(p,q) = 1 \\newline\n",
    "\\text{3.}&\\; d > 0 \\newline\n",
    "\\text{4.}&\\; q > (6d + 1)p\n",
    "\\end{aligned}\n",
    "$$\n",
    "is based on operations in three **polynomial rings** $R, R_p, R_q$\n",
    "$$\n",
    "R = \\frac{\\mathbb{Z}[X]}{X^{N} - 1}, \\quad R_p = \\frac{\\mathbb{Z}_p[X]}{X^{N} - 1}, \\quad R_q = \\frac{\\mathbb{Z}_q[X]}{X^{N} - 1}\n",
    "$$\n",
    "and **ternary polynomials** defined for any positive integers $d_1$ and $d_2$ as\n",
    "$$\n",
    "\\mathcal{T}(d_1, d_2) = \\left\\{ a(x) \\in R \\; \\middle| \\;\n",
    "\\begin{array}{l}\n",
    "a(x) \\; \\text{has} \\; d_1 \\; \\text{coefficients equal to 1,}\\newline\n",
    "a(x) \\; \\text{has} \\; d_2 \\; \\text{coefficients equal to -1,}\\newline\n",
    "a(x) \\; \\text{has all other coefficients equal to 0}\n",
    "\\end{array}\n",
    "\\right\\}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f494f2ca-7411-4a1b-8189-132b5cdfef2e",
   "metadata": {},
   "source": [
    "# Key creation \n",
    "Let $(N, p, q, d)$ be publicly known parameters of the NTRU cryptosystem, chosen by some trusted authority.  \n",
    "**Private  key** consists of two randomly chosen polynomials\n",
    "$$\n",
    "f(x) \\in \\mathcal{T}(d + 1, d) \\qquad \\text{and} \\qquad g(x) \\in \\mathcal{T}(d, d)\n",
    "$$\n",
    "To set up public key we first need to calculate inverses\n",
    "$$\n",
    "F_q(x) = f(x)^{-1} \\; \\text{in} \\; R_q \\qquad \\text{and} \\qquad F_p = f(x)^{-1} \\; \\text{in} \\; R_p\n",
    "$$\n",
    "**The $f(x)$ must be chosen such that inverses $F_q(x)$ and $F_p(x)$ exists**.\n",
    "\n",
    "The polynomial $h(x)$ defined as\n",
    "$$\n",
    "h(x) = F_q(x) \\cdot g(x) \\; \\text{in} \\; R_q\n",
    "$$\n",
    "will be a **public key**\n",
    "\n",
    "Private key: $(f(x), g(x))$  \n",
    "Public key: $h(x)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "733a74e9-31af-4fd4-b030-3b55575b069c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pqlattice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "25a9aaa7-8bc4-4681-9bbd-6cd79d961f90",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pqlattice as pq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f4b31d66-99d3-4a6b-93b6-245343f46f78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's choose NTRU parameters and private key (f,g)\n",
    "N, p, q, d = 7, 3, 41, 2\n",
    "Rq = pq.polynomial.construct_ring(\"-\", N, q) #(X^N + 1)\n",
    "Rp = pq.polynomial.construct_ring(\"-\", N, p) #(X^N + 1)\n",
    "\n",
    "f = pq.as_integer([-1, 0, 1, 1, -1, 0, 1]) #f(x) = x^6 - x^4 + x^3 + x^2 - 1\n",
    "g = pq.as_integer([0, -1, -1, 0, 1, 0, 1]) #g(x) = x^6 + x^4 - x^2 - x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "084885e9-79b9-4784-b91c-73ff36b42794",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "f^-1 in Rq:\n",
      "[37 2 40 21 31 26 8]\n",
      "f * f^-1: [1]\n",
      "\n",
      "f^-1 in Rp:\n",
      "[1 1 1 1 0 2 1]\n",
      "f * f^-1: [1]\n",
      "\n",
      "h = Fq * g:\n",
      "[30 26 8 38 2 40 20]\n",
      "\n",
      "Private key: (f(x), Fp(x))\n",
      "f(x) =  [-1 0 1 1 -1 0 1]\n",
      "g(x) =  [0 -1 -1 0 1 0 1]\n",
      "\n",
      "Public key: h(x)\n",
      "h(x) =  [30 26 8 38 2 40 20]\n"
     ]
    }
   ],
   "source": [
    "# Now lets compute inverses Fq and Fp\n",
    "Fq = Rq.inv(f)\n",
    "Fp = Rp.inv(f)\n",
    "\n",
    "# Lets check if the inverses are correct\n",
    "print(\"f^-1 in Rq:\")\n",
    "print(Fq)\n",
    "print(\"f * f^-1:\", Rq.mul(f, Fq))\n",
    "print()\n",
    "print(\"f^-1 in Rp:\")\n",
    "print(Fp)\n",
    "print(\"f * f^-1:\", Rp.mul(f, Fp))\n",
    "print()\n",
    "\n",
    "\n",
    "# Now let's calculate public key h\n",
    "h = Rq.mul(Fq, g)\n",
    "print(\"h = Fq * g:\")\n",
    "print(h)\n",
    "print()\n",
    "\n",
    "print(\"Private key: (f(x), Fp(x))\")\n",
    "print(\"f(x) = \", f)\n",
    "print(\"g(x) = \", g)\n",
    "print()\n",
    "print(\"Public key: h(x)\")\n",
    "print(\"h(x) = \", h)\n",
    "\n",
    "sk = (f, Fp)\n",
    "pk = h"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "716aa902-a7aa-4cf3-abaf-162f3463cecd",
   "metadata": {},
   "source": [
    "# Encryption\n",
    "Plaintext $m$ has a form of polynomial $m(x) \\in R$ whose coefficients satisfy $ \\; -\\frac{1}{2}p < m_i \\leq \\frac{1}{2}p $.  \n",
    "In order to encrypt $m$ we also need a random polynomial $r(x) \\in \\mathcal{T}(d,d)$ then\n",
    "$$\n",
    "e(x) \\equiv p h(x) \\cdot r(x) + m(x) \\quad \\mod q\n",
    "$$\n",
    "is a **ciphertext** and $e(x) \\in R_q$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "830ea18c-3a7e-4ca5-bbf7-622f89cc9145",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "plaintext m(x) = [1 -1 1 1 0 -1]\n",
      "random element r(x) = [-1 1 0 0 0 -1 1]\n",
      "\n",
      "ciphertext e(x) = [25 3 40 2 4 19 31]\n"
     ]
    }
   ],
   "source": [
    "# plaintext\n",
    "m = pq.as_integer([ 1,-1, 1, 1, 0,-1])\n",
    "print(\"plaintext m(x) =\", m)\n",
    "# random element\n",
    "r = pq.as_integer([-1, 1, 0, 0, 0,-1, 1])\n",
    "print(\"random element r(x) =\", r) \n",
    "\n",
    "print()\n",
    "# ciphertext\n",
    "e = Rq.add(p * Rq.mul(r, pk), m)\n",
    "print(\"ciphertext e(x) =\", e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5a285d1-d2c1-426d-a3c5-97c02276654d",
   "metadata": {},
   "source": [
    "# Decryption\n",
    "To decrypt ciphertext $e(x)$ we first need to compute\n",
    "$$\n",
    "a(x) \\equiv f(x) \\cdot e(x) \\quad \\mod q\n",
    "$$\n",
    "then\n",
    "$$\n",
    "b(x) \\equiv F_p(x) \\cdot \\text{center\\_lift} (a(x)) \\quad \\mod p\n",
    "$$\n",
    "is equal to the plaintext $m(x)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3c4d6e0b-0647-485a-8888-1b93c8d6b734",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Decrypted message d:\n",
      "[1 -1 1 1 0 -1]\n",
      "Original message m:\n",
      "[1 -1 1 1 0 -1]\n"
     ]
    }
   ],
   "source": [
    "a = Rq.center_lift(Rq.mul(sk[0], e))\n",
    "d = Rp.center_lift(Rp.mul(sk[1], a))\n",
    "print(\"Decrypted message d:\")\n",
    "print(d)\n",
    "print(\"Original message m:\")\n",
    "print(m)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ece18e76-49ab-47b2-9823-79be4c9cfd2e",
   "metadata": {},
   "source": [
    "# NTRU as a lattice\n",
    "Given NTRU cryptosystem with parameters $(N,p,q,d)$ we are going to identify each pair of polynomials\n",
    "$$\n",
    "a(x) = a_0 + \\ldots + a_{N-1}x^{N-1} \\qquad \\text{and} \\qquad b(x) = b_0 + \\ldots + b_{N-1}x^{N-1}\n",
    "$$\n",
    "in $R$ with $2N$-dimensional vector\n",
    "$$\n",
    "(a,b) = (a_0, \\ldots, a_{N-1}, b_0, \\ldots, b_{N-1}) \\in \\mathbb{Z}^{2N}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bee31610-3d14-4d4f-a7d2-befb2ed80acd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "28718743-6bd7-4b9c-ba3b-3c52bc40dd50",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NTRU as lattice:\n",
      "Matrix of integers with shape: 14 x 14\n",
      "======================================\n",
      "      [0]  [1]  [2]  [3]  [4]  [5]  [6]  [7]  [8]  [9]  [10]  [11]  [12]  [13]\n",
      " [0]    1    0    0    0    0    0    0   30   26    8    38     2    40    20\n",
      " [1]    0    1    0    0    0    0    0   20   30   26     8    38     2    40\n",
      " [2]    0    0    1    0    0    0    0   40   20   30    26     8    38     2\n",
      " [3]    0    0    0    1    0    0    0    2   40   20    30    26     8    38\n",
      " [4]    0    0    0    0    1    0    0   38    2   40    20    30    26     8\n",
      " [5]    0    0    0    0    0    1    0    8   38    2    40    20    30    26\n",
      " [6]    0    0    0    0    0    0    1   26    8   38     2    40    20    30\n",
      " [7]    0    0    0    0    0    0    0   41    0    0     0     0     0     0\n",
      " [8]    0    0    0    0    0    0    0    0   41    0     0     0     0     0\n",
      " [9]    0    0    0    0    0    0    0    0    0   41     0     0     0     0\n",
      "[10]    0    0    0    0    0    0    0    0    0    0    41     0     0     0\n",
      "[11]    0    0    0    0    0    0    0    0    0    0     0    41     0     0\n",
      "[12]    0    0    0    0    0    0    0    0    0    0     0     0    41     0\n",
      "[13]    0    0    0    0    0    0    0    0    0    0     0     0     0    41\n"
     ]
    }
   ],
   "source": [
    "L11 = pq.as_integer(np.identity(N))\n",
    "L12 = pq.as_integer([pq.polynomial.poly.pad(Rq.mul(pq.polynomial.poly.monomial(1, i), h), N - 1) for i in range(N)])\n",
    "L21 = pq.as_integer(np.zeros((N,N)))\n",
    "L22 = q * pq.as_integer(np.identity(N))\n",
    "\n",
    "NTRU_B = np.block([[L11, L12], [L21, L22]])\n",
    "print(\"NTRU as lattice:\")\n",
    "pq.show(NTRU_B, max_rows=14, max_cols=14)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4efd7e87-6267-4ba5-9fe8-97df64d4f860",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Short vector from LLL reduced basis:\n",
      "[1 0 -1 1 0 -1 -1 -1 0 -1 0 1 1 0]\n",
      "f'^-1 in Rq:\n",
      "[20 10 15 33 4 39 1]\n",
      "f'^-1 * f': [1]\n",
      "\n",
      "f'^-1 in Rp:\n",
      "[2 0 1 2 2 2 2]\n",
      "f'^-1 * f': [1]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "lll_B = pq.lattice.lll(NTRU_B)\n",
    "w = lll_B[1]\n",
    "print(\"Short vector from LLL reduced basis:\")\n",
    "print(w)\n",
    "f_prim = w[:N]\n",
    "g_prim = w[N:]\n",
    "\n",
    "f_prim_inv_q = Rq.inv(f_prim)\n",
    "f_prim_inv_p = Rp.inv(f_prim)\n",
    "\n",
    "print(\"f'^-1 in Rq:\")\n",
    "print(f_prim_inv_q)\n",
    "print(\"f'^-1 * f':\", Rq.mul(f_prim_inv_q, f_prim))\n",
    "print()\n",
    "\n",
    "print(\"f'^-1 in Rp:\")\n",
    "print(f_prim_inv_p)\n",
    "print(\"f'^-1 * f':\", Rp.mul(f_prim_inv_p, f_prim))\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0ab8fb25-0499-457f-9cdf-cc7b61ae33e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Decrypted message d:\n",
      "[1 -1 1 1 0 -1]\n",
      "Original message m:\n",
      "[1 -1 1 1 0 -1]\n"
     ]
    }
   ],
   "source": [
    "a_prim = Rq.center_lift(Rq.mul(f_prim, e))\n",
    "d_prim = Rp.center_lift(Rp.mul(f_prim_inv_p, a_prim))\n",
    "\n",
    "print(\"Decrypted message d:\")\n",
    "print(d_prim)\n",
    "print(\"Original message m:\")\n",
    "print(m)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
