Diffie--Hellman and Elgamal.ipynb

diffie-hellman-and-elgamal.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/Pratyush/60480cb567cd3ebc87bbefd0476f9eff/diffie-hellman-and-elgamal.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u96I3gJ5h-sJ"
      },
      "source": [
        "## Setup and introduction to groups\n",
        "\n",
        "We'll begin by installing `ark_algebra_py`, and will then go over the group APIs.\n",
        "$\\newcommand{\\KeyGen}{\\mathsf{KeyGen}}$\n",
        "$\\newcommand{\\Encrypt}{\\mathsf{Enc}}$\n",
        "$\\newcommand{\\Decrypt}{\\mathsf{Dec}}$\n",
        "$\\newcommand{\\pk}{\\mathsf{pk}}$\n",
        "$\\newcommand{\\sk}{\\mathsf{sk}}$\n",
        "$\\newcommand{\\Fp}{\\mathbb{Z}_p}$\n",
        "$\\newcommand{\\Group}{\\mathbb{G}}$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ASq6EaEmh48a"
      },
      "outputs": [],
      "source": [
        "!pip install --upgrade ark_algebra_py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dhqsf2tFiTxu"
      },
      "source": [
        "### Group APIs\n",
        "\n",
        "We'll use *additive* notation for our group operations. So, instead of writing $g^x \\cdot g^y$, we'll write $x \\cdot G + y \\cdot G$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fo5N4b1EiX9j"
      },
      "outputs": [],
      "source": [
        "from ark_algebra_py.ark_algebra_py import G1 as G, Scalar\n",
        "\n",
        "G = G();\n",
        "# can generate random scalars/numbers via `.rand()`.\n",
        "# This is *not* secure; we only use it for this example\n",
        "a = Scalar.rand();\n",
        "b = Scalar.rand();\n",
        "print(a)\n",
        "print(b)\n",
        "\n",
        "# We can multiply the generator by scalar field elements\n",
        "# (This is the same as exponentiation)\n",
        "aG = a * G\n",
        "bG = b * G;\n",
        "\n",
        "# (don't expect any of the print outputs to make sense; we are working with elliptic curve points)\n",
        "print(\"a + b = \", aG + bG); # we can add...\n",
        "print(\"a - b = \", aG - G); # ... subtract ...\n",
        "print(\"a + -a = \", aG + -aG) # ... negate ...\n",
        "print(\"a.double() = \", aG.double()) # ... and double points.\n",
        "\n",
        "assert(aG.double() == aG + aG)\n",
        "assert(aG + -aG == G.identity())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pmDZan8QnJkx"
      },
      "source": [
        "## Diffie--Hellman Key Exchange\n",
        "\n",
        "Recall that the Diffie--Hellman key exchange looks like the following:\n",
        "\n",
        "1. Alice samples $a \\gets \\Fp$, and sends $A = a \\cdot G$.\n",
        "2. Bob samples $b \\gets \\Fp$, and sends $B = b \\cdot G$.\n",
        "3. Alice computes $K := a \\cdot B = ab \\cdot G$.\n",
        "4. Bob computes $K := b \\cdot A = ab \\cdot G$.\n",
        "\n",
        "Let's implement Alice's and Bob's algorithms."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xB1FJZwRn1Hw"
      },
      "outputs": [],
      "source": [
        "def alice_1():\n",
        "  # your code here;\n",
        "  return (a, A)\n",
        "\n",
        "def bob_1():\n",
        "  # your code here;\n",
        "  return (b, B)\n",
        "\n",
        "def alice_2(a, B):\n",
        "  # your code here\n",
        "  return K\n",
        "\n",
        "def bob_2(b, A):\n",
        "  # your code here\n",
        "  return K\n",
        "\n",
        "(a, A) = alice_1()\n",
        "(b, B) = bob_1()\n",
        "\n",
        "k_alice = alice_2(a, B)\n",
        "k_bob = bob_2(b, A)\n",
        "\n",
        "print(k_alice)\n",
        "print(k_bob)\n",
        "assert(k_alice == k_bob)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xtqic8vaoTzh"
      },
      "source": [
        "## Elgamal Encryption\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LyUItI6dpCmy"
      },
      "source": [
        "### Elgamal Key Generation\n",
        "\n",
        "Let's begin by implementing key generation for Elgamal. The pseudocode is below.\n",
        "\n",
        "$\\KeyGen(1^n)$:\n",
        "1. Sample $a \\gets \\Fp$.\n",
        "2. Set $\\sk := a$ and $\\pk := a \\cdot G$.\n",
        "3. Output $(\\sk, \\pk)$.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kbD02qJtoQgF"
      },
      "outputs": [],
      "source": [
        "def keygen():\n",
        "  # hint: use `Scalar.rand()`\n",
        "  # your code here\n",
        "  return (sk, pk)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "--wXU18VtJtQ"
      },
      "source": [
        "### Elgamal Encryption\n",
        "\n",
        "Great, we can now generate keys. What about encrypting a message? Let's try to implement that. Recall that the pseudocode for Elgamal encryption looks like the following:\n",
        "\n",
        "\n",
        "$\\Encrypt(\\pk = A, m \\in \\Group) \\to c$:\n",
        "1. Sample $r \\gets \\Fp$.\n",
        "2. Output $c := (rG, m + r \\cdot A)$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4dDV_-EUs0om"
      },
      "outputs": [],
      "source": [
        "def encrypt(pk, m):\n",
        "  # your code here\n",
        "  return c"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dn3CD2V54z2B"
      },
      "source": [
        "### Elgamal Decryption\n",
        "\n",
        "\n",
        "Recall that the opening algorithm looks like the following:\n",
        "\n",
        "\n",
        "$\\Decrypt(\\sk = a, c = (R, C)) \\to m$:\n",
        "1. Output $m := C - a \\cdot R$.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MV3MsHTf5hlG"
      },
      "outputs": [],
      "source": [
        "def decrypt(sk, c):\n",
        "  # your code here\n",
        "  return m"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S5if6wwAAO_x"
      },
      "source": [
        "### Correctness\n",
        "\n",
        "Now let's check that our algorithms work!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "22JcJ_SpAWtx"
      },
      "outputs": [],
      "source": [
        "(sk, pk) = keygen()\n",
        "print((sk, pk))\n",
        "\n",
        "m = Scalar(10) * G # encode our message (10) into the group\n",
        "print(m)\n",
        "c = encrypt(pk, m)\n",
        "m_prime = decrypt(sk, c)\n",
        "\n",
        "print(m_prime)\n",
        "assert(m == m_prime)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1QUy2rt3AOyN"
      },
      "source": [
        "### Security\n",
        "\n",
        "We can't directly check that the scheme is secure, because if it is secure, then we can't find an efficient distinguishing adversary. But let's try something simpler: using a different key to decrypt."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SFQlqJIrD6XC"
      },
      "outputs": [],
      "source": [
        "(sk_2, pk_2) = keygen()\n",
        "print((sk_2, pk_2))\n",
        "\n",
        "m_2 = decrypt(sk_2, c)\n",
        "print(m_2)\n",
        "assert(m != m_2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Insecure encryption\n",
        "\n",
        "Nila is trying to encrypt a message to Rambo, but she's not the smartest dog, and so she thinks it's okay to use limited amounts of randomness for encryption. In particular, she uses the following *insecure* encryption algorithm:\n",
        "\n",
        "\n",
        "$\\mathsf{InsecureEncrypt}(\\pk = A, m \\in \\Group) \\to c$:\n",
        "1. Sample $r \\gets \\{0, 1\\}^10$.\n",
        "2. Output $c := (rG, m + r \\cdot A)$.\n",
        "\n",
        "Implement this algorithm, using the provided snippet for generating 10-bit random scalars."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import random\n",
        "\n",
        "def ten_bit_random_scalar():\n",
        "    return Scalar(random.getrandbits(10))\n",
        "\n",
        "def insecure_encrypt(pk, m):\n",
        "  # your code here\n",
        "  return c"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now, let's try to attack an encryption of a message $m$ that Nila sent to Rambo. Implement the following function:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def attack(pk, c):\n",
        "    # Your code here\n",
        "    return m\n",
        "\n",
        "(sk, pk) = keygen()\n",
        "print((sk, pk))\n",
        "\n",
        "m = Scalar(random.getrandbits(32)) * G # encode our message into the group\n",
        "print(m)\n",
        "c = insecure_encrypt(pk, m)\n",
        "m_prime = attack(pk, c)\n",
        "assert(m_prime == m)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "authorship_tag": "ABX9TyNVtbZitLXONrrgZMLaolJw",
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": ".venv",
      "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.13.1"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}