key_exchanges_demo_poc.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyN6UReHMOFpB85F06bzXfI2",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ericyoc/key_exchanges_demo_poc/blob/main/key_exchanges_demo_poc.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#!pip install prettytable"
      ],
      "metadata": {
        "id": "lejTZ1S5dbSo"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "J4HyI3MAc8c1"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "from prettytable import PrettyTable"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Diffie-Hellman Key Exchange functions\n",
        "def generate_dh_prime():\n",
        "    return 23  # Small prime for demonstration\n",
        "\n",
        "def generate_base(prime):\n",
        "    return 5  # Should be a primitive root modulo p in practice\n",
        "\n",
        "def generate_private_key(prime):\n",
        "    return random.randint(1, prime - 1)\n",
        "\n",
        "def compute_public_key(prime, base, private_key):\n",
        "    return pow(base, private_key, prime)\n",
        "\n",
        "def compute_shared_secret(prime, public_key, private_key):\n",
        "    return pow(public_key, private_key, prime)"
      ],
      "metadata": {
        "id": "GnTY_tTJhxO1"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# RSA functions\n",
        "def generate_rsa_primes():\n",
        "    return 11, 17  # Small primes for demonstration\n",
        "\n",
        "def mod_inverse(e, phi):\n",
        "    for d in range(3, phi):\n",
        "        if (d * e) % phi == 1:\n",
        "            return d\n",
        "    raise ValueError(\"mod_inverse does not exist\")\n",
        "\n",
        "def generate_rsa_keypair():\n",
        "    p, q = generate_rsa_primes()\n",
        "    n = p * q\n",
        "    phi = (p - 1) * (q - 1)\n",
        "    e = 7  # Common choice for small examples\n",
        "    d = mod_inverse(e, phi)\n",
        "    return ((e, n), (d, n))\n",
        "\n",
        "def rsa_encrypt(public_key, plaintext):\n",
        "    e, n = public_key\n",
        "    return pow(plaintext, e, n)\n",
        "\n",
        "def rsa_decrypt(private_key, ciphertext):\n",
        "    d, n = private_key\n",
        "    return pow(ciphertext, d, n)\n",
        "\n",
        "def rsa_sign(private_key, message):\n",
        "    return rsa_encrypt(private_key, message)  # In practice, we'd hash the message first\n",
        "\n",
        "def rsa_verify(public_key, message, signature):\n",
        "    return rsa_decrypt(public_key, signature) == message"
      ],
      "metadata": {
        "id": "1-Db-Nk-h27Z"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def demonstrate_scenarios():\n",
        "    # Scenario 1: Diffie-Hellman Key Exchange\n",
        "    print(\"Scenario 1: Diffie-Hellman Key Exchange\")\n",
        "    prime = generate_dh_prime()\n",
        "    base = generate_base(prime)\n",
        "\n",
        "    alice_private = generate_private_key(prime)\n",
        "    alice_public = compute_public_key(prime, base, alice_private)\n",
        "\n",
        "    bob_private = generate_private_key(prime)\n",
        "    bob_public = compute_public_key(prime, base, bob_private)\n",
        "\n",
        "    alice_shared = compute_shared_secret(prime, bob_public, alice_private)\n",
        "    bob_shared = compute_shared_secret(prime, alice_public, bob_private)\n",
        "\n",
        "    table = PrettyTable()\n",
        "    table.title = \"Diffie-Hellman Key Exchange\"\n",
        "    table.field_names = [\"Step\", \"Alice\", \"Bob\", \"Eve\"]\n",
        "    table.add_row([\"1. Agree on parameters\", f\"p = {prime}, g = {base}\", f\"p = {prime}, g = {base}\", f\"Observes p = {prime}, g = {base}\"])\n",
        "    table.add_row([\"2. Generate private key\", f\"a = {alice_private}\", f\"b = {bob_private}\", \"Cannot know a or b\"])\n",
        "    table.add_row([\"3. Compute public key\", f\"A = g^a mod p = {alice_public}\", f\"B = g^b mod p = {bob_public}\", f\"Observes A = {alice_public}, B = {bob_public}\"])\n",
        "    table.add_row([\"4. Exchange public keys\", \"Sends A to Bob\", \"Sends B to Alice\", \"Intercepts A and B\"])\n",
        "    table.add_row([\"5. Compute shared secret\", f\"s = B^a mod p = {alice_shared}\", f\"s = A^b mod p = {bob_shared}\", \"Cannot compute s\"])\n",
        "    print(table)\n",
        "\n",
        "    # Scenario 2: RSA Key Exchange\n",
        "    print(\"\\nScenario 2: RSA Key Exchange\")\n",
        "    alice_public, alice_private = generate_rsa_keypair()\n",
        "    bob_public, bob_private = generate_rsa_keypair()\n",
        "\n",
        "    # Bob generates a symmetric key\n",
        "    symmetric_key = random.randint(1, 100)  # Small number for demonstration\n",
        "    encrypted_key = rsa_encrypt(alice_public, symmetric_key)\n",
        "    decrypted_key = rsa_decrypt(alice_private, encrypted_key)\n",
        "\n",
        "    table = PrettyTable()\n",
        "    table.title = \"RSA Key Exchange\"\n",
        "    table.field_names = [\"Step\", \"Alice\", \"Bob\", \"Eve\"]\n",
        "    table.add_row([\"1. Generate key pairs\", f\"Public: {alice_public}\\nPrivate: {alice_private}\", f\"Public: {bob_public}\\nPrivate: {bob_private}\", \"Observes public keys\"])\n",
        "    table.add_row([\"2. Exchange public keys\", \"Sends public key to Bob\", \"Sends public key to Alice\", \"Intercepts public keys\"])\n",
        "    table.add_row([\"3. Generate symmetric key\", \"Not involved\", f\"k = {symmetric_key}\", \"Unknown\"])\n",
        "    table.add_row([\"4. Encrypt symmetric key\", \"Not involved\", f\"c = k^e mod n = {encrypted_key}\", f\"Observes c = {encrypted_key}\"])\n",
        "    table.add_row([\"5. Send encrypted key\", \"Receives c\", \"Sends c to Alice\", \"Intercepts c\"])\n",
        "    table.add_row([\"6. Decrypt symmetric key\", f\"k = c^d mod n = {decrypted_key}\", \"Not involved\", \"Cannot decrypt\"])\n",
        "    print(table)\n",
        "\n",
        "    # Scenario 3: RSA Encryption and Decryption\n",
        "    print(\"\\nScenario 3: RSA Encryption and Decryption\")\n",
        "    alice_public, alice_private = generate_rsa_keypair()\n",
        "    message = 42  # Small number for demonstration\n",
        "    encrypted = rsa_encrypt(alice_public, message)\n",
        "    decrypted = rsa_decrypt(alice_private, encrypted)\n",
        "\n",
        "    table = PrettyTable()\n",
        "    table.title = \"RSA Encryption and Decryption\"\n",
        "    table.field_names = [\"Step\", \"Alice\", \"Bob\", \"Eve\"]\n",
        "    table.add_row([\"1. Generate key pair\", f\"Public: {alice_public}\\nPrivate: {alice_private}\", \"Not involved\", \"Observes public key\"])\n",
        "    table.add_row([\"2. Share public key\", \"Sends public key to Bob\", \"Receives Alice's public key\", \"Observes public key exchange\"])\n",
        "    table.add_row([\"3. Prepare message\", \"Not involved\", f\"m = {message}\", \"Unknown\"])\n",
        "    table.add_row([\"4. Encrypt message\", \"Not involved\", f\"c = m^e mod n = {encrypted}\", f\"Observes c = {encrypted}\"])\n",
        "    table.add_row([\"5. Send encrypted message\", \"Receives c\", \"Sends c to Alice\", \"Intercepts c\"])\n",
        "    table.add_row([\"6. Decrypt message\", f\"m = c^d mod n = {decrypted}\", \"Not involved\", \"Cannot decrypt\"])\n",
        "    print(table)\n",
        "\n",
        "    # Scenario 4: RSA Signing\n",
        "    print(\"\\nScenario 4: RSA Signing\")\n",
        "    message_to_sign = 42  # Small number for demonstration\n",
        "    signature = rsa_sign(alice_private, message_to_sign)\n",
        "    is_valid = rsa_verify(alice_public, message_to_sign, signature)\n",
        "\n",
        "    table = PrettyTable()\n",
        "    table.title = \"RSA Signing\"\n",
        "    table.field_names = [\"Step\", \"Alice\", \"Bob\", \"Eve\"]\n",
        "    table.add_row([\"1. Prepare message\", f\"m = {message_to_sign}\", \"Not involved\", \"Unknown\"])\n",
        "    table.add_row([\"2. Sign message\", f\"s = m^d mod n = {signature}\", \"Not involved\", \"Cannot sign\"])\n",
        "    table.add_row([\"3. Send message and signature\", \"Sends m and s to Bob\", \"Receives m and s\", \"Intercepts m and s\"])\n",
        "    table.add_row([\"4. Verify signature\", \"Not involved\", f\"m == s^e mod n: {is_valid}\", \"Can verify, cannot forge\"])\n",
        "    print(table)"
      ],
      "metadata": {
        "id": "Kqtp7TdVh6Xu"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "if __name__ == \"__main__\":\n",
        "    demonstrate_scenarios()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SdxC1bwwh-Lg",
        "outputId": "6be48620-af59-4c76-b2f8-8a9a1aa7d6d2"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Scenario 1: Diffie-Hellman Key Exchange\n",
            "+----------------------------------------------------------------------------------------------+\n",
            "|                                 Diffie-Hellman Key Exchange                                  |\n",
            "+--------------------------+--------------------+--------------------+-------------------------+\n",
            "|           Step           |       Alice        |        Bob         |           Eve           |\n",
            "+--------------------------+--------------------+--------------------+-------------------------+\n",
            "|  1. Agree on parameters  |   p = 23, g = 5    |   p = 23, g = 5    |  Observes p = 23, g = 5 |\n",
            "| 2. Generate private key  |       a = 11       |       b = 5        |    Cannot know a or b   |\n",
            "|  3. Compute public key   | A = g^a mod p = 22 | B = g^b mod p = 20 | Observes A = 22, B = 20 |\n",
            "| 4. Exchange public keys  |   Sends A to Bob   |  Sends B to Alice  |    Intercepts A and B   |\n",
            "| 5. Compute shared secret | s = B^a mod p = 22 | s = A^b mod p = 22 |     Cannot compute s    |\n",
            "+--------------------------+--------------------+--------------------+-------------------------+\n",
            "\n",
            "Scenario 2: RSA Key Exchange\n",
            "+----------------------------------------------------------------------------------------------------------+\n",
            "|                                             RSA Key Exchange                                             |\n",
            "+---------------------------+-------------------------+---------------------------+------------------------+\n",
            "|            Step           |          Alice          |            Bob            |          Eve           |\n",
            "+---------------------------+-------------------------+---------------------------+------------------------+\n",
            "|   1. Generate key pairs   |     Public: (7, 187)    |      Public: (7, 187)     |  Observes public keys  |\n",
            "|                           |    Private: (23, 187)   |     Private: (23, 187)    |                        |\n",
            "|  2. Exchange public keys  | Sends public key to Bob | Sends public key to Alice | Intercepts public keys |\n",
            "| 3. Generate symmetric key |       Not involved      |           k = 35          |        Unknown         |\n",
            "|  4. Encrypt symmetric key |       Not involved      |     c = k^e mod n = 18    |    Observes c = 18     |\n",
            "|   5. Send encrypted key   |        Receives c       |      Sends c to Alice     |      Intercepts c      |\n",
            "|  6. Decrypt symmetric key |    k = c^d mod n = 35   |        Not involved       |     Cannot decrypt     |\n",
            "+---------------------------+-------------------------+---------------------------+------------------------+\n",
            "\n",
            "Scenario 3: RSA Encryption and Decryption\n",
            "+------------------------------------------------------------------------------------------------------------------+\n",
            "|                                          RSA Encryption and Decryption                                           |\n",
            "+---------------------------+-------------------------+-----------------------------+------------------------------+\n",
            "|            Step           |          Alice          |             Bob             |             Eve              |\n",
            "+---------------------------+-------------------------+-----------------------------+------------------------------+\n",
            "|    1. Generate key pair   |     Public: (7, 187)    |         Not involved        |     Observes public key      |\n",
            "|                           |    Private: (23, 187)   |                             |                              |\n",
            "|    2. Share public key    | Sends public key to Bob | Receives Alice's public key | Observes public key exchange |\n",
            "|     3. Prepare message    |       Not involved      |            m = 42           |           Unknown            |\n",
            "|     4. Encrypt message    |       Not involved      |      c = m^e mod n = 15     |       Observes c = 15        |\n",
            "| 5. Send encrypted message |        Receives c       |       Sends c to Alice      |         Intercepts c         |\n",
            "|     6. Decrypt message    |    m = c^d mod n = 42   |         Not involved        |        Cannot decrypt        |\n",
            "+---------------------------+-------------------------+-----------------------------+------------------------------+\n",
            "\n",
            "Scenario 4: RSA Signing\n",
            "+--------------------------------------------------------------------------------------------------------+\n",
            "|                                              RSA Signing                                               |\n",
            "+-------------------------------+----------------------+----------------------+--------------------------+\n",
            "|              Step             |        Alice         |         Bob          |           Eve            |\n",
            "+-------------------------------+----------------------+----------------------+--------------------------+\n",
            "|       1. Prepare message      |        m = 42        |     Not involved     |         Unknown          |\n",
            "|        2. Sign message        | s = m^d mod n = 168  |     Not involved     |       Cannot sign        |\n",
            "| 3. Send message and signature | Sends m and s to Bob |   Receives m and s   |    Intercepts m and s    |\n",
            "|      4. Verify signature      |     Not involved     | m == s^e mod n: True | Can verify, cannot forge |\n",
            "+-------------------------------+----------------------+----------------------+--------------------------+\n"
          ]
        }
      ]
    }
  ]
}