1. Intro (lab).ipynb

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)