1. Intro (lab).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3bb2a3ef",
   "metadata": {},
   "source": [
    "# Tutorial 1: Cryptography in practice\n",
    "\n",
    "# Lab Content\n",
    "\n",
    "In this tutorial, we will explore how abstract cryptographic operations are implemented in software. Our goal is to identify the \"critical paths\" where secret data (keys) interact with the algorithm. Note that hardware-focused implementations and implementations in hardware may look differently.\n",
    "\n",
    "## Tasks\n",
    "1. Implement core components of AES, RSA, and ECC.\n",
    "2. Find and examine components of open-source implementations.\n",
    "\n",
    "## Task 1: AES-128 - The S-Box and Key Addition\n",
    "\n",
    "Your task is to implement the SBOX lookup table, the `sub_bytes` and `add_round_key` functions. The test at the end of the cell should pass.\n",
    "\n",
    "Consult the [AES standard](https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.197-upd1.pdf) and [Wikipedia](https://en.wikipedia.org/wiki/Rijndael_S-box).\n",
    "\n",
    "**Note**: The solution is available in the [aes](aes.py) module. You can look at it if you are struggling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e350810",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Standard AES S-Box\n",
    "# Task 1: Find the AES SBOX table (256 bytes) and paste it here.\n",
    "SBOX = []\n",
    "\n",
    "def sub_bytes(state: list) -> list:\n",
    "    \"\"\"\n",
    "    The SubBytes step of AES.\n",
    "    \"\"\"\n",
    "    # Task 1: Implement the SubBytes step here\n",
    "    raise NotImplementedError\n",
    "\n",
    "def add_round_key(state: list, round_key: list) -> list:\n",
    "    \"\"\"\n",
    "    The AddRoundKey step of AES.\n",
    "    \"\"\"\n",
    "    # Task 1: Implement the AddRoundKey step here\n",
    "    raise NotImplementedError\n",
    "\n",
    "def shift_rows(state: list) -> list:\n",
    "    \"\"\"\n",
    "    The ShiftRows step of AES.\n",
    "    \"\"\"\n",
    "    return [state[0], state[5], state[10], state[15], state[4], state[9], state[14], state[3], state[8], state[13], state[2], state[7], state[12], state[1], state[6], state[11]]\n",
    "\n",
    "def xtime(a: int) -> int:\n",
    "    return (((a << 1) ^ 0x1B) & 0xFF) if (a & 0x80) else (a << 1)\n",
    "\n",
    "def mix_columns(state: list) -> list:\n",
    "    \"\"\"\n",
    "    The MixColumns step of AES.\n",
    "    \"\"\"\n",
    "    res = [0] * 16\n",
    "    for i in range(4):\n",
    "        c = state[i*4 : i*4+4]\n",
    "        res[i*4+0] = xtime(c[0]) ^ (xtime(c[1]) ^ c[1]) ^ c[2] ^ c[3]\n",
    "        res[i*4+1] = c[0] ^ xtime(c[1]) ^ (xtime(c[2]) ^ c[2]) ^ c[3]\n",
    "        res[i*4+2] = c[0] ^ c[1] ^ xtime(c[2]) ^ (xtime(c[3]) ^ c[3])\n",
    "        res[i*4+3] = (xtime(c[0]) ^ c[0]) ^ c[1] ^ c[2] ^ xtime(c[3])\n",
    "    return res\n",
    "\n",
    "def expand_key(key: bytes) -> list:\n",
    "    \"\"\"\n",
    "    AES-128 Key Expansion.\n",
    "    \"\"\"\n",
    "    Rcon = [0x00, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1B, 0x36]\n",
    "    # Put the key into a grid\n",
    "    w = [list(key[i:i+4]) for i in range(0, 16, 4)]\n",
    "    for i in range(4, 44):\n",
    "        temp = list(w[i-1])\n",
    "        if i % 4 == 0:\n",
    "            temp = [SBOX[b] for b in temp[1:] + temp[:1]]\n",
    "            temp[0] ^= Rcon[i//4]\n",
    "        w.append([w[i-4][j] ^ temp[j] for j in range(4)])\n",
    "    return [sum(w[i:i+4], []) for i in range(0, 44, 4)]\n",
    "\n",
    "def aes_encrypt(plaintext: bytes, key: bytes):\n",
    "    \"\"\"\n",
    "    An AES-128 encryption loop.\n",
    "    \"\"\"\n",
    "    round_keys = expand_key(key)\n",
    "    state = list(plaintext)\n",
    "    \n",
    "    # Initial Round\n",
    "    state = add_round_key(state, round_keys[0])\n",
    "    \n",
    "    # Rounds 1 to 9\n",
    "    for i in range(1, 10):\n",
    "        state = sub_bytes(state)\n",
    "        state = shift_rows(state)\n",
    "        state = mix_columns(state)\n",
    "        state = add_round_key(state, round_keys[i])\n",
    "        \n",
    "    # Final Round (No MixColumns)\n",
    "    state = sub_bytes(state)\n",
    "    state = shift_rows(state)\n",
    "    state = add_round_key(state, round_keys[10])\n",
    "    \n",
    "    return bytes(state)\n",
    "\n",
    "\n",
    "# --- TEST ---\n",
    "key = bytes.fromhex(\"2b7e151628aed2a6abf7158809cf4f3c\")\n",
    "pt = bytes.fromhex(\"6bc1bee22e409f96e93d7e117393172a\")\n",
    "expected_ct = \"3ad77bb40d7a3660a89ecaf32466ef97\"\n",
    "\n",
    "try:\n",
    "    if not SBOX:\n",
    "        raise ValueError(\"Empty SBOX\")\n",
    "    ct = aes_encrypt(pt, key)\n",
    "    print(f\"Plaintext:  {pt.hex()}\")\n",
    "    print(f\"Ciphertext: {ct.hex()}\")\n",
    "    print(f\"Success:    {ct.hex() == expected_ct}\")\n",
    "except Exception as e:\n",
    "    print(f\"Encryption failed: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e575a5ed",
   "metadata": {},
   "source": [
    "## Task 2: RSA - Modular exponentiation\n",
    "\n",
    "The core operation in RSA is modular exponentiation. There are many ways to implement it. What they have in common is that they process the private key in some rgular fashion (i.e., left-to-right, right-to-left, in \"windows\", as a \"comb\", or sliding).\n",
    "\n",
    "Your task is to implement this algorithm. You can choose whichever algorithm you want, it just has to be correct. I suggest a simple one.\n",
    "\n",
    "Consult [Exponentiation by squaring](https://en.wikipedia.org/wiki/Exponentiation_by_squaring), examine the [mod](mod.py) file for how the Mod class works. It represent an integer modulo a given $n$.\n",
    "\n",
    "**Note**: The solution is available in the [rsa](rsa.py) module. You can look at it if you are struggling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ffa9874",
   "metadata": {},
   "outputs": [],
   "source": [
    "import secrets\n",
    "from mod import Mod\n",
    "\n",
    "def square_and_multiply(base: int, exponent: int, modulus: int) -> int:\n",
    "    \"\"\"\n",
    "    Computes (base**exponent) % modulus using square-and-multiply.\n",
    "    \"\"\"\n",
    "    # Task 2: Implement square and multiply.\n",
    "    raise NotImplementedError\n",
    "\n",
    "def pkcs1_v15_pad(message: bytes, modulus_len: int) -> int:\n",
    "    \"\"\"\n",
    "    A PKCS#1 v1.5 padding scheme (for encryption).\n",
    "    Structure: 0x00 || 0x02 || PS || 0x00 || Message\n",
    "    \n",
    "\t\thttps://www.rfc-editor.org/rfc/rfc8017#section-7.2.1\n",
    "    \n",
    "    Warning: This is not a good padding scheme! It is vulnerable\n",
    "    to Bleichenbacher's attack (see lecture 5 of Real World Crypto Engineering).\n",
    "    \"\"\"\n",
    "    if len(message) > modulus_len - 11:\n",
    "        raise ValueError(\"Message too long for modulus\")\n",
    "    \n",
    "    # Padding string (non-zero random bytes)\n",
    "    ps_len = modulus_len - len(message) - 3\n",
    "    ps = b\"\"\n",
    "    while len(ps) < ps_len:\n",
    "        new_byte = secrets.token_bytes(1)\n",
    "        if new_byte != b\"\\x00\":\n",
    "            ps += new_byte\n",
    "            \n",
    "    padded = b\"\\x00\\x02\" + ps + b\"\\x00\" + message\n",
    "    return int.from_bytes(padded, \"big\")\n",
    "\n",
    "def pkcs1_v15_unpad(padded_int: int, modulus_len: int) -> bytes:\n",
    "    \"\"\"\n",
    "    Removes the PKCS#1 v1.5 padding (for decryption).\n",
    "\t    \n",
    "\t    https://www.rfc-editor.org/rfc/rfc8017#section-7.2.2\n",
    "    \n",
    "    Warning: This is not a good padding scheme! It is vulnerable\n",
    "    to Bleichenbacher's attack (see lecture 5 of Real World Crypto Engineering).\n",
    "    \"\"\"\n",
    "    padded = padded_int.to_bytes(modulus_len, \"big\")\n",
    "    print(padded)\n",
    "    if padded[0:2] != b\"\\x00\\x02\":\n",
    "        raise ValueError(\"Invalid padding prefix\")\n",
    "    \n",
    "    # Find the separator\n",
    "    sep_idx = padded.find(b\"\\x00\", 2)\n",
    "    if sep_idx == -1:\n",
    "        raise ValueError(\"No separator found\")\n",
    "        \n",
    "    return padded[sep_idx + 1:]\n",
    "\n",
    "def rsa_encrypt(m: bytes, e: int, n: int) -> int:\n",
    "    \"\"\"\n",
    "    RSA Encryption with padding.\n",
    "    \"\"\"\n",
    "    mod_len = (n.bit_length() + 7) // 8\n",
    "    padded_m = pkcs1_v15_pad(m, mod_len)\n",
    "    return square_and_multiply(padded_m, e, n)\n",
    "\n",
    "def rsa_decrypt(c: int, d: int, n: int) -> bytes:\n",
    "    \"\"\"\n",
    "    RSA Decryption and unpadding.\n",
    "    \"\"\"\n",
    "    mod_len = (n.bit_length() + 7) // 8\n",
    "    padded_m = square_and_multiply(c, d, n)\n",
    "    return pkcs1_v15_unpad(padded_m, mod_len)\n",
    "\n",
    "\n",
    "# --- TEST ---\n",
    "p = 0x7c644070a1b9b3155ae937f9688144299d857299d5fb47bb7ac5b55aaa684b151c4b4b415f8256d655febb6991a8877fb6c00b6a67a7a9c83f3580d4d333eac99541fdaa951ce38d37abd857d94f034a97666ae1899e4dca933b79ab3d304ee310bcef35a85e889ce1808497f5d9f5f257fdebe6b8aada36f06f674b7f81121b\n",
    "q = 0x9cef0fe8c6a49c2a6df5ca3042bcf01e0226847bac81d8656c2203b6486cbcf955c70a224395a5e5bc7693a5b2fdcd38479a3265c4b4898bb1c9228b1d8f1dfd2398f9572f2193d759e976211cd68bfd3f927f2279979d17c7248c199b0aa60b8dce47285e9b098b1a0fe59844b60d74389f477400a131e404866c8ecac9368b\n",
    "n = p * q\n",
    "e = 17\n",
    "d = 0xd74ed47f5ea5a6903da8bb3f47ef76578627a2f72c1eee2ebf28572cf0ebb10b4df8445afd38255d465d98eebffea614d7cd821219c7770934c49ae87ee7597bb1b36e687123e181d60f6dcb0a15a6a0adafb37fe2188730cfc5884869a7172042772b31065f61f87728ea0332032a3d522df377368e357764a9eddff6d11afa6a611e06c4e97951515c9e0a83d731faa8298db0f954b59e5ad4a3e1bb0f3fc7aea9c310dfd6104242b1616507ee42aba9b3ae24a15e9e52236e2631adb35fa60ebd4602da92a3af6a8649ddf870e8166c8cf9e5c7d9f0d69b7bfe06c8b205859c2bb103aa989f4136eb605be5cd756f64a25a9cca4470df8dd743b00e8563d\n",
    "msg = b\"Hi!\"\n",
    "\n",
    "try:\n",
    "    cipher = rsa_encrypt(msg, e, n)\n",
    "    decrypted = rsa_decrypt(cipher, d, n)\n",
    "    print(f\"Original:  {msg}\")\n",
    "    print(f\"Encrypted: {cipher}\")\n",
    "    print(f\"Decrypted: {decrypted}\")\n",
    "    print(f\"Success:   {msg == decrypted}\")\n",
    "except NotImplementedError:\n",
    "    print(\"Implement square_and_multiply.\")\n",
    "except Exception as e:\n",
    "    print(f\"RSA failed: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bfd88ee",
   "metadata": {},
   "source": [
    "## Task 3: ECC - Scalar Multiplication\n",
    "\n",
    "Elliptic Curve Cryptography uses a similar loop structure, but the operations are point addition and point doubling. Again, there are many algorithms for this.\n",
    "\n",
    "Consult [Scalar multiplication](https://en.wikipedia.org/wiki/Elliptic_curve_point_multiplication) and an excellent [scalar multiplication](https://cryptojedi.org/peter/data/eccss-20130911b.pdf) talk by Peter Schwabe, also examine the [mod](mod.py) file for how the Mod class works. It represent an integer modulo a given $n$.\n",
    "\n",
    "**Note**: The solution is available in the [ecc](ecc.py) module. You can look at it if you are struggling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cf2c768",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "from hashlib import sha1\n",
    "from typing import Tuple, Union\n",
    "\n",
    "from mod import Mod\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Point(object):\n",
    "    \"\"\"A point on an elliptic curve.\"\"\"\n",
    "    x: Mod\n",
    "    y: Mod\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Curve(object):\n",
    "    \"\"\"An elliptic curve.\"\"\"\n",
    "\n",
    "    p: int\n",
    "    a: Mod\n",
    "    b: Mod\n",
    "    g: Point\n",
    "    n: int\n",
    "\n",
    "    def __init__(self, p: int, a: int, b: int, gx: int, gy: int, n: int):\n",
    "        self.p = p\n",
    "        self.a = Mod(a, p)\n",
    "        self.b = Mod(b, p)\n",
    "        self.g = Point(Mod(gx, p), Mod(gy, p))\n",
    "        self.n = n\n",
    "\n",
    "    def add(self, p: Point, q: Point):\n",
    "        \"\"\"\n",
    "        Add two points on an elliptic curve (If the points are equal, the result is erroneous).\n",
    "        See <https://en.wikipedia.org/wiki/Elliptic_curve_point_multiplication#Point_addition>.\n",
    "        \"\"\"\n",
    "        λ = (q.y - p.y) / (q.x - p.x)\n",
    "        x = λ**2 - p.x - q.x\n",
    "        y = λ * (p.x - x) - p.y\n",
    "        return Point(x, y)\n",
    "\n",
    "    def dbl(self, p: Point):\n",
    "        \"\"\"\n",
    "        Double a point on an elliptic curve.\n",
    "        See <https://en.wikipedia.org/wiki/Elliptic_curve_point_multiplication#Point_doubling>.\n",
    "        \"\"\"\n",
    "        λ = (3 * p.x**2 + self.a) / (2 * p.y)\n",
    "        x = λ**2 - p.x - p.x\n",
    "        y = λ * (p.x - x) - p.y\n",
    "        return Point(x, y)\n",
    "\n",
    "\n",
    "# A popular curve: https://std.neuromancer.sk/secg/secp256r1\n",
    "curve_secp256r1 = Curve(\n",
    "    0xFFFFFFFF00000001000000000000000000000000FFFFFFFFFFFFFFFFFFFFFFFF,\n",
    "    0xFFFFFFFF00000001000000000000000000000000FFFFFFFFFFFFFFFFFFFFFFFC,\n",
    "    0x5AC635D8AA3A93E7B3EBBD55769886BC651D06B0CC53B0F63BCE3C3E27D2604B,\n",
    "    0x6B17D1F2E12C4247F8BCE6E563A440F277037D812DEB33A0F4A13945D898C296,\n",
    "    0x4FE342E2FE1A7F9B8EE7EB4A7C0F9E162BCE33576B315ECECBB6406837BF51F5,\n",
    "    0xFFFFFFFF00000000FFFFFFFFFFFFFFFFBCE6FAADA7179E84F3B9CAC2FC632551,\n",
    ")\n",
    "\n",
    "\n",
    "def scalarmult(point: Point, scalar: Union[int, Mod], curve: Curve) -> Point:\n",
    "    \"\"\"\n",
    "    Perform scalar multiplication of the `point` with the `scalar` on the elliptic curve given by `curve`.\n",
    "    See <https://en.wikipedia.org/wiki/Elliptic_curve_point_multiplication> for some methods.\n",
    "    Returns `[scalar]point`.\n",
    "    \"\"\"\n",
    "    s = int(scalar)\n",
    "    # Task 3: Implement the scalar multiplication algorithm here.\n",
    "    raise NotImplementedError\n",
    "\n",
    "def keygen(curve: Curve) -> Tuple[Mod, Point]:\n",
    "    \"\"\"\n",
    "    Generate an ECC keypair on the `curve`.\n",
    "    Returns a tuple of `private key, public key`.\n",
    "    \"\"\"\n",
    "    private = Mod.random(curve.n)\n",
    "    public = scalarmult(curve.g, int(private), curve)\n",
    "    return private, public\n",
    "\n",
    "def sign(message: bytes, private: Mod, curve: Curve) -> Tuple[Mod, Mod]:\n",
    "    \"\"\"\n",
    "    Sign the `message` using the `private` key on the `curve`.\n",
    "    Returns the signature tuple `r, s`.\n",
    "\n",
    "    https://en.wikipedia.org/wiki/Elliptic_Curve_Digital_Signature_Algorithm\n",
    "    \"\"\"\n",
    "    # h = SHA1(message) and then the ECDSA trimming.\n",
    "    h = int(sha1(message).hexdigest(), 16) >> (\n",
    "        0 if 160 <= curve.n.bit_length() else 160 - curve.n.bit_length()\n",
    "    )\n",
    "    # k = random integer from ℤₙ, the Mod.random function should be considered constant-time.\n",
    "    k = Mod.random(curve.n)\n",
    "    # r = ([k]G)_x mod n\n",
    "    r = Mod(int(scalarmult(curve.g, int(k), curve).x), curve.n)\n",
    "    # s = k^-1 * (H(message) + r * x) mod n\n",
    "    s = k ** (-1) * (h + r * private)\n",
    "    return r, s\n",
    "\n",
    "def verify(signature: Tuple[Mod, Mod], message: bytes, public: Point, curve: Curve) -> bool:\n",
    "    \"\"\"\n",
    "    Verify the `signature` on the `message` using the `public` key on the `curve`.\n",
    "\n",
    "    Note, this is missing some checks.\n",
    "\n",
    "    https://en.wikipedia.org/wiki/Elliptic_Curve_Digital_Signature_Algorithm\n",
    "    \"\"\"\n",
    "    # h = SHA1(message) and then the ECDSA trimming.\n",
    "    h = int(sha1(message).hexdigest(), 16) >> (\n",
    "        0 if 160 <= curve.n.bit_length() else 160 - curve.n.bit_length()\n",
    "    )\n",
    "    r, s = signature\n",
    "    if r == 0 or s == 0:\n",
    "        raise ValueError(\"Malformed signature\")\n",
    "    sinv = s ** (-1)\n",
    "    u1 = h * sinv\n",
    "    u2 = r * sinv\n",
    "    u1G = scalarmult(curve.g, u1, curve)\n",
    "    u2P = scalarmult(public, u2, curve)\n",
    "    res = curve.add(u1G, u2P)\n",
    "    return int(r) == int(res.x)\n",
    "\n",
    "\n",
    "# --- TEST ---\n",
    "message = b\"Hi!\"\n",
    "\n",
    "try:\n",
    "    privkey = Mod(0xd836c8d5317e144e95f2fd225c75c7eb8f3e386514342dbaf4d89ffd0cc5b0ba, curve_secp256r1.n)\n",
    "    pubkey = scalarmult(curve_secp256r1.g, privkey, curve_secp256r1)\n",
    "    \n",
    "    signature = sign(message, privkey, curve_secp256r1)\n",
    "    verfied = verify(signature, message, pubkey, curve_secp256r1)\n",
    "\n",
    "    print(f\"Message = {message}\")\n",
    "    print(f\"Signature = {signature}\")\n",
    "    print(f\"Verified = {verfied}\")\n",
    "except NotImplementedError:\n",
    "    print(\"Implement scalarmult.\")\n",
    "except Exception as e:\n",
    "    print(f\"ECDSA failed: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b13c2c8f",
   "metadata": {},
   "source": [
    "## Examining open-source crypto libraries\n",
    "\n",
    "See the slides for the next steps."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}

aes.py

# Task 1: Find the AES SBOX table (256 bytes) and paste it here.
SBOX = [
    0x63,
    0x7C,
    0x77,
    0x7B,
    0xF2,
    0x6B,
    0x6F,
    0xC5,
    0x30,
    0x01,
    0x67,
    0x2B,
    0xFE,
    0xD7,
    0xAB,
    0x76,
    0xCA,
    0x82,
    0xC9,
    0x7D,
    0xFA,
    0x59,
    0x47,
    0xF0,
    0xAD,
    0xD4,
    0xA2,
    0xAF,
    0x9C,
    0xA4,
    0x72,
    0xC0,
    0xB7,
    0xFD,
    0x93,
    0x26,
    0x36,
    0x3F,
    0xF7,
    0xCC,
    0x34,
    0xA5,
    0xE5,
    0xF1,
    0x71,
    0xD8,
    0x31,
    0x15,
    0x04,
    0xC7,
    0x23,
    0xC3,
    0x18,
    0x96,
    0x05,
    0x9A,
    0x07,
    0x12,
    0x80,
    0xE2,
    0xEB,
    0x27,
    0xB2,
    0x75,
    0x09,
    0x83,
    0x2C,
    0x1A,
    0x1B,
    0x6E,
    0x5A,
    0xA0,
    0x52,
    0x3B,
    0xD6,
    0xB3,
    0x29,
    0xE3,
    0x2F,
    0x84,
    0x53,
    0xD1,
    0x00,
    0xED,
    0x20,
    0xFC,
    0xB1,
    0x5B,
    0x6A,
    0xCB,
    0xBE,
    0x39,
    0x4A,
    0x4C,
    0x58,
    0xCF,
    0xD0,
    0xEF,
    0xAA,
    0xFB,
    0x43,
    0x4D,
    0x33,
    0x85,
    0x45,
    0xF9,
    0x02,
    0x7F,
    0x50,
    0x3C,
    0x9F,
    0xA8,
    0x51,
    0xA3,
    0x40,
    0x8F,
    0x92,
    0x9D,
    0x38,
    0xF5,
    0xBC,
    0xB6,
    0xDA,
    0x21,
    0x10,
    0xFF,
    0xF3,
    0xD2,
    0xCD,
    0x0C,
    0x13,
    0xEC,
    0x5F,
    0x97,
    0x44,
    0x17,
    0xC4,
    0xA7,
    0x7E,
    0x3D,
    0x64,
    0x5D,
    0x19,
    0x73,
    0x60,
    0x81,
    0x4F,
    0xDC,
    0x22,
    0x2A,
    0x90,
    0x88,
    0x46,
    0xEE,
    0xB8,
    0x14,
    0xDE,
    0x5E,
    0x0B,
    0xDB,
    0xE0,
    0x32,
    0x3A,
    0x0A,
    0x49,
    0x06,
    0x24,
    0x5C,
    0xC2,
    0xD3,
    0xAC,
    0x62,
    0x91,
    0x95,
    0xE4,
    0x79,
    0xE7,
    0xC8,
    0x37,
    0x6D,
    0x8D,
    0xD5,
    0x4E,
    0xA9,
    0x6C,
    0x56,
    0xF4,
    0xEA,
    0x65,
    0x7A,
    0xAE,
    0x08,
    0xBA,
    0x78,
    0x25,
    0x2E,
    0x1C,
    0xA6,
    0xB4,
    0xC6,
    0xE8,
    0xDD,
    0x74,
    0x1F,
    0x4B,
    0xBD,
    0x8B,
    0x8A,
    0x70,
    0x3E,
    0xB5,
    0x66,
    0x48,
    0x03,
    0xF6,
    0x0E,
    0x61,
    0x35,
    0x57,
    0xB9,
    0x86,
    0xC1,
    0x1D,
    0x9E,
    0xE1,
    0xF8,
    0x98,
    0x11,
    0x69,
    0xD9,
    0x8E,
    0x94,
    0x9B,
    0x1E,
    0x87,
    0xE9,
    0xCE,
    0x55,
    0x28,
    0xDF,
    0x8C,
    0xA1,
    0x89,
    0x0D,
    0xBF,
    0xE6,
    0x42,
    0x68,
    0x41,
    0x99,
    0x2D,
    0x0F,
    0xB0,
    0x54,
    0xBB,
    0x16,
]


def sub_bytes(state: list) -> list:
    """
    The SubBytes step of AES.
    """
    # Task 2: Implement the SubBytes step here
    return [SBOX[b] for b in state]


def add_round_key(state: list, round_key: list) -> list:
    """
    The AddRoundKey step of AES.
    """
    # Task 3: Implement the AddRoundKey step here
    return [s ^ r for s, r in zip(state, round_key)]


def shift_rows(state: list) -> list:
    """
    The ShiftRows step of AES.
    """
    return [
        state[0],
        state[5],
        state[10],
        state[15],
        state[4],
        state[9],
        state[14],
        state[3],
        state[8],
        state[13],
        state[2],
        state[7],
        state[12],
        state[1],
        state[6],
        state[11],
    ]


def xtime(a: int) -> int:
    return (((a << 1) ^ 0x1B) & 0xFF) if (a & 0x80) else (a << 1)


def mix_columns(state: list) -> list:
    """
    The MixColumns step of AES.
    """
    res = [0] * 16
    for i in range(4):
        c = state[i * 4 : i * 4 + 4]
        res[i * 4 + 0] = xtime(c[0]) ^ (xtime(c[1]) ^ c[1]) ^ c[2] ^ c[3]
        res[i * 4 + 1] = c[0] ^ xtime(c[1]) ^ (xtime(c[2]) ^ c[2]) ^ c[3]
        res[i * 4 + 2] = c[0] ^ c[1] ^ xtime(c[2]) ^ (xtime(c[3]) ^ c[3])
        res[i * 4 + 3] = (xtime(c[0]) ^ c[0]) ^ c[1] ^ c[2] ^ xtime(c[3])
    return res


def expand_key(key: bytes) -> list:
    """
    AES-128 Key Expansion.
    """
    Rcon = [0x00, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1B, 0x36]
    w = [list(key[i : i + 4]) for i in range(0, 16, 4)]
    for i in range(4, 44):
        temp = list(w[i - 1])
        if i % 4 == 0:
            temp = [SBOX[b] for b in temp[1:] + temp[:1]]
            temp[0] ^= Rcon[i // 4]
        w.append([w[i - 4][j] ^ temp[j] for j in range(4)])
    return [sum(w[i : i + 4], []) for i in range(0, 44, 4)]


def aes_encrypt(plaintext: bytes, key: bytes):
    """
    An AES-128 encryption loop.
    """
    round_keys = expand_key(key)
    state = list(plaintext)

    # Initial Round
    state = add_round_key(state, round_keys[0])

    # Rounds 1 to 9
    for i in range(1, 10):
        state = sub_bytes(state)
        state = shift_rows(state)
        state = mix_columns(state)
        state = add_round_key(state, round_keys[i])

    # Final Round (No MixColumns)
    state = sub_bytes(state)
    state = shift_rows(state)
    state = add_round_key(state, round_keys[10])

    return bytes(state)

ecc.py

from dataclasses import dataclass
from hashlib import sha1
from typing import Tuple, Union

from mod import Mod


@dataclass
class Point(object):
    """A point on an elliptic curve."""

    x: Mod
    y: Mod


@dataclass
class Curve(object):
    """An elliptic curve."""

    p: int
    a: Mod
    b: Mod
    g: Point
    n: int

    def __init__(self, p: int, a: int, b: int, gx: int, gy: int, n: int):
        self.p = p
        self.a = Mod(a, p)
        self.b = Mod(b, p)
        self.g = Point(Mod(gx, p), Mod(gy, p))
        self.n = n

    def add(self, p: Point, q: Point):
        """
        Add two points on an elliptic curve (If the points are equal, the result is erroneous).
        See <https://en.wikipedia.org/wiki/Elliptic_curve_point_multiplication#Point_addition>.
        """
        λ = (q.y - p.y) / (q.x - p.x)
        x = λ**2 - p.x - q.x
        y = λ * (p.x - x) - p.y
        return Point(x, y)

    def dbl(self, p: Point):
        """
        Double a point on an elliptic curve.
        See <https://en.wikipedia.org/wiki/Elliptic_curve_point_multiplication#Point_doubling>.
        """
        λ = (3 * p.x**2 + self.a) / (2 * p.y)
        x = λ**2 - p.x - p.x
        y = λ * (p.x - x) - p.y
        return Point(x, y)


curve_secp256r1 = Curve(
    0xFFFFFFFF00000001000000000000000000000000FFFFFFFFFFFFFFFFFFFFFFFF,
    0xFFFFFFFF00000001000000000000000000000000FFFFFFFFFFFFFFFFFFFFFFFC,
    0x5AC635D8AA3A93E7B3EBBD55769886BC651D06B0CC53B0F63BCE3C3E27D2604B,
    0x6B17D1F2E12C4247F8BCE6E563A440F277037D812DEB33A0F4A13945D898C296,
    0x4FE342E2FE1A7F9B8EE7EB4A7C0F9E162BCE33576B315ECECBB6406837BF51F5,
    0xFFFFFFFF00000000FFFFFFFFFFFFFFFFBCE6FAADA7179E84F3B9CAC2FC632551,
)


def scalarmult(point: Point, scalar: Union[int, Mod], curve: Curve) -> Point:
    """
    Perform scalar multiplication of the `point` with the `scalar` on the elliptic curve given by `curve`.
    See <https://en.wikipedia.org/wiki/Elliptic_curve_point_multiplication> for some methods.
    Returns `[scalar]point`.
    """
    r0 = point
    r1 = curve.dbl(point)
    s = int(scalar)
    for i in range(s.bit_length() - 2, -1, -1):
        if s & (1 << i) == 0:
            r1 = curve.add(r0, r1)
            r0 = curve.dbl(r0)
        else:
            r0 = curve.add(r0, r1)
            r1 = curve.dbl(r1)
    return r0


def keygen(curve: Curve) -> Tuple[Mod, Point]:
    """
    Generate an ECC keypair on the `curve`.
    Returns a tuple of `private key, public key`.
    """
    private = Mod.random(curve.n)
    public = scalarmult(curve.g, int(private), curve)
    return private, public


def sign(message: bytes, private: Mod, curve: Curve) -> Tuple[Mod, Mod]:
    """
    Sign the `message` using the `private` key on the `curve`.
    Returns the signature tuple `r, s`.

    https://en.wikipedia.org/wiki/Elliptic_Curve_Digital_Signature_Algorithm
    """
    # h = SHA1(message) and then the ECDSA trimming.
    h = int(sha1(message).hexdigest(), 16) >> (
        0 if 160 <= curve.n.bit_length() else 160 - curve.n.bit_length()
    )
    # k = random integer from ℤₙ, the Mod.random function should be considered constant-time.
    k = Mod.random(curve.n)
    # r = ([k]G)_x mod n
    r = Mod(int(scalarmult(curve.g, int(k), curve).x), curve.n)
    # s = k^-1 * (H(message) + r * x) mod n
    s = k ** (-1) * (h + r * private)
    return r, s

def verify(signature: Tuple[Mod, Mod], message: bytes, public: Point, curve: Curve) -> bool:
    """
    Verify the `signature` on the `message` using the `public` key on the `curve`.

    Note, this is missing some checks.

    https://en.wikipedia.org/wiki/Elliptic_Curve_Digital_Signature_Algorithm
    """
    # h = SHA1(message) and then the ECDSA trimming.
    h = int(sha1(message).hexdigest(), 16) >> (
        0 if 160 <= curve.n.bit_length() else 160 - curve.n.bit_length()
    )
    r, s = signature
    if r == 0 or s == 0:
        raise ValueError("Malformed signature")
    sinv = s ** (-1)
    u1 = h * sinv
    u2 = r * sinv
    u1G = scalarmult(curve.g, u1, curve)
    u2P = scalarmult(public, u2, curve)
    res = curve.add(u1G, u2P)
    return int(r) == int(res.x)

mod.py

import secrets
from functools import wraps


def check(func):
    @wraps(func)
    def method(self, other):
        if type(other) is type(self):
            if self.n != other.n:
                raise ValueError
        elif isinstance(other, int):
            other = self.__class__(other, self.n)
        else:
            raise ValueError
        return func(self, other)

    return method


def extgcd(a, b):
    """Extended Euclid's greatest common denominator algorithm."""
    if abs(b) > abs(a):
        (x, y, d) = extgcd(b, a)
        return y, x, d

    if abs(b) == 0:
        return 1, 0, a

    x1, x2, y1, y2 = 0, 1, 1, 0
    while abs(b) > 0:
        q, r = divmod(a, b)
        x = x2 - q * x1
        y = y2 - q * y1
        a, b, x2, x1, y2, y1 = b, r, x1, x, y1, y

    return x2, y2, a


class Mod(object):
    """
    An element x of ℤₙ.

    Construct like this:

    Mod(5, 7)

    Works like this:

    Mod(5, 7) + Mod(3, 7) = Mod(1, 7)

    Then, get the integer out like this:

    int(Mod(3, 7)) = 3
    """

    def __init__(self, x: int, n: int):
        self.x: int = x % n
        self.n: int = n

    @check
    def __add__(self, other):
        return Mod((self.x + other.x) % self.n, self.n)

    @check
    def __radd__(self, other):
        return self + other

    @check
    def __sub__(self, other):
        return Mod((self.x - other.x) % self.n, self.n)

    @check
    def __rsub__(self, other):
        return -self + other

    def __neg__(self):
        return Mod(self.n - self.x, self.n)

    def inverse_gcd(self):
        x, y, d = extgcd(self.x, self.n)
        return Mod(x, self.n)

    def inverse_pow(self):
        return self ** (self.n - 2)

    def inverse(self):
        return self.inverse_gcd()

    def __invert__(self):
        return self.inverse()

    @check
    def __mul__(self, other):
        return Mod((self.x * other.x) % self.n, self.n)

    @check
    def __rmul__(self, other):
        return self * other

    @check
    def __truediv__(self, other):
        return self * ~other

    @check
    def __rtruediv__(self, other):
        return ~self * other

    @check
    def __floordiv__(self, other):
        return self * ~other

    @check
    def __rfloordiv__(self, other):
        return ~self * other

    @check
    def __div__(self, other):
        return self.__floordiv__(other)

    @check
    def __rdiv__(self, other):
        return self.__rfloordiv__(other)

    @check
    def __divmod__(self, divisor):
        q, r = divmod(self.x, divisor.x)
        return Mod(q, self.n), Mod(r, self.n)

    def __bytes__(self):
        return self.x.to_bytes((self.n.bit_length() + 7) // 8, byteorder="big")

    @staticmethod
    def random(n: int):
        return Mod(secrets.randbelow(n), n)

    def __int__(self):
        return self.x

    def __index__(self):
        return self.x

    def __eq__(self, other):
        if type(other) is int:
            return self.x == (other % self.n)
        if type(other) is not Mod:
            return False
        return self.x == other.x and self.n == other.n

    def __ne__(self, other):
        return not self == other

    def __repr__(self):
        return str(self.x)

    def __pow__(self, n):
        if type(n) is not int:
            raise TypeError
        if n == 0:
            return Mod(1, self.n)
        if n < 0:
            return self.inverse() ** (-n)
        if n == 1:
            return Mod(self.x, self.n)

        q = self
        r = self if n & 1 else Mod(1, self.n)

        i = 2
        while i <= n:
            q = q * q
            if n & i == i:
                r = q * r
            i = i << 1
        return r

requirements.txt

jupyter
notebook
jupytext

rsa.py

import secrets

from mod import Mod


def square_and_multiply(base: int, exponent: int, modulus: int) -> int:
    """
    Computes (base**exponent) % modulus using square-and-multiply.
    """
    # Task 4: Implement square and multiply
    res = Mod(1, modulus)
    bits = bin(exponent)[2:]

    for bit in bits:
        res = res**2
        if bit == "1":
            res = res * base

    return int(res)


def pkcs1_v15_pad(message: bytes, modulus_len: int) -> int:
    """
    A PKCS#1 v1.5 padding scheme (for encryption).
    Structure: 0x00 || 0x02 || PS || 0x00 || Message

        https://www.rfc-editor.org/rfc/rfc8017#section-7.2.1

    Warning: This is not a good padding scheme! It is vulnerable
    to Bleichenbacher's attack (see lecture 5 of Real World Crypto Engineering).
    """
    if len(message) > modulus_len - 11:
        raise ValueError("Message too long for modulus")

    # Padding string (non-zero random bytes)
    ps_len = modulus_len - len(message) - 3
    ps = b""
    while len(ps) < ps_len:
        new_byte = secrets.token_bytes(1)
        if new_byte != b"\x00":
            ps += new_byte

    padded = b"\x00\x02" + ps + b"\x00" + message
    return int.from_bytes(padded, "big")


def pkcs1_v15_unpad(padded_int: int, modulus_len: int) -> bytes:
    """
    Removes the PKCS#1 v1.5 padding (for decryption).

        https://www.rfc-editor.org/rfc/rfc8017#section-7.2.2

    Warning: This is not a good padding scheme! It is vulnerable
    to Bleichenbacher's attack (see lecture 5 of Real World Crypto Engineering).
    """
    padded = padded_int.to_bytes(modulus_len, "big")
    if padded[0:2] != b"\x00\x02":
        raise ValueError("Invalid padding prefix")

    # Find the separator
    sep_idx = padded.find(b"\x00", 2)
    if sep_idx == -1:
        raise ValueError("No separator found")

    return padded[sep_idx + 1 :]


def rsa_encrypt(m: bytes, e: int, n: int) -> int:
    """
    RSA Encryption with padding.
    """
    mod_len = (n.bit_length() + 7) // 8
    padded_m = pkcs1_v15_pad(m, mod_len)
    return square_and_multiply(padded_m, e, n)


def rsa_decrypt(c: int, d: int, n: int) -> bytes:
    """
    RSA Decryption and unpadding.
    """
    mod_len = (n.bit_length() + 7) // 8
    padded_m = square_and_multiply(c, d, n)
    return pkcs1_v15_unpad(padded_m, mod_len)