12 - Homomorphic Encryption with Paillier.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Homomorphic Encryption with Paillier\n",
    "\n",
    "This notebook provides an introduction to homomorphic encryption basics. The goal is to build a mental model that can help you notice use cases where homomorphic encyption would be a good fit. Again, please do not actually roll your own crypto! Instead, play around with the neat math on display here and build your understanding and intuition. :)\n",
    "\n",
    "Most of this notebook is code written by [Morten Dahl](https://github.com/mortendahl) in his [Private ML repo](https://github.com/mortendahl/privateml/) and [Paillier post](https://github.com/mortendahl/mortendahl.github.io/blob/master/_drafts/2019-04-15-paillier-encryption.md). Putting it together required several conversations where he helped me through how to understand the properties. Thank you Morten for the support in making this notebook and book possible. :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import random"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Required Primatives & Utility Functions\n",
    "\n",
    "In this section, I'll use Morten's work to define some building blocks we need for Paillier!\n",
    "\n",
    "\n",
    "A brief introduction to some of the methods:\n",
    "\n",
    "- [The Rabin-Miller primality test](https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test) is used to probablistically determine if a number is a prime.\n",
    "- There is a method to sample randomness, so that we can actually safely encrypt.\n",
    "- [The Extended Euclidean algorithm (egcd)](https://en.wikipedia.org/wiki/Extended_Euclidean_algorithm) that is used to calculate greatest common divisor, but also additional parameters which we will use to do encrypted math in a field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# see https://inventwithpython.com/rabinMiller.py\n",
    "\n",
    "n = 137\n",
    "SMALL_PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, \n",
    "                67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, \n",
    "                139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211, \n",
    "                223, 227, 229, 233, 239, 241, 251, 257, 263, 269, 271, 277, 281, 283, \n",
    "                293, 307, 311, 313, 317, 331, 337, 347, 349, 353, 359, 367, 373, 379, \n",
    "                383, 389, 397, 401, 409, 419, 421, 431, 433, 439, 443, 449, 457, 461, \n",
    "                463, 467, 479, 487, 491, 499, 503, 509, 521, 523, 541, 547, 557, 563, \n",
    "                569, 571, 577, 587, 593, 599, 601, 607, 613, 617, 619, 631, 641, 643, \n",
    "                647, 653, 659, 661, 673, 677, 683, 691, 701, 709, 719, 727, 733, 739, \n",
    "                743, 751, 757, 761, 769, 773, 787, 797, 809, 811, 821, 823, 827, 829, \n",
    "                839, 853, 857, 859, 863, 877, 881, 883, 887, 907, 911, 919, 929, \n",
    "                937, 941, 947, 953, 967, 971, 977, 983, 991, 997]\n",
    "\n",
    "def rewrite(num):\n",
    "    s = num - 1\n",
    "    t = 0\n",
    "    while s % 2 == 0:\n",
    "        s = s // 2\n",
    "        t += 1\n",
    "    return s, t\n",
    "\n",
    "def rabin_miller(num, iterations=10):\n",
    "    s, t = rewrite(num)\n",
    "    for _ in range(iterations):\n",
    "        a = random.randrange(2, num - 1)\n",
    "        v = pow(a, s, num)\n",
    "        if v != 1:\n",
    "            i = 0\n",
    "            while v != (num - 1):\n",
    "                if i == t - 1:\n",
    "                    return False\n",
    "                else:\n",
    "                    i = i + 1\n",
    "                    v = pow(v, 2, num)\n",
    "    return True\n",
    "\n",
    "def is_prime(num):\n",
    "    if (num < 2): return False\n",
    "    for prime in SMALL_PRIMES:\n",
    "        if num == prime: return True\n",
    "        if num % prime == 0: return False\n",
    "    return rabin_miller(num)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_randomness(ek):\n",
    "    while True:\n",
    "        r = random.randrange(ek.n)\n",
    "        if math.gcd(r, ek.n) == 1:\n",
    "            return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_prime(bitsize):\n",
    "    lower = 1 << (bitsize-1)\n",
    "    upper = 1 << (bitsize)\n",
    "    while True:\n",
    "        candidate = random.randrange(lower, upper)\n",
    "        if is_prime(candidate):\n",
    "            return candidate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from http://www.ucl.ac.uk/~ucahcjm/combopt/ext_gcd_python_programs.pdf\n",
    "def egcd_binary(a, b):\n",
    "    u, v, s, t, r = 1, 0, 0, 1, 0\n",
    "    while (a % 2 == 0) and (b % 2 == 0):\n",
    "        a, b, r = a//2, b//2, r+1\n",
    "    alpha, beta = a, b\n",
    "    while (a % 2 == 0):\n",
    "        a = a//2\n",
    "        if (u % 2 == 0) and (v % 2 == 0):\n",
    "            u, v = u//2, v//2\n",
    "        else:\n",
    "            u, v = (u + beta)//2, (v - alpha)//2\n",
    "    while a != b:\n",
    "        if (b % 2 == 0):\n",
    "            b = b//2\n",
    "            if (s % 2 == 0) and (t % 2 == 0):\n",
    "                s, t = s//2, t//2\n",
    "            else:\n",
    "                s, t = (s + beta)//2, (t - alpha)//2\n",
    "        elif b < a:\n",
    "            a, b, u, v, s, t = b, a, s, t, u, v\n",
    "        else:\n",
    "            b, s, t = b - a, s - u, t - v\n",
    "    return (2 ** r) * a, s, t\n",
    "\n",
    "\n",
    "def inverse(a, field):\n",
    "    _, b, _ = egcd_binary(a, field)\n",
    "    return b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Looking at multiplicative field inverses\n",
    "\n",
    "Now that we have the helper functions, we can take a look at the field properties. As you learned in the book, the multiplicative inverse works in a field to support multiplication primatives. It is a multiplicative inverse if the it can be multiplied with the original number and then modulo the field to equal 1. Check out the list below to see this in action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n = 137\n",
      "\n",
      "x: 1\t inverse(x): 1  \t x * inverse(x) mod n: 1\n",
      "x: 2\t inverse(x): 69  \t x * inverse(x) mod n: 1\n",
      "x: 3\t inverse(x): 46  \t x * inverse(x) mod n: 1\n",
      "x: 4\t inverse(x): 103  \t x * inverse(x) mod n: 1\n",
      "x: 5\t inverse(x): 55  \t x * inverse(x) mod n: 1\n",
      "x: 6\t inverse(x): 23  \t x * inverse(x) mod n: 1\n",
      "x: 7\t inverse(x): -39  \t x * inverse(x) mod n: 1\n",
      "x: 8\t inverse(x): 120  \t x * inverse(x) mod n: 1\n",
      "x: 9\t inverse(x): 61  \t x * inverse(x) mod n: 1\n",
      "x: 10\t inverse(x): 96  \t x * inverse(x) mod n: 1\n",
      "x: 11\t inverse(x): 25  \t x * inverse(x) mod n: 1\n",
      "x: 12\t inverse(x): 80  \t x * inverse(x) mod n: 1\n",
      "x: 13\t inverse(x): 116  \t x * inverse(x) mod n: 1\n",
      "x: 14\t inverse(x): 49  \t x * inverse(x) mod n: 1\n",
      "x: 15\t inverse(x): 64  \t x * inverse(x) mod n: 1\n",
      "x: 16\t inverse(x): 60  \t x * inverse(x) mod n: 1\n",
      "x: 17\t inverse(x): 129  \t x * inverse(x) mod n: 1\n",
      "x: 18\t inverse(x): 99  \t x * inverse(x) mod n: 1\n",
      "x: 19\t inverse(x): 101  \t x * inverse(x) mod n: 1\n"
     ]
    }
   ],
   "source": [
    "print(\"n = {}\\n\".format(n))\n",
    "\n",
    "for x in range(1,20):\n",
    "    print(\"x: {}\\t inverse(x): {}  \\t x * inverse(x) mod n: {}\".format(\n",
    "        x, inverse(x, n), x * inverse(x, n) % n))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "17 * 129 % n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Paillier Encryption & Decryption Keys\n",
    "\n",
    "These helper functions (from Morten!) define the encryption and decryption keys for the Paillier cryptosystem. You will notice some of the calculations you already learned in Chapter 7 in these definitions!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncryptionKey:\n",
    "    def __init__(self, n):\n",
    "        self.n = n\n",
    "        self.nn = n * n\n",
    "        self.g = 1 + n\n",
    "        \n",
    "    def __repr__(self):\n",
    "        return \"Encryption Key <n: {}\\n nn: {}\\n g: {}>\".format(\n",
    "            self.n, self.nn, self.g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecryptionKey:\n",
    "    def __init__(self, p, q):\n",
    "        n = p * q\n",
    "\n",
    "        self.n = p * q\n",
    "        self.nn = n * n\n",
    "        self.g = 1 + n\n",
    "\n",
    "        order_of_n = (p - 1) * (q - 1)\n",
    "        self.d1 = order_of_n\n",
    "        self.d2 = inverse(order_of_n, n)\n",
    "        self.e = inverse(n, order_of_n)\n",
    "        \n",
    "    def __repr__(self):\n",
    "        return \"Decryption-Key <n: {}, d1: {}, e: {}>\".format(\n",
    "            self.n, self.d1, self.e)        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def keygen(n_bitlength=512): # should be 2048\n",
    "    p = sample_prime(n_bitlength // 2)\n",
    "    q = sample_prime(n_bitlength // 2)\n",
    "    n = p * q\n",
    "\n",
    "    return EncryptionKey(n), DecryptionKey(p, q)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Encrypting, Decrypting and Extracting the Randomness (r)\n",
    "\n",
    "These functions allow you to encrypt and decrypt using the above keys and extract the randomness if that is required. Usually this would only be done to prove the \"correctness\" of the computation, and is therefore not always required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def enc(ek: EncryptionKey, x, r):\n",
    "    gx = pow(ek.g, x, ek.nn)\n",
    "    rn = pow(r, ek.n, ek.nn)\n",
    "    c = (gx * rn) % ek.nn\n",
    "    return c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dec(dk: DecryptionKey, c):\n",
    "    gxd = pow(c, dk.d1, dk.nn)\n",
    "    xd = dlog(gxd, dk.n)\n",
    "    x = (xd * dk.d2) % dk.n\n",
    "    return x if x < dk.n/2 else x - dk.n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dlog(gy, n):\n",
    "    y = (gy - 1) // n\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract(dk: DecryptionKey, c):\n",
    "    x = dec(dk, c)\n",
    "    gx = pow(dk.g, x, ek.nn)\n",
    "    gx_inv = inverse(gx, ek.nn)\n",
    "    rn = (c * gx_inv) % ek.nn\n",
    "    r = pow(rn, dk.e, dk.n)\n",
    "    return r"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Encrypting and Decrypting with Paillier\n",
    "\n",
    "Now you're ready to actually encrypt and decrypt using the Paillier cryptosystem. You'll generate your encryption and decryption keys, sample randomness in order to encrypt your message into ciphertext and then decrypt the message using your decryption key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "ek, dk = keygen()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Encryption Key <n: 6308324349728227799169268047033184195881371779885026845869399967914822134165716999609022763633695481334038499768973875074267051433386261148736040420917503\n",
       " nn: 39794956101374068115766666914697141943453546314743347274423294609860951223379024314340048792754131263520085527191478582015693574913144306405236416760926946880647144903986950437870812486346552275961475778935783786588398962729721190663270676696972571841356825843768225905367214242321349345371302987784331755009\n",
       " g: 6308324349728227799169268047033184195881371779885026845869399967914822134165716999609022763633695481334038499768973875074267051433386261148736040420917504>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ek"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = sample_randomness(ek)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "msg = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "ciphertext = enc(ek, msg, r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "13318562040696409702758528812924345335818299164953617929650042923992327244313640064803569416193632088674806904695878999880761923933432455325763026904104798673804522548252179782189132064552724542421613876038350839917013343021961311303770831146608043027815871740619537991742484293856289271519982275283282986213"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ciphertext"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dec(dk, ciphertext) == msg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: you might want to extract the randomness provided as part of the encryption to prove the correctness of the decryption. For one implementation of how you might use this, take a look at [the tf-encrypted implementation of secure aggregation](https://medium.com/dropoutlabs/building-secure-aggregation-into-tensorflow-federated-4514fca40cc0) where it is used to prove correct decryption."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "extract(dk, ciphertext) == r"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Homomorphic Properties\n",
    "\n",
    "Okay, so you can encrypt and decrypt, cool! But really, the important part of the Paillier cryptosystem is that you can use it for homomorphic encryption. Let's see how you can do this via a few helper functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_cipher(ek, c1, c2):\n",
    "    c = (c1 * c2) % ek.nn\n",
    "    return c\n",
    "\n",
    "def add_plain(ek, c1, x2):\n",
    "    c2 = pow(ek.g, x2, ek.nn)\n",
    "    c = (c1 * c2) % ek.nn\n",
    "    return c\n",
    "\n",
    "def neg(ek, c):\n",
    "    return inverse(c, ek.nn) \n",
    "\n",
    "def sub_cipher(ek, c1, c2):\n",
    "    c = add_cipher(ek, c1, neg(ek, c2))\n",
    "    return c\n",
    "\n",
    "def sub_plain(ek, c1, x2):\n",
    "    c = add_plain(ek, c1, ek.n - x2)\n",
    "    return c\n",
    "\n",
    "def mul_plain(ek, c1, x2):\n",
    "    c = pow(c1, x2, ek.nn)\n",
    "    return c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "msg_one, msg_two = 45, 234"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "r1 = sample_randomness(ek)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "r2 = sample_randomness(ek)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "c1 = enc(ek, msg_one, r1)\n",
    "c2 = enc(ek, msg_two, r2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(8300495113123882981928995611828881848870500488811192567106333973552994675153967728301429493081586585645685831033361750025256343879840490091038837627647576831632164879332711852520618080888121301881134281374204062468850439810425971974733881003239259689285683978092152056294200674796153145116647799919448217892,\n",
       " 33452402939873505743893139364694959902023028400516601475505862327229878012365643816181605356897921228726624232821175484794497579216648870602743534131101510282069072588225503679084360452010980311459832832285567611991063143624052719024408987810400271746237133259246888193032702937161383783733229312870416499129)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c1, c2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_addition = add_cipher(ek, c1, c2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "35095335123754385146727887296882861810552271776352568124652569998196032488955963715413728062632770206255525021056222528340129894596214981800335348265350405359318940795022893285765858690297764984882458146852205354288793002905878139434817941414723317902866109956826897234935159144629779820845063512909288772804"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result_addition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dec(dk, result_addition) == msg_one + msg_two"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_subtraction = sub_cipher(ek, c1, c2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dec(dk, result_subtraction) == msg_one - msg_two"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wow! Pretty neat, eh? :D\n",
    "\n",
    "### Plaintext methods\n",
    "\n",
    "You can also perform plaintext operations, like multiplying, adding and subtracting publicly known values as part of the computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dec(dk, mul_plain(ek, c1, 5)) == 5 * msg_one"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dec(dk, sub_plain(ek, c2, 1000)) == msg_two - 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dec(dk, add_plain(ek, c1, 4893849327)) == msg_one + 4893849327"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}