lattice-based-encryption.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Alisa Todorova"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to lattice-based encryption"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tasks\n",
    "1. Write the LWE public-key encryption function\n",
    "2. Write the RLWE public-key encryption and decryption functions.\n",
    "3. Write the attack recovering the secret $s$ from the knowledge of $a$, $t$ and $e$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LWE encryption\n",
    "\n",
    "We first recall basic lattice-based encryption, starting from the Regev scheme [Reg05].\n",
    "\n",
    "Let $q \\in {\\mathbb Z}$ be a prime number. Let $\\vec{s} \\in {\\mathbb Z}^n$ be the secret-key. \n",
    "An LWE ciphertext for a message $m \\in \\{0,1\\}$ is a vector $\\vec{c} \\in {\\mathbb F}_q^n$ such that\n",
    "$$ \\langle \\vec{c},\\vec{s} \\rangle= e + m \\cdot \\lfloor q/2 \\rceil \\pmod{q}$$\n",
    "where for the error $e$, we can take the centered binomial distribution $\\chi$ with\n",
    "parameter $\\kappa$. \n",
    "\n",
    "One can let $e=h(u)-h(v)$\n",
    "where $u,v \\leftarrow \\{0,1\\}^\\kappa$, where $h$ is the Hamming weight function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hw(x):\n",
    "  return sum(x.digits(base=2))\n",
    "\n",
    "def binom(kappa=2):\n",
    "  return hw(ZZ.random_element(2^kappa))-hw(ZZ.random_element(2^kappa))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For simplicity, we first consider the symmetric scheme. We consider a binary secret-key $\\vec{s}=(s_1,\\cdots,s_n)$, with $s_1=1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def genKey(n=10,kappa=2,nq=6):\n",
    "  q=next_prime(2^nq)\n",
    "  s=vector(ZZ,[1]+[ZZ.random_element(2) for i in range(n-1)])\n",
    "  return kappa,q,s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'kappa=2, q=67, s=(1, 1, 0, 0, 0, 0, 0, 1, 0, 1)'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kappa,q,s=genKey()\n",
    "\"kappa=%d, q=%d, s=%s\" % (kappa,q,s.__str__())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An LWE ciphertext is a vector of elements in ${\\mathbb Z}_q$. We have taken $s_1=1$, so \n",
    "$$\\langle \\vec{c},\\vec{s} \\rangle = c_1 + \\sum\\limits_{i=2}^n c_i \\cdot s_i = e + m \\cdot \\lfloor q/2 \\rceil \\pmod{q}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def LWEencSym(mes,q,kappa,s):\n",
    "  a=vector(GF(q),[ZZ.random_element(q) for j in range(s.length())])\n",
    "  a[0]=-a[1:]*s[1:]+binom(kappa)+ZZ((q-1)/2)*mes\n",
    "  return a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Decryption is performed by computing\n",
    "$m={\\sf th}(\\langle \\vec{c},\\vec{s} \\rangle \\bmod q)$, where ${\\sf th}(x)=1$ if $q/4 \\leq x \\leq 3q/4$, and $0$ otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def th(x,q):\n",
    "  if x>=q/4 and x<3*q/4: return 1\n",
    "  return 0\n",
    "\n",
    "def LWEdecrypt(c,q,s):\n",
    "  return th((c*s).lift(),q)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To encrypt from the public-key, one can use a matrix \n",
    "${\\bf A} \\in {\\mathbb F}_q^{\\ell \\times n}$ of row vectors  ${\\vec a}_i \\in {\\mathbb F}_q^n$, such that \n",
    "$$ \\langle \\vec{a}_i,\\vec{s} \\rangle= e_i \\pmod q$$\n",
    "for $e_i \\leftarrow \\chi$, for all $1 \\leq i \\leq \\ell$.\n",
    "\n",
    "This can be written ${\\bf A} \\cdot \\vec{s}= \\vec{e} \\pmod{q}$.\n",
    "Therefore, every row of $\\vec{A}$ is an LWE encryption of $0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vecBinom(n,kappa=2):\n",
    "  return vector([binom(kappa) for i in range(n)])\n",
    "\n",
    "def genKeyPK(n=10,ell=40,kappa=2,nq=12):\n",
    "  q=next_prime(2^nq)\n",
    "  A=Matrix(GF(q),[[ZZ.random_element(q) for j in range(n)] for i in range(ell)])\n",
    "  s=vector(ZZ,[1]+[ZZ.random_element(2) for i in range(n-1)])\n",
    "  A[:,0]=-A[:,1:]*s[1:]+vecBinom(ell,kappa)\n",
    "  return kappa,q,A,s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'kappa=2, q=4099, s=(1, 1, 0, 0, 0, 1, 0, 1, 0, 1)'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kappa,q,A,s=genKeyPK()\n",
    "\"kappa=%d, q=%d, s=%s\" % (kappa,q,s.__str__())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To encrypt a  message $m \\in \\{0,1\\}$, one generates a linear\n",
    "combination of the row vectors ${\\vec a}_i$: \n",
    "\\begin{align*}\n",
    "\\vec{c}  =\\left\\lfloor \\frac{q}{2} \\right\\rceil \\cdot\n",
    "(m,0,\\ldots,0)+ \\vec{u} \\cdot \\vec{A} \\pmod q \n",
    "\\end{align*}\n",
    "where $\\vec{u} \\leftarrow \\chi^\\ell$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original Message: 0\n",
      "Encrypted Message: (511, 828, 3125, 1960, 1959, 3777, 2396, 3914, 1390, 3261)\n",
      "Decrypted Message: 0\n"
     ]
    }
   ],
   "source": [
    "# mes is the message\n",
    "def LWEenc(mes, A, kappa, q):\n",
    "    # Generates vec{u} <- χ^l\n",
    "    u = vecBinom(A.nrows(), kappa) #where A.nrows()=l\n",
    "    # Based on the formula above \n",
    "    c = vector((q // 2) * mes for _ in range(A.ncols())) + u * A\n",
    "    c = c % q\n",
    "    return c\n",
    "\n",
    "def testLWE():\n",
    "    message = ZZ.random_element(2)\n",
    "    print(\"Original Message:\", message)\n",
    "\n",
    "    # Public-key encryption\n",
    "    c_pk = LWEenc(message, A, kappa, q)\n",
    "    # Public-key decryption\n",
    "    decrypted_pk = LWEdecrypt(c_pk, q, s)\n",
    "    print(\"Encrypted Message:\", c_pk)\n",
    "    print(\"Decrypted Message:\", decrypted_pk)\n",
    "    \n",
    "testLWE()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RLWE encryption\n",
    "\n",
    "For RLWE encryption, we replace ${\\mathbb Z}_q$ by \n",
    "the polynomial ring $R_q={\\mathbb Z}_q[x] / <x^k+1>$, where $k$\n",
    "is a power of $2$. \n",
    "\n",
    " Addition and multiplication of polynomials are performed modulo\n",
    "  $x^k+1$ and prime $q$.\n",
    " We can take $m \\in R_2={\\mathbb Z}_2[x] / <\\!x^k+1\\!>$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "R:Univariate Polynomial Ring in x over Integer Ring,k:8,p:x^8 + 1,q:67,Rq:Univariate Polynomial Ring in x over Finite Field of size 67\n"
     ]
    }
   ],
   "source": [
    "def genRings(k=8,nq=6):\n",
    "  R=ZZ['x']\n",
    "  p=R.gen()^k+1\n",
    "  q=next_prime(2^nq)\n",
    "  Rq=R.change_ring(GF(q))\n",
    "  return R,k,p,q,Rq\n",
    "\n",
    "R,k,p,q,Rq=genRings()\n",
    "print(f\"R:{R},k:{k},p:{p},q:{q},Rq:{Rq}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The public-key is $(a,t)$ with $t=a \\cdot s+e$\n",
    "for random $a$ in $R_q$ and small $s$, $e$ in $R$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def binomialRLWE(R,k,kappa):\n",
    "  return R([binom(kappa) for i in range(k)])\n",
    "\n",
    "def genKeyRLWE(kappa=2,k=8,nq=12):\n",
    "  R,k,p,q,Rq=genRings(k,nq)\n",
    "  s=binomialRLWE(R,k,kappa)\n",
    "  a=Rq([ZZ.random_element(q) for i in range(k)])\n",
    "  e=binomialRLWE(R,k,kappa)\n",
    "  t=(a*s+e) % p\n",
    "  return q,Rq,R,p,kappa,k,a,t,s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "q:4099,Rq:Univariate Polynomial Ring in x over Finite Field of size 4099,R:Univariate Polynomial Ring in x over Integer Ring,p:x^8 + 1,kappa:2,k:8,a:2621*x^7 + 1986*x^6 + 1432*x^5 + 21*x^4 + 2786*x^3 + 24*x^2 + 3455*x + 772,t:2197*x^7 + 1610*x^6 + 2840*x^5 + 273*x^4 + 1602*x^3 + 1491*x^2 + 2077*x + 564,s:x^5 + x^4 + x^3 + x - 1\n"
     ]
    }
   ],
   "source": [
    "q,Rq,R,p,kappa,k,a,t,s=genKeyRLWE()\n",
    "print(f\"q:{q},Rq:{Rq},R:{R},p:{p},kappa:{kappa},k:{k},a:{a},t:{t},s:{s}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To encrypt\n",
    "a message $m \\in R_2$, compute\n",
    "$ c=(a \\cdot r+e_1,~t \\cdot r+e_2 + \\lfloor q/2 \\rceil m)$,\n",
    "for small $e_1$, $e_2$ and $r$ in $R$. \n",
    "\n",
    "To decrypt,\n",
    "compute $m={\\sf th}(v- s \\cdot u)$.\n",
    "\\begin{align*}\n",
    " v- { s} \\cdot u & =\n",
    "     { t} \\cdot r+e_2 + \\lfloor q/2 \\rceil m\n",
    "     - { s} \\cdot ( { a} \\cdot\n",
    "     r+e_1) \\\\\n",
    "     & = ({ t} -{ a}\n",
    "     \\cdot { s} ) \\cdot r+e_2 + \\lfloor q/2 \\rceil m\n",
    "     - { s} \\cdot e_1 \\\\\n",
    "     & =\\lfloor q/2 \\rceil m +\n",
    "\\underbrace{\n",
    "     { e} \\cdot  r + e_2\n",
    "     - { s} \\cdot e_1}_{\\text{small}}\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original Message: 1507*x^7 + 3203*x^6 + 2245*x^5 + 1331*x^4 + 1837*x^3 + 3119*x^2 + 2822*x + 1262\n",
      "Encrypted Message: (1320*x^7 + 3567*x^6 + 2389*x^5 + 2323*x^4 + 2849*x^3 + 2798*x^2 + 2898*x + 1552, 3312*x^7 + 3714*x^6 + 4010*x^5 + 942*x^4 + 1851*x^3 + 1510*x^2 + 1257*x + 1230)\n",
      "Decrypted Message: x^7 + x^5 + x^4 + x^3 + x + 1\n"
     ]
    }
   ],
   "source": [
    "def encryptRLWE(m,R,q,p,kappa,k,a,t):\n",
    "    \n",
    "    # Adds random noise\n",
    "    e1 = binomialRLWE(R, k, kappa) \n",
    "    e2 = binomialRLWE(R, k, kappa) \n",
    "    # Small random polynomial\n",
    "    r = binomialRLWE(R, k, kappa) \n",
    "\n",
    "    # mod p (i.e., %p) ensures the result is within the polynomial ring\n",
    "    c0 = (a*r+e1)%p\n",
    "    c1 = (t*r+e2+((q // 2) * m))%p #(q // 2) adds noise\n",
    "    \n",
    "    return c0,c1\n",
    "\n",
    "\n",
    "def decryptRLWE(c,R,q,p,s):\n",
    "    c0, c1 = c \n",
    "    u = (s*c0)%p\n",
    "    v = ((c1-u)%p)//(q // 2) #//(q // 2) removes the noise\n",
    "    \n",
    "    # Apply th() to each term in v\n",
    "    m = v.map_coefficients(lambda x: th(x.lift(), q))\n",
    "    m = Rq(m) #makes sure we're working with a polynomial ring\n",
    "    \n",
    "    return m\n",
    "\n",
    "\n",
    "def testRLWE():\n",
    "    message = Rq([ZZ.random_element(q) for i in range(k)])\n",
    "    print(f\"Original Message: {message}\")\n",
    "\n",
    "    encrypted_message = encryptRLWE(message,R,q,p,kappa,k,a,t)\n",
    "    print(f\"Encrypted Message: {encrypted_message}\")\n",
    "    \n",
    "    decrypted_message = decryptRLWE(encrypted_message, R, q, p, s)\n",
    "    print(f\"Decrypted Message: {decrypted_message}\")\n",
    "\n",
    "\n",
    "testRLWE()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The error $e$ must be kept secret, otherwise we can recover the secret-key $s$ using\n",
    "$$ s=(t-e)/a \\pmod{x^k+1,q}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s = x^5 - x^4 + x^3 + x\n",
      "Recovered s = x^5 + 4098*x^4 + x^3 + x\n"
     ]
    }
   ],
   "source": [
    "def from_e():\n",
    "    s=binomialRLWE(R,k,kappa)\n",
    "    print(\"s = %s\" % str(s))\n",
    "    a=Rq([ZZ.random_element(q) for i in range(k)])\n",
    "    e=binomialRLWE(R,k,kappa)\n",
    "    t=(a*s+e) % p\n",
    "    \n",
    "    # show how to recover s from the knowledge of a, t and e.\n",
    "    s_recover = ((t - e) * a.inverse_mod(p)) % p\n",
    "    s_recover = Rq(s_recover)\n",
    "    \n",
    "    print(\"Recovered s = %s\" % str(s_recover))\n",
    "\n",
    "\n",
    "from_e()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "SageMath 10.1",
   "language": "sage",
   "name": "sagemath"
  },
  "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.11.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}