docs/tuto-p4_bitsliceAES.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bit-slicing AES\n",
    "In this example, we build a circuit of bit-slicing AES. This implementation is based on the white-box AES of [BU18](https://github.com/cryptolu/whitebox/blob/master/synthesis/whitebox/ciphers/AES/bitaes.py)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Operations on vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "\n",
    "class Vector(list):\n",
    "    ZERO = 0\n",
    "    WIDTH = None\n",
    "\n",
    "    @classmethod\n",
    "    def make(cls, lst):\n",
    "        lst = list(lst)\n",
    "        if cls.WIDTH is not None:\n",
    "            assert len(lst) == cls.WIDTH\n",
    "        return cls(lst)\n",
    "\n",
    "    def split(self, n=2):\n",
    "        assert len(self) % n == 0\n",
    "        w = len(self) // n\n",
    "        return Vector(self.make(self[i:i+w]) for i in range(0, len(self), w))\n",
    "\n",
    "    def rol(self, n=1):\n",
    "        n %= len(self)\n",
    "        return self.make(self[n:] + self[:n])\n",
    "\n",
    "    def ror(self, n=1):\n",
    "        return self.rol(-n)\n",
    "\n",
    "    def __repr__(self):\n",
    "        return \"<Vector len=%d list=%r>\" % (len(self), list(self))\n",
    "\n",
    "    def flatten(self):\n",
    "        if isinstance(self[0], Vector):\n",
    "            return self[0].concat(*self[1:])\n",
    "        return reduce(operator.add, list(self))\n",
    "\n",
    "    def map(self, f, with_coord=False):\n",
    "        if with_coord:\n",
    "            return self.make(f(i, v) for i, v in enumerate(self))\n",
    "        else:\n",
    "            return self.make(f(v) for v in self)\n",
    "\n",
    "    def __xor__(self, other):\n",
    "        assert isinstance(other, Vector)\n",
    "        assert len(self) == len(other)\n",
    "        return self.make(a ^ b for a, b in zip(self, other))\n",
    "\n",
    "    def __or__(self, other):\n",
    "        assert isinstance(other, Vector)\n",
    "        assert len(self) == len(other)\n",
    "        return self.make(a | b for a, b in zip(self, other))\n",
    "\n",
    "    def __and__(self, other):\n",
    "        assert isinstance(other, Vector)\n",
    "        assert len(self) == len(other)\n",
    "        return self.make(a & b for a, b in zip(self, other))\n",
    "\n",
    "    def set(self, x, val):\n",
    "        return self.make(v if i != x else val for i, v in enumerate(self))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Operations on matrices (Rect)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class Rect(object):\n",
    "    def __init__(self, vec, h=None, w=None):\n",
    "        assert h or w\n",
    "        if h:\n",
    "            w = len(vec) // h\n",
    "        elif w:\n",
    "            h = len(vec) // w\n",
    "        assert w * h == len(vec)\n",
    "        self.w, self.h = w, h\n",
    "\n",
    "        self.lst = []\n",
    "        for i in range(0, len(vec), w):\n",
    "            self.lst.append(list(vec[i:i+w]))\n",
    "\n",
    "    @classmethod\n",
    "    def from_rect(cls, rect):\n",
    "        self = object.__new__(cls)\n",
    "        self.lst = rect\n",
    "        self.h = len(rect)\n",
    "        self.w = len(rect[0])\n",
    "        return self\n",
    "\n",
    "    def __getitem__(self, pos):\n",
    "        y, x = pos\n",
    "        return self.lst[y][x]\n",
    "\n",
    "    def __setitem__(self, pos, val):\n",
    "        y, x = pos\n",
    "        self.lst[y][x] = val\n",
    "\n",
    "    def row(self, i):\n",
    "        return Vector(self.lst[i])\n",
    "\n",
    "    def col(self, i):\n",
    "        return Vector(self.lst[y][i] for y in range(self.h))\n",
    "\n",
    "    def set_row(self, y, vec):\n",
    "        for x in range(self.w):\n",
    "            self.lst[y][x] = vec[x]\n",
    "        return self\n",
    "\n",
    "    def set_col(self, x, vec):\n",
    "        for y in range(self.h):\n",
    "            self.lst[y][x] = vec[y]\n",
    "        return self\n",
    "\n",
    "    def apply(self, f, with_coord=False):\n",
    "        for y in range(self.h):\n",
    "            if with_coord:\n",
    "                self.lst[y] = [f(y, x, v) for x, v in enumerate(self.lst[y])]\n",
    "            else:\n",
    "                self.lst[y] = list(map(f, self.lst[y]))\n",
    "        return self\n",
    "\n",
    "    def apply_row(self, x, func):\n",
    "        return self.set_row(x, func(self.row(x)))\n",
    "\n",
    "    def apply_col(self, x, func):\n",
    "        return self.set_col(x, func(self.col(x)))\n",
    "\n",
    "    def flatten(self):\n",
    "        lst = []\n",
    "        for v in self.lst:\n",
    "            lst += v\n",
    "        return Vector(lst)\n",
    "\n",
    "    def zipwith(self, f, other):\n",
    "        assert isinstance(other, Rect)\n",
    "        assert self.h == other.h\n",
    "        assert self.w == other.w\n",
    "        return Rect(\n",
    "            [f(a, b) for a, b in zip(self.flatten(), other.flatten())],\n",
    "            h=self.h, w=self.w\n",
    "        )\n",
    "\n",
    "    def transpose(self):\n",
    "        rect = [[self.lst[y][x] for y in range(self.h)] for x in range(self.w)]\n",
    "        return Rect.from_rect(rect=rect)\n",
    "\n",
    "    def __repr__(self):\n",
    "        return \"<Rect %dx%d>\" % (self.h, self.w)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bit-slicing implementation of Sbox"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Not(x):\n",
    "    return 1^x\n",
    "\n",
    "def GF_SQ_2(A): return A[1], A[0]\n",
    "def GF_SCLW_2(A): return A[1], A[1] ^ A[0]\n",
    "def GF_SCLW2_2(A): return A[1] ^ A[0], A[0]\n",
    "\n",
    "def GF_MULS_2(A, ab, B, cd):\n",
    "    abcd = (ab & cd)\n",
    "    p = ((A[1] & B[1])) ^ abcd\n",
    "    q = ((A[0] & B[0])) ^ abcd\n",
    "    return q, p\n",
    "\n",
    "def GF_MULS_SCL_2(A, ab, B, cd):\n",
    "    t = (A[0] & B[0])\n",
    "    p = ((ab & cd)) ^ t\n",
    "    q = ((A[1] & B[1])) ^ t\n",
    "    return q, p\n",
    "\n",
    "def XOR_LIST(a, b):\n",
    "    return [a ^ b for a, b in zip(a, b)]\n",
    "\n",
    "def NotOr(a, b):\n",
    "    # return Not(a | b)\n",
    "    return Not(a) & Not(b)\n",
    "\n",
    "def GF_INV_4(A):\n",
    "    a = A[2:4]\n",
    "    b = A[0:2]\n",
    "    sa = a[1] ^ a[0]\n",
    "    sb = b[1] ^ b[0]\n",
    "\n",
    "    ab = GF_MULS_2(a, sa, b, sb)\n",
    "    ab2 = GF_SQ_2(XOR_LIST(a, b))\n",
    "    ab2N = GF_SCLW2_2(ab2)\n",
    "    d = GF_SQ_2(XOR_LIST(ab, ab2N))\n",
    "\n",
    "    c = [\n",
    "        NotOr(sa, sb) ^ (Not(a[0] & b[0])),\n",
    "        NotOr(a[1], b[1]) ^ (Not(sa & sb)),\n",
    "    ]\n",
    "\n",
    "    sd = d[1] ^ d[0]\n",
    "    p = GF_MULS_2(d, sd, b, sb)\n",
    "    q = GF_MULS_2(d, sd, a, sa)\n",
    "    return q + p\n",
    "\n",
    "def GF_SQ_SCL_4(A):\n",
    "    a = A[2:4]\n",
    "    b = A[0:2]\n",
    "    ab2 = GF_SQ_2(a ^ b)\n",
    "    b2 = GF_SQ_2(b)\n",
    "    b2N2 = GF_SCLW_2(b2)\n",
    "    return b2N2 + ab2\n",
    "\n",
    "def GF_MULS_4(A, a, Al, Ah, aa, B, b, Bl, Bh, bb):\n",
    "    ph = GF_MULS_2(A[2:4], Ah, B[2:4], Bh)\n",
    "    pl = GF_MULS_2(A[0:2], Al, B[0:2], Bl)\n",
    "    p = GF_MULS_SCL_2(a, aa, b, bb)\n",
    "    return XOR_LIST(pl, p) + XOR_LIST(ph, p) #(pl ^ p), (ph ^ p)\n",
    "\n",
    "def GF_INV_8(A):\n",
    "    a = A[4:8]\n",
    "    b = A[0:4]\n",
    "    sa = XOR_LIST(a[2:4], a[0:2])\n",
    "    sb = XOR_LIST(b[2:4], b[0:2])\n",
    "    al = a[1] ^ a[0]\n",
    "    ah = a[3] ^ a[2]\n",
    "    aa = sa[1] ^ sa[0]\n",
    "    bl = b[1] ^ b[0]\n",
    "    bh = b[3] ^ b[2]\n",
    "    bb = sb[1] ^ sb[0]\n",
    "\n",
    "    c1 = (ah & bh)\n",
    "    c2 = (sa[0] & sb[0])\n",
    "    c3 = (aa & bb)\n",
    "\n",
    "    c = [\n",
    "        (NotOr(a[0] , b[0] ) ^ ((al & bl))) ^ ((sa[1] & sb[1])) ^ Not(c2), #0\n",
    "        (NotOr(al   , bl   ) ^ (Not(a[1] & b[1]))) ^ c2 ^ c3 , #1\n",
    "        (NotOr(sa[1], sb[1]) ^ (Not(a[2] & b[2]))) ^ c1 ^ c2 , #2\n",
    "        (NotOr(sa[0], sb[0]) ^ (Not(a[3] & b[3]))) ^ c1 ^ c3 , #3\n",
    "    ]\n",
    "    d = GF_INV_4(c)\n",
    "\n",
    "    sd = XOR_LIST(d[2:4], d[0:2])\n",
    "    dl = d[1] ^ d[0]\n",
    "    dh = d[3] ^ d[2]\n",
    "    dd = sd[1] ^ sd[0]\n",
    "    p = GF_MULS_4(d, sd, dl, dh, dd, b, sb, bl, bh, bb)\n",
    "    q = GF_MULS_4(d, sd, dl, dh, dd, a, sa, al, ah, aa)\n",
    "    return q + p\n",
    "\n",
    "\n",
    "def MUX21I(A, B, s): #return ((~A & s) ^ (~B & ~s)\n",
    "    return Not(A if s else B)\n",
    "\n",
    "def SELECT_NOT_8( A, B, s):\n",
    "    Q = [None] * 8\n",
    "    for i in range(8):\n",
    "        Q[i] = MUX21I(A[i], B[i], s)\n",
    "    return Q\n",
    "\n",
    "def Sbox(A, encrypt):\n",
    "    R1 = A[7] ^ A[5]\n",
    "    R2 = A[7] ^ Not(A[4])\n",
    "    R3 = A[6] ^ A[0]\n",
    "    R4 = A[5] ^ Not(R3)\n",
    "    R5 = A[4] ^ R4\n",
    "    R6 = A[3] ^ A[0]\n",
    "    R7 = A[2] ^ R1\n",
    "    R8 = A[1] ^ R3\n",
    "    R9 = A[3] ^ R8\n",
    "\n",
    "    B = [None] * 8\n",
    "    B[7] = R7 ^ Not(R8)\n",
    "    B[6] = R5\n",
    "    B[5] = A[1] ^ R4\n",
    "    B[4] = R1 ^ Not(R3)\n",
    "    B[3] = A[1]^ R2 ^ R6\n",
    "    B[2] = Not( A[0])\n",
    "    B[1] = R4\n",
    "    B[0] = A[2] ^ Not(R9)\n",
    "\n",
    "    Y = [None] * 8\n",
    "    Y[7] = R2\n",
    "    Y[6] = A[4] ^ R8\n",
    "    Y[5] = A[6] ^ A[4]\n",
    "    Y[4] = R9\n",
    "    Y[3] = A[6] ^ Not(R2)\n",
    "    Y[2] = R7\n",
    "    Y[1] = A[4] ^ R6\n",
    "    Y[0] = A[1] ^ R5\n",
    "\n",
    "    Z = SELECT_NOT_8(B, Y, encrypt)\n",
    "    C = GF_INV_8(Z)\n",
    "\n",
    "    T1 = C[7] ^ C[3]\n",
    "    T2 = C[6] ^ C[4]\n",
    "    T3 = C[6] ^ C[0]\n",
    "    T4 = C[5] ^ Not(C[3])\n",
    "    T5 = C[5] ^ Not(T1)\n",
    "    T6 = C[5] ^ Not(C[1])\n",
    "    T7 = C[4] ^ Not(T6)\n",
    "    T8 = C[2] ^ T4\n",
    "    T9 = C[1] ^ T2\n",
    "    T10 = T3 ^ T5\n",
    "\n",
    "    D = [None] * 8\n",
    "    D[7] = T4\n",
    "    D[6] = T1\n",
    "    D[5] = T3\n",
    "    D[4] = T5\n",
    "    D[3] = T2 ^ T5\n",
    "    D[2] = T3 ^ T8\n",
    "    D[1] = T7\n",
    "    D[0] = T9\n",
    "\n",
    "    X = [None] * 8\n",
    "    X[7] = C[4] ^ Not(C[1])\n",
    "    X[6] = C[1] ^ T10\n",
    "    X[5] = C[2] ^ T10\n",
    "    X[4] = C[6] ^ Not(C[1])\n",
    "    X[3] = T8 ^ T9\n",
    "    X[2] = C[7] ^ Not(T7)\n",
    "    X[1] = T6\n",
    "    X[0] = Not(C[2])\n",
    "    return SELECT_NOT_8(D, X, encrypt)\n",
    "\n",
    "\n",
    "def bitSbox(A, inverse=False):\n",
    "    res = Sbox(A[::-1], encrypt=1-inverse)[::-1]\n",
    "    return res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bit-slicing implementations of ShiftRow and MixColumn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ShiftRow(row, nr, inverse=False):\n",
    "    if inverse:\n",
    "        nr = -nr\n",
    "    off = nr % 4\n",
    "    return row[off:] + row[:off]\n",
    "\n",
    "def MixColumn(col, inverse=False):\n",
    "    res = [[0] * 8 for _ in range(4)]\n",
    "    table = MCi_TABLE if inverse else MC_TABLE\n",
    "    for yi in range(4):\n",
    "        for yj in range(8):\n",
    "            y = yi * 8 + yj\n",
    "            for x in table[y]:\n",
    "                xi, xj = divmod(x, 8)\n",
    "                res[yi][yj] ^= col[xi][xj]\n",
    "    return res\n",
    "\n",
    "# y -> set of x indices to xor\n",
    "MC_TABLE = [{1, 8, 9, 16, 24}, {2, 9, 10, 17, 25}, {3, 10, 11, 18, 26}, {0, 4, 8, 11, 12, 19, 27}, {0, 5, 8, 12, 13, 20, 28}, {6, 13, 14, 21, 29}, {0, 7, 8, 14, 15, 22, 30}, {0, 8, 15, 23, 31}, {0, 9, 16, 17, 24}, {1, 10, 17, 18, 25}, {2, 11, 18, 19, 26}, {3, 8, 12, 16, 19, 20, 27}, {4, 8, 13, 16, 20, 21, 28}, {5, 14, 21, 22, 29}, {6, 8, 15, 16, 22, 23, 30}, {7, 8, 16, 23, 31}, {0, 8, 17, 24, 25}, {1, 9, 18, 25, 26}, {2, 10, 19, 26, 27}, {3, 11, 16, 20, 24, 27, 28}, {4, 12, 16, 21, 24, 28, 29}, {5, 13, 22, 29, 30}, {6, 14, 16, 23, 24, 30, 31}, {7, 15, 16, 24, 31}, {0, 1, 8, 16, 25}, {1, 2, 9, 17, 26}, {2, 3, 10, 18, 27}, {0, 3, 4, 11, 19, 24, 28}, {0, 4, 5, 12, 20, 24, 29}, {5, 6, 13, 21, 30}, {0, 6, 7, 14, 22, 24, 31}, {0, 7, 15, 23, 24}]\n",
    "MCi_TABLE = [{1, 2, 3, 8, 9, 11, 16, 18, 19, 24, 27}, {0, 2, 3, 4, 8, 9, 10, 12, 16, 17, 19, 20, 24, 25, 28}, {1, 3, 4, 5, 8, 9, 10, 11, 13, 17, 18, 20, 21, 24, 25, 26, 29}, {2, 4, 5, 6, 8, 9, 10, 11, 12, 14, 16, 18, 19, 21, 22, 25, 26, 27, 30}, {1, 2, 5, 6, 7, 10, 12, 13, 15, 16, 17, 18, 20, 22, 23, 24, 26, 28, 31}, {1, 6, 7, 8, 9, 13, 14, 17, 21, 23, 24, 25, 29}, {2, 7, 8, 9, 10, 14, 15, 16, 18, 22, 25, 26, 30}, {0, 1, 2, 8, 10, 15, 17, 18, 23, 26, 31}, {0, 3, 9, 10, 11, 16, 17, 19, 24, 26, 27}, {0, 1, 4, 8, 10, 11, 12, 16, 17, 18, 20, 24, 25, 27, 28}, {0, 1, 2, 5, 9, 11, 12, 13, 16, 17, 18, 19, 21, 25, 26, 28, 29}, {1, 2, 3, 6, 10, 12, 13, 14, 16, 17, 18, 19, 20, 22, 24, 26, 27, 29, 30}, {0, 2, 4, 7, 9, 10, 13, 14, 15, 18, 20, 21, 23, 24, 25, 26, 28, 30, 31}, {0, 1, 5, 9, 14, 15, 16, 17, 21, 22, 25, 29, 31}, {1, 2, 6, 10, 15, 16, 17, 18, 22, 23, 24, 26, 30}, {2, 7, 8, 9, 10, 16, 18, 23, 25, 26, 31}, {0, 2, 3, 8, 11, 17, 18, 19, 24, 25, 27}, {0, 1, 3, 4, 8, 9, 12, 16, 18, 19, 20, 24, 25, 26, 28}, {1, 2, 4, 5, 8, 9, 10, 13, 17, 19, 20, 21, 24, 25, 26, 27, 29}, {0, 2, 3, 5, 6, 9, 10, 11, 14, 18, 20, 21, 22, 24, 25, 26, 27, 28, 30}, {0, 1, 2, 4, 6, 7, 8, 10, 12, 15, 17, 18, 21, 22, 23, 26, 28, 29, 31}, {1, 5, 7, 8, 9, 13, 17, 22, 23, 24, 25, 29, 30}, {0, 2, 6, 9, 10, 14, 18, 23, 24, 25, 26, 30, 31}, {1, 2, 7, 10, 15, 16, 17, 18, 24, 26, 31}, {0, 1, 3, 8, 10, 11, 16, 19, 25, 26, 27}, {0, 1, 2, 4, 8, 9, 11, 12, 16, 17, 20, 24, 26, 27, 28}, {0, 1, 2, 3, 5, 9, 10, 12, 13, 16, 17, 18, 21, 25, 27, 28, 29}, {0, 1, 2, 3, 4, 6, 8, 10, 11, 13, 14, 17, 18, 19, 22, 26, 28, 29, 30}, {2, 4, 5, 7, 8, 9, 10, 12, 14, 15, 16, 18, 20, 23, 25, 26, 29, 30, 31}, {0, 1, 5, 6, 9, 13, 15, 16, 17, 21, 25, 30, 31}, {0, 1, 2, 6, 7, 8, 10, 14, 17, 18, 22, 26, 31}, {0, 2, 7, 9, 10, 15, 18, 23, 24, 25, 26}]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bit-slicing implementation of Key Schedule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "Rcon = [0x8d, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36,\n",
    "        0x6c, 0xd8, 0xab, 0x4d, 0x9a, 0x2f, 0x5e, 0xbc, 0x63, 0xc6, 0x97,\n",
    "        0x35, 0x6a, 0xd4, 0xb3, 0x7d, 0xfa, 0xef, 0xc5, 0x91, 0x39, 0x72,\n",
    "        0xe4, 0xd3, 0xbd, 0x61, 0xc2, 0x9f, 0x25, 0x4a, 0x94, 0x33, 0x66,\n",
    "        0xcc, 0x83, 0x1d, 0x3a, 0x74, 0xe8, 0xcb, 0x8d, 0x01, 0x02, 0x04,\n",
    "        0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36, 0x6c, 0xd8, 0xab, 0x4d,\n",
    "        0x9a, 0x2f, 0x5e, 0xbc, 0x63, 0xc6, 0x97, 0x35, 0x6a, 0xd4, 0xb3,\n",
    "        0x7d, 0xfa, 0xef, 0xc5, 0x91, 0x39, 0x72, 0xe4, 0xd3, 0xbd, 0x61,\n",
    "        0xc2, 0x9f, 0x25, 0x4a, 0x94, 0x33, 0x66, 0xcc, 0x83, 0x1d, 0x3a,\n",
    "        0x74, 0xe8, 0xcb, 0x8d, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40,\n",
    "        0x80, 0x1b, 0x36, 0x6c, 0xd8, 0xab, 0x4d, 0x9a, 0x2f, 0x5e, 0xbc,\n",
    "        0x63, 0xc6, 0x97, 0x35, 0x6a, 0xd4, 0xb3, 0x7d, 0xfa, 0xef, 0xc5,\n",
    "        0x91, 0x39, 0x72, 0xe4, 0xd3, 0xbd, 0x61, 0xc2, 0x9f, 0x25, 0x4a,\n",
    "        0x94, 0x33, 0x66, 0xcc, 0x83, 0x1d, 0x3a, 0x74, 0xe8, 0xcb, 0x8d,\n",
    "        0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36, 0x6c,\n",
    "        0xd8, 0xab, 0x4d, 0x9a, 0x2f, 0x5e, 0xbc, 0x63, 0xc6, 0x97, 0x35,\n",
    "        0x6a, 0xd4, 0xb3, 0x7d, 0xfa, 0xef, 0xc5, 0x91, 0x39, 0x72, 0xe4,\n",
    "        0xd3, 0xbd, 0x61, 0xc2, 0x9f, 0x25, 0x4a, 0x94, 0x33, 0x66, 0xcc,\n",
    "        0x83, 0x1d, 0x3a, 0x74, 0xe8, 0xcb, 0x8d, 0x01, 0x02, 0x04, 0x08,\n",
    "        0x10, 0x20, 0x40, 0x80, 0x1b, 0x36, 0x6c, 0xd8, 0xab, 0x4d, 0x9a,\n",
    "        0x2f, 0x5e, 0xbc, 0x63, 0xc6, 0x97, 0x35, 0x6a, 0xd4, 0xb3, 0x7d,\n",
    "        0xfa, 0xef, 0xc5, 0x91, 0x39, 0x72, 0xe4, 0xd3, 0xbd, 0x61, 0xc2,\n",
    "        0x9f, 0x25, 0x4a, 0x94, 0x33, 0x66, 0xcc, 0x83, 0x1d, 0x3a, 0x74,\n",
    "        0xe8, 0xcb ]\n",
    "\n",
    "def tobin(x, n):\n",
    "    return tuple(map(int, bin(x).lstrip(\"0b\").rjust(n, \"0\")))\n",
    "\n",
    "def c8(c):\n",
    "    return Vector(tobin(c, 8))\n",
    "\n",
    "BitRcon = list(map(c8, Rcon))\n",
    "\n",
    "def ks_rotate(word):\n",
    "    return word[1:] + word[:1]\n",
    "\n",
    "def ks_core(word, iteration):\n",
    "    word = word.rol(1)\n",
    "    word = word.map(lambda b: Vector(bitSbox(b)))\n",
    "    word = word.set(0, word[0] ^ BitRcon[iteration])\n",
    "    return word\n",
    "\n",
    "def KS_round(kstate, rno):\n",
    "    t = ks_core(kstate.col(3), rno+1)\n",
    "    kstate.apply_col(0, lambda c: c ^ t)\n",
    "    t = kstate.col(0)\n",
    "    kstate.apply_col(1, lambda c: c ^ t)\n",
    "    t = kstate.col(1)\n",
    "    kstate.apply_col(2, lambda c: c ^ t)\n",
    "    t = kstate.col(2)\n",
    "    kstate.apply_col(3, lambda c: c ^ t)\n",
    "    return kstate\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bit-slicing AES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def BitAES(plaintext, key, rounds=10):\n",
    "    bx = Vector(plaintext).split(16)\n",
    "    bk = Vector(key).split(16)\n",
    "\n",
    "    state = Rect(bx, w=4, h=4).transpose()\n",
    "    kstate = Rect(bk, w=4, h=4).transpose()\n",
    "\n",
    "    for rno in range(rounds):\n",
    "        state = AK(state, kstate)\n",
    "        state = SB(state)\n",
    "        state = SR(state)\n",
    "        if rno < rounds-1:\n",
    "            state = MC(state)\n",
    "        kstate = KS(kstate, rno)\n",
    "    state = AK(state, kstate)\n",
    "\n",
    "    state = state.transpose()\n",
    "    kstate = kstate.transpose()\n",
    "    bits = sum( map(list, state.flatten()), [])\n",
    "    kbits = sum( map(list, kstate.flatten()), [])\n",
    "    return bits, kbits\n",
    "\n",
    "def AK(state, kstate):\n",
    "    return state.zipwith(lambda a, b: a ^ b, kstate)\n",
    "\n",
    "def SB(state, inverse=False):\n",
    "    return state.apply(lambda v: Vector(bitSbox(v, inverse=inverse)))\n",
    "\n",
    "def SR(state, inverse=False):\n",
    "    for y in range(4):\n",
    "        state.apply_row(y, lambda row: ShiftRow(row, y, inverse=inverse))\n",
    "    return state\n",
    "\n",
    "def MC(state, inverse=False):\n",
    "    for x in range(4):\n",
    "        state.apply_col(x, lambda v: list(map(Vector, MixColumn(v))))\n",
    "    return state\n",
    "\n",
    "def KS(kstate, rno):\n",
    "    return KS_round(kstate, rno)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build a circuit for bit-slicing AES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "skip"
    ]
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from circkit.boolean import BooleanCircuit\n",
    "\n",
    "C = BooleanCircuit()\n",
    "pt = C.add_inputs(n=128, format=\"m%d\")\n",
    "k = C.add_inputs(n=128, format=\"k%d\")\n",
    "ct, k10 = BitAES(pt, k, rounds=10)\n",
    "C.add_output(ct)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can evaluate the circuit to test its correctness."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ciphertext: 69c4e0d86a7b0430d8cdb78070b4c55a\n"
     ]
    }
   ],
   "source": [
    "def str2bin(s):\n",
    "    return list(map(int, \"\".join(bin(ord(c))[2:].zfill(8) for c in s)))\n",
    "\n",
    "def hex2bin(s):\n",
    "    return list(map(int, \"\".join(bin(c)[2:].zfill(8) for c in bytes.fromhex(s))))\n",
    "\n",
    "def bin2hex(s):\n",
    "    assert len(s) % 8 == 0\n",
    "    v = int(\"\".join(map(str, s)), 2)\n",
    "    v = (\"%x\" %v).zfill(len(s) // 4)\n",
    "    return v\n",
    "\n",
    "# Expected ciphertext: 69c4e0d86a7b0430d8cdb78070b4c55a\n",
    "# See the AES documentation of NIST: https://nvlpubs.nist.gov/nistpubs/fips/nist.fips.197.pdf\n",
    "MSG = \"00112233445566778899aabbccddeeff\"\n",
    "KEY = \"000102030405060708090a0b0c0d0e0f\"\n",
    "\n",
    "inp = hex2bin(MSG) + hex2bin(KEY)\n",
    "out = C.evaluate(inp)\n",
    "out = bin2hex(out)\n",
    "print(f\"Ciphertext: {out}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a white-box implementation, we can treat the key as 128 constants of bits instead of 128 input nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ciphertext: 69c4e0d86a7b0430d8cdb78070b4c55a\n"
     ]
    }
   ],
   "source": [
    "C = BooleanCircuit()\n",
    "pt = C.add_inputs(n=128, format=\"m%d\")\n",
    "k = [C.add_const(b) for b in hex2bin(KEY)]\n",
    "ct, k10 = BitAES(pt, k, rounds=10)\n",
    "C.add_output(ct)\n",
    "\n",
    "inp = hex2bin(MSG)\n",
    "out = C.evaluate(inp)\n",
    "out = bin2hex(out)\n",
    "print(f\"Ciphertext: {out}\")"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "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.9.5"
  },
  "vscode": {
   "interpreter": {
    "hash": "986ed710fd0861d223fd21335cd6da488ecd593829391b548355316a07d356bf"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}