NuCypher_sumbased.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9e6a4a6f",
   "metadata": {},
   "source": [
    "# NuCyper\n",
    "\n",
    "notes\n",
    "- KEM : key encapsulation mechanism\n",
    "- DEM : data encapsulation mechanism\n",
    "- AEAD: authentication encryption with associated data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfb17e03",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "3492a49b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x = 5931, inv_num = 6959, x*inv_num = 1\n",
      "polynomial = a =[5, 6, 45, 45, 3, 2, 3]\n",
      "evaluate_polynomial(a,0) == a[0] = True\n",
      "evaluate_polynomial(a,1) == sum(a)%q = True\n",
      "opp_mod(6,11) = 5\n",
      "add_mod(opp_mod(6,11),6,11) == 0? True\n",
      "test_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]\n",
      "test_bytes = b'\\x01\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x03\\x00\\x00\\x00\\x04\\x00\\x00\\x00\\x05\\x00\\x00\\x00\\x06\\x00\\x00\\x00\\x07\\x00\\x00\\x00\\x08\\x00\\x00\\x00\\t\\x00\\x00\\x00\\n\\x00\\x00\\x00'\n",
      "bytes_to_list(test_bytes) = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]\n"
     ]
    }
   ],
   "source": [
    "# utils \n",
    "\n",
    "import hashlib\n",
    "import struct # for packing and unpacking\n",
    "import random # for random numbers, not cryptographically secure\n",
    "from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305 # for AEAD\n",
    "import os # for os.urandom for nonce\n",
    "\n",
    "q = 7919 # group of prime order, 1000 primese here https://primes.utm.edu/lists/small/1000.txt\n",
    "g = 1 # generator in Gq\n",
    "U = 1 # generator in Gq\n",
    "\n",
    "def inv_mod(x,mod=q):\n",
    "    if q == 0:\n",
    "        return 1/x\n",
    "    else:\n",
    "        # return x^-1 mod q\n",
    "        return pow(x,mod-2,q)\n",
    "\n",
    "def add_mod(x,y,mod=q):\n",
    "    return (x+y) % mod\n",
    "\n",
    "def sub_mod(x,y,mod=q):\n",
    "    return (x-y) % mod\n",
    "\n",
    "def opp_mod(x,mod=q):\n",
    "    return (-x) % mod\n",
    "\n",
    "def mul_mod(x,y,mod=q):\n",
    "    return (x*y) % mod\n",
    "\n",
    "def div_mod(x,y,mod=q):\n",
    "    return (x*inv_mod(y,mod)) % mod\n",
    "\n",
    "def evaluate_polynomial(coeff,val,mod=q):\n",
    "    s = 0\n",
    "    for i in range(len(coeff)):\n",
    "        s += (coeff[i]*pow(val,i,mod)) %mod\n",
    "        s=s%mod\n",
    "    \n",
    "    return s\n",
    "\n",
    "def list_to_bytes(l):\n",
    "    return b''.join([struct.pack('i', i) for i in l])\n",
    "\n",
    "def bytes_to_list(b):\n",
    "    return [struct.unpack('i', b[i:i+4])[0] for i in range(0, len(b), 4)]\n",
    "\n",
    "def lambda_compute(s_vec,idx,mod=q):\n",
    "    # s_vec is a list of secret shares\n",
    "    # returns lambda\n",
    "    l = 1\n",
    "    for j, s_j in enumerate(s_vec):\n",
    "        if j != idx:\n",
    "            inv_num = inv_mod(s_j-s_vec[idx],mod)\n",
    "            s = s_j\n",
    "            l = (l * s) * inv_num \n",
    "            #print(f\"i = {i}, inv_num = {inv_num}, l = {l}\")\n",
    "    return l%q\n",
    "\n",
    "\n",
    "# test inv_mod\n",
    "x = random.randint(1,q-1)\n",
    "inv_num = inv_mod(x,q)\n",
    "\n",
    "print(f\"x = {x}, inv_num = {inv_num}, x*inv_num = {(x*inv_num) % q}\")\n",
    "\n",
    "# test evaluate_polynomial\n",
    "a = [5,6,45,45,3,2,3]\n",
    "print(f\"polynomial = a ={a}\")\n",
    "print(f\"evaluate_polynomial(a,0) == a[0] = {evaluate_polynomial(a,0) == a[0]}\")\n",
    "print(f\"evaluate_polynomial(a,1) == sum(a)%q = {evaluate_polynomial(a,1) == sum(a)%q}\")\n",
    "\n",
    "x = 6\n",
    "mod = 11\n",
    "# test opp_mod\n",
    "print(f\"opp_mod({x},{mod}) = {opp_mod(x,mod)}\")\n",
    "print(f\"add_mod(opp_mod({x},{mod}),{x},{mod}) == 0? {add_mod(opp_mod(x,mod),x,mod) == 0}\")\n",
    "\n",
    "# test list_to_bytes and bytes_to_list\n",
    "test_list = [1,2,3,4,5,6,7,8,9,10]\n",
    "test_bytes = list_to_bytes(test_list)\n",
    "print(f\"test_list = {test_list}\")\n",
    "print(f\"test_bytes = {test_bytes}\")\n",
    "print(f\"bytes_to_list(test_bytes) = {bytes_to_list(test_bytes)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "5ef34ebf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bytes_data = b'hello world'\n",
      "hash_to_q(bytes_data) = 5576\n",
      "\n",
      "a = 1, b = 2, c = 3, d = 4\n",
      "h2(a,b) = 6762\n",
      "h3(a,b,c) = 3127\n",
      "h4(a,b,c,d) = 3110\n",
      "\n",
      "b'\\x8e\\xe7\\xe0C\\xad&\\xa4+\\x99\\xd8\\xf9S\\x172\\x9d\\xc1\\xb3\\xacv\\xa4>K\\xf46\\x81\\xec\\x0fZL\\x1av\\x15'\n"
     ]
    }
   ],
   "source": [
    "def hash_to_q(bytes_data,q=q):\n",
    "    hexdig = hashlib.sha256(bytes_data).hexdigest()\n",
    "    \n",
    "    # convert to int mod q\n",
    "    return int(hexdig,16) % q\n",
    "\n",
    "def H2(num1,num2,q=q):\n",
    "    num1_bytes = num1.to_bytes(32,byteorder='big')\n",
    "    num2_bytes = num2.to_bytes(32,byteorder='big')\n",
    "    return hash_to_q(num1_bytes+num2_bytes,q)\n",
    "\n",
    "def H3(num1,num2,num3,q=q):\n",
    "    num1_bytes = num1.to_bytes(32,byteorder='big')\n",
    "    num2_bytes = num2.to_bytes(32,byteorder='big')\n",
    "    num3_bytes = num3.to_bytes(32,byteorder='big')\n",
    "    return hash_to_q(num1_bytes+num2_bytes+num3_bytes,q)\n",
    "\n",
    "def H4(num1,num2,num3,num4,q=q):\n",
    "    num4 = num4 % q\n",
    "    num1_bytes = num1.to_bytes(32,byteorder='big')\n",
    "    num2_bytes = num2.to_bytes(32,byteorder='big')\n",
    "    num3_bytes = num3.to_bytes(32,byteorder='big')\n",
    "    num4_bytes = num4.to_bytes(32,byteorder='big')\n",
    "    return hash_to_q(num1_bytes+num2_bytes+num3_bytes+num4_bytes,q)\n",
    "\n",
    "def H6(num1,num2,num3,num4,num5,num6,q=q):\n",
    "    num4 = num4 % q\n",
    "    num5 = num5 % q\n",
    "    num6 = num6 % q\n",
    "    num1_bytes = num1.to_bytes(32,byteorder='big')\n",
    "    num2_bytes = num2.to_bytes(32,byteorder='big')\n",
    "    num3_bytes = num3.to_bytes(32,byteorder='big')\n",
    "    num4_bytes = num4.to_bytes(32,byteorder='big')\n",
    "    num5_bytes = num5.to_bytes(32,byteorder='big')\n",
    "    num6_bytes = num6.to_bytes(32,byteorder='big')\n",
    "    return hash_to_q(num1_bytes+num2_bytes+num3_bytes+num4_bytes+num5_bytes+num6_bytes,q)   \n",
    "\n",
    "def KDF(num,q=q):\n",
    "    num_bytes = num.to_bytes(32,byteorder='big')\n",
    "    salt = b'NuCypher'\n",
    "    # KDF returns bytes as per the spec\n",
    "    return hashlib.sha256(num_bytes+salt).digest()\n",
    "\n",
    "# test hash_to_q\n",
    "bytes_data = b'hello world'\n",
    "h = hash_to_q(bytes_data)\n",
    "print(f\"bytes_data = {bytes_data}\\nhash_to_q(bytes_data) = {h}\\n\")\n",
    "\n",
    "# make 4 random numbers and hash them\n",
    "a = 1\n",
    "b = 2\n",
    "c = 3\n",
    "d = 4\n",
    "\n",
    "print(f\"a = {a}, b = {b}, c = {c}, d = {d}\")\n",
    "\n",
    "h2 = H2(a,b)\n",
    "h3 = H3(a,b,c)\n",
    "h4 = H4(a,b,c,d)\n",
    "\n",
    "print(f\"h2(a,b) = {h2}\\nh3(a,b,c) = {h3}\\nh4(a,b,c,d) = {h4}\\n\")\n",
    "\n",
    "# KDF\n",
    "kdf = KDF(h)\n",
    "print(kdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7bd8db5",
   "metadata": {},
   "source": [
    "## Key Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "49cb34f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Alice's sk: 2341, pk: 2341\n",
      "Bob's sk: 7454, pk: 7454\n",
      "d:  7050\n",
      "D:  4697\n",
      "i 0 sx:  2859\n",
      "i 1 sx:  339\n",
      "i 2 sx:  6498\n",
      "i 3 sx:  1914\n",
      "i 4 sx:  3323\n",
      "i 5 sx:  6732\n",
      "i 6 sx:  7898\n",
      "i 7 sx:  1729\n",
      "i 8 sx:  6588\n",
      "i 9 sx:  6499\n",
      "[[6509, 7221, 4829, 7221, 4740, 3185], [7026, 5969, 4829, 5969, 3593, 2457], [241, 2916, 4829, 2916, 5090, 1547], [1089, 3926, 4829, 3926, 7335, 2236], [5452, 2262, 4829, 2262, 4221, 7843], [1399, 2205, 4829, 2205, 398, 3997], [6850, 5104, 4829, 5104, 7163, 4899], [726, 6947, 4829, 6947, 5158, 4743], [7353, 2265, 4829, 2265, 423, 4891], [7562, 3971, 4829, 3971, 5035, 6921]]\n"
     ]
    }
   ],
   "source": [
    "def key_gen(mod=q,g=g):\n",
    "    sk = random.randint(1,q-1)\n",
    "    pk = mul_mod(sk, g, mod)\n",
    "    return sk,pk\n",
    "\n",
    "def re_key_gen(sk, pk, N, t,mod=q,g=g):\n",
    "    # N is the number of re-encryption keys\n",
    "    # t is the threshold\n",
    "    assert t <= N\n",
    "    assert sk < mod\n",
    "    assert pk < mod\n",
    "    \n",
    "    x = random.randint(1, q-1)\n",
    "    X = mul_mod(x, g,mod)\n",
    "    d = H3(X,pk, mul_mod(x, pk,mod))\n",
    "    print(\"d: \",d)\n",
    "\n",
    "    # generate t-1 random numbers\n",
    "    f = [random.randint(1,mod-1) for i in range(t-1)]\n",
    "\n",
    "    # f0=add_mod(sk, opp_mod(d,mod),mod)\n",
    "    # print(\"f0: \",f0)\n",
    "    f0=mul_mod(sk,inv_mod(d,mod), mod)\n",
    "    poly_coeffs = [f0] + f\n",
    "    pkA = mul_mod(sk, g,mod)\n",
    "    D = H3(pkA,pk, mul_mod(pk,sk,mod))\n",
    "    print(\"D: \",D)\n",
    "\n",
    "    KF = []\n",
    "\n",
    "    for i in range(N):\n",
    "        y = random.randint(1,mod-1)\n",
    "        id = random.randint(1,mod-1)\n",
    "        sx = H2(id, D)\n",
    "        print(\"i\", i, \"sx: \",sx)\n",
    "        Y = mul_mod(g,y,mod)\n",
    "        rk = evaluate_polynomial(poly_coeffs,sx,mod)\n",
    "        U1 = mul_mod(U,rk,q)\n",
    "        z1 = H6(Y,id,pkA,pk,U1,X)    \n",
    "        z2 = (y - z1 * sk) % q\n",
    "\n",
    "        kFrag = [id,rk, X, U1, z1, z2]\n",
    "        KF.append(kFrag)\n",
    "    \n",
    "    return KF\n",
    "\n",
    "\n",
    "alice_sk, alice_pk = key_gen()\n",
    "bob_sk, bob_pk = key_gen()\n",
    "\n",
    "print(f\"Alice's sk: {alice_sk}, pk: {alice_pk}\")\n",
    "print(f\"Bob's sk: {bob_sk}, pk: {bob_pk}\")\n",
    "\n",
    "# re-encryption key generation\n",
    "N = 10\n",
    "t = 6\n",
    "KF = re_key_gen(alice_sk, bob_pk, N, t)\n",
    "\n",
    "print(KF)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ddfadf77",
   "metadata": {},
   "source": [
    "## Encapsulation and Decapsulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "d41d6649",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iinKDF: 1539\n",
      "inKDF: 1539\n",
      "is K == K2? True\n"
     ]
    }
   ],
   "source": [
    "def encapsulate(alice_pk,mod=q,g=g):\n",
    "    r = random.randint(1,mod-1)\n",
    "    u = random.randint(1,mod-1)\n",
    "\n",
    "    E = mul_mod(g,r,mod)\n",
    "    V = mul_mod(g,u,mod)\n",
    "\n",
    "    h2 = H2(E,V)\n",
    "    s = add_mod(u, mul_mod(r, h2,mod), mod) \n",
    "    iinKDF = mul_mod(alice_pk,add_mod(r,u,mod),mod)\n",
    "    print(f\"iinKDF: {iinKDF}\")\n",
    "    K = KDF(iinKDF)\n",
    "    capsule = [E,V,s]\n",
    "\n",
    "    return K, capsule\n",
    "\n",
    "def check_capsule(capsule,mod=q,g=g):\n",
    "    E,V,s = capsule\n",
    "  \n",
    "    h2 = H2(E,V)\n",
    " \n",
    "    if mul_mod(g,s,mod) == add_mod(V, mul_mod(E,h2,mod),mod):\n",
    "        return True\n",
    "    else:\n",
    "        print(f\"mul_mod(g,s,mod) {mul_mod(g,s,mod)}\")\n",
    "        print(f\"add_mod(V, mul_mod(E,h2,mod),mod: {add_mod(V, mul_mod(E,h2,mod),mod)}\")\n",
    "        return False\n",
    "\n",
    "def decapsulate(alice_sk, capsule,mod=q):\n",
    "    E,V,s = capsule\n",
    "    if not check_capsule(capsule):\n",
    "        raise ValueError(\"Invalid capsule\")\n",
    "    inKDF = mul_mod(E+V,alice_sk,mod)\n",
    "    print(f\"inKDF: {inKDF}\")\n",
    "    K = KDF(inKDF)\n",
    "    return K\n",
    "\n",
    "\n",
    "K, capsule = encapsulate(alice_pk)\n",
    "K2 = decapsulate(alice_sk, capsule)\n",
    "print(f\"is K == K2? {K == K2}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2b072f87",
   "metadata": {},
   "source": [
    "## Re-Encapsulation and Fragments Decapsulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "c7f02877",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "selected indices: [4, 8, 5, 0, 7, 6]\n",
      "cFrag 4: [4210, 6424, 5452, 4829]\n",
      "cFrag 8: [1905, 4374, 7353, 4829]\n",
      "cFrag 5: [491, 5779, 1399, 4829]\n",
      "cFrag 0: [3084, 7106, 6509, 4829]\n",
      "cFrag 7: [2434, 6923, 726, 4829]\n",
      "cFrag 6: [3611, 1906, 6850, 4829]\n",
      "\n",
      "D: 4697\n",
      "S: [3323, 6588, 6732, 2859, 1729, 7898]\n",
      "d: 7050\n",
      "E_prime: 4105\n",
      "V_prime: 7430\n",
      "inKDF: 1539\n",
      "is K == K3? **True**\n"
     ]
    }
   ],
   "source": [
    "def re_encapsulate(kFrag, capsule,mod=q):\n",
    "    if not check_capsule(capsule):\n",
    "        raise ValueError(\"Invalid capsule\")\n",
    "    \n",
    "    E,V,s = capsule\n",
    "    id,rk,X,U1,z1,z2 = kFrag\n",
    "\n",
    "    E1 = mul_mod(E,rk,mod)\n",
    "    V1 = mul_mod(V,rk,mod)\n",
    "    cFrag = [E1,V1,id,X]\n",
    "\n",
    "    return cFrag\n",
    "\n",
    "def decapsulateFrags(skB,pkA,cFrag_vec,mod=q,g=g):\n",
    "    \n",
    "    t = len(cFrag_vec)\n",
    "    pkB = mul_mod(g,skB,mod)\n",
    "    assert pkB == bob_pk\n",
    "    D = H3(pkA,pkB, mul_mod(pkA,skB,mod))\n",
    "    print(f\"D: {D}\" )\n",
    "    S = [H2(cFrag_vec[i][2],D) for i in range(t)] # H2(id_i,D), id is the third element of cFrag\n",
    "    print(f\"S: {S}\")\n",
    "\n",
    "    E_prime = 0\n",
    "    V_prime = 0\n",
    "\n",
    "    for i in range(t):\n",
    "        lam = lambda_compute(S,i)\n",
    "        E = cFrag_vec[i][0] # E is the first element of cFrag\n",
    "        V = cFrag_vec[i][1] # V is the second element of cFrag\n",
    "        E_prime = add_mod(E_prime , mul_mod(E,lam, mod),mod)\n",
    "        V_prime = add_mod(V_prime,mul_mod(V,lam, mod), mod)\n",
    "    X = cFrag_vec[0][3]\n",
    "    d = H3(X,pkB, mul_mod(X,skB,mod))\n",
    "    inKDF = mul_mod(add_mod(E_prime,V_prime,mod),d,mod)\n",
    "    print(f\"d: {d}\")\n",
    "    print(f\"E_prime: {E_prime}\")\n",
    "    print(f\"V_prime: {V_prime}\")\n",
    "    print(f\"inKDF: {inKDF}\")\n",
    "    K = KDF(inKDF)\n",
    "\n",
    "    return K\n",
    "\n",
    "tprime =t\n",
    "L = random.sample(range(N), tprime)\n",
    "# L = [0,1,2]\n",
    "# L=sorted([7, 8, 6, 4, 3, 5, 2])\n",
    "print(f\"selected indices: {L}\")\n",
    "cFrag_vec = [re_encapsulate(KF[i], capsule) for i in L]\n",
    "\n",
    "for i in range(tprime):\n",
    "    print(f\"cFrag {L[i]}: {cFrag_vec[i]}\")\n",
    "\n",
    "print()\n",
    "\n",
    "K3 = decapsulateFrags(bob_sk, alice_pk, cFrag_vec)\n",
    "print(f\"is K == K3? **{K == K3}**\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "640cc00b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test_aead_key: b'\\x0f\\xe6\\x16\\xab\\xf6\\tR\\xda\\xa7{JC\"E\\x0f\\x00w\\x07\\xf6\\x124G\\x1fh\\xe2\\xa8\\xb2\\x0f\\x16\\x87\\xfdo'\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "b'a secret message'"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## testing the ChaCha20Poly1305 AEAD\n",
    "test_aead_data = b\"a secret message\"\n",
    "test_aead_aad = b\"authenticated but unencrypted data\"\n",
    "test_aead_key = ChaCha20Poly1305.generate_key()\n",
    "\n",
    "print(f\"test_aead_key: {test_aead_key}\")\n",
    "test_aead_chacha = ChaCha20Poly1305(test_aead_key)\n",
    "test_aead_nonce = os.urandom(12)\n",
    "test_aead_ct = test_aead_chacha.encrypt(test_aead_nonce, test_aead_data, test_aead_aad)\n",
    "test_aead_chacha.decrypt(test_aead_nonce, test_aead_ct, test_aead_aad)\n",
    "# b'a secret message'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6e54130d",
   "metadata": {},
   "source": [
    "## KEM/DEM construction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "38ee0bca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iinKDF: 5675\n",
      "enc_data: b'f(\\x1d\\x8a\\x8c&\\x9f\\xa5\\x9a\\xca?\\xea\\xf6 \\x16h#\\x7f\\xfc\\x81Lw\\xbc\\xea\\xe4\\xb8\\xd1\\x08\\x8a?v\\xbc\\xfc\\x1b\\x19I\\xfa\\xb6'\n",
      "capsule: [5299, 471, 3599]\n",
      "nonce: b'\\x1e\\xae,\\xc1\\xb8m{\\xd9f\\x9eR\\xc7'\n",
      "inKDF: 5675\n",
      "decrypted_message: b'this is a test message'\n",
      "test_message == decrypted_message? True\n"
     ]
    }
   ],
   "source": [
    "def encrypt(pkA,M):\n",
    "    K, capsule = encapsulate(pkA)\n",
    "    bytes_capsule = list_to_bytes(capsule)\n",
    "    nonce = os.urandom(12)\n",
    "    enc_data = ChaCha20Poly1305(K).encrypt(nonce, M, bytes_capsule)\n",
    "    return (capsule, enc_data, nonce)\n",
    "\n",
    "def decrypt(skA,cyphertext):\n",
    "    capsule, enc_data, nonce= cyphertext\n",
    "    bytes_capsule = list_to_bytes(capsule)\n",
    "    K = decapsulate(skA, capsule)\n",
    "    return ChaCha20Poly1305(K).decrypt(nonce, enc_data, bytes_capsule)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "test_message = b\"this is a test message\"\n",
    "capsule, enc_data, nonce = encrypt(alice_pk, test_message)\n",
    "print(f\"enc_data: {enc_data}\")\n",
    "print(f\"capsule: {capsule}\")\n",
    "print(f\"nonce: {nonce}\")\n",
    "\n",
    "decrypted_message = decrypt(alice_sk, (capsule, enc_data,nonce ))\n",
    "print(f\"decrypted_message: {decrypted_message}\")\n",
    "print(f\"test_message == decrypted_message? {test_message == decrypted_message}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "c54cbfc0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cFrag_vec: [[7390, 3840, 6509, 4829], [1245, 154, 7026, 4829], [1915, 3449, 241, 4829], [661, 4019, 1089, 4829], [4891, 4256, 5452, 4829], [3770, 1166, 1399, 4829]]\n",
      "D: 4697\n",
      "S: [2859, 339, 6498, 1914, 3323, 6732]\n",
      "d: 7050\n",
      "E_prime: 5117\n",
      "V_prime: 6313\n",
      "inKDF: 5675\n",
      "decrypted_message2: b'this is a test message'\n",
      "test_message == decrypted_message2? True\n"
     ]
    }
   ],
   "source": [
    "def re_encrypt(kFrag, cyphertext):\n",
    "    capsule, enc_data, nonce= cyphertext\n",
    "    cFrag = re_encapsulate(kFrag, capsule)\n",
    "    Cprime = (cFrag, enc_data, nonce)\n",
    "    return Cprime\n",
    "\n",
    "\n",
    "def decrypt_frags(skB,pkA,cyphertexts, capsule,mod=q,g=g):\n",
    "    cFrag_vec = [cyphertexts[i][0] for i in range(len(cyphertexts))]\n",
    "    print(f\"cFrag_vec: {cFrag_vec}\")\n",
    "    K = decapsulateFrags(skB,pkA,cFrag_vec,mod=q,g=g)\n",
    "    _, enc_data, nonce = cyphertexts[0]\n",
    "    bytes_capsule = list_to_bytes(capsule)\n",
    "    dec_data = ChaCha20Poly1305(K).decrypt(nonce, enc_data, bytes_capsule)\n",
    "    return dec_data\n",
    "\n",
    "cyphertexts = [re_encrypt(KF[i], (capsule, enc_data, nonce)) for i in range(tprime)]\n",
    "\n",
    "decrypted_message2 = decrypt_frags(bob_sk, alice_pk, cyphertexts, capsule)\n",
    "print(f\"decrypted_message2: {decrypted_message2}\")\n",
    "print(f\"test_message == decrypted_message2? {test_message == decrypted_message2}\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7 (main, Sep 15 2022, 01:51:29) [Clang 14.0.0 (clang-1400.0.29.102)]"
  },
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}