# Sunshine CTF 2022 - aeschall

## AES Overview

AES has 4 steps

1. SubBytes
2. ShiftRows
3. MixColumns
4. AddRoundKey

If we got rid of the SubBytes and AddRoundKey steps, then AES would be a completely linear (not affine) process, meaning we can model it as a function $E(x) = A*x$.

$x$ can be viewed as a 128 x 1 column vector in ${\mathbb {F}}_2^{128}$, whose entries are $\in \mathbb {F}_2$.

This is because AES operates off of 16-byte chunks and thus we can view a message as a vector $x$ of its bits.

Below is an implementation of a toy AES cipher in Sage with just a ShiftRows and MixColumns step to show the linearity of the encryption.

## Toy AES Implementation

P.<x> = PolynomialRing(GF(2))
T.<z> = GF(2^8, modulus=x^8 + x^4 + x^3 + x + 1)
PR = PolynomialRing(T, [f'm{i}' for i in range(16)])
Mgens = PR.gens()

class AES:
    def __init__(self):
        pass

    def encrypt(self, plaintext):
        self.plain_state = plaintext

        for i in range(1, 10):
            self.plain_state = self.__shift_rows(self.plain_state)
            self.plain_state = self.__mix_columns(self.plain_state)

        self.plain_state = self.__shift_rows(self.plain_state)

        return self.plain_state
        

    def __shift_rows(self, M):
        s = [list(r) for r in M.rows()]
        s[0][1], s[1][1], s[2][1], s[3][1] = s[1][1], s[2][1], s[3][1], s[0][1]
        s[0][2], s[1][2], s[2][2], s[3][2] = s[2][2], s[3][2], s[0][2], s[1][2]
        s[0][3], s[1][3], s[2][3], s[3][3] = s[3][3], s[0][3], s[1][3], s[2][3]
        return Matrix(s)


    def __mix_columns(self, M):
        S = Matrix(T, [
            [T.fetch_int(2), T.fetch_int(3), T.fetch_int(1), T.fetch_int(1)],
            [T.fetch_int(1), T.fetch_int(2), T.fetch_int(3), T.fetch_int(1)],
            [T.fetch_int(1), T.fetch_int(1), T.fetch_int(2), T.fetch_int(3)],
            [T.fetch_int(3), T.fetch_int(1), T.fetch_int(1), T.fetch_int(2)],
        ])
        return S*M

cipher = AES()
plaintext = Matrix(PR, 4, 4, Mgens)
print(type(Mgens[0]))
print(type(plaintext[0, 0]))
print(plaintext)
print("------------------------")
enc = cipher.encrypt(plaintext)
print(enc)
print("------------------------")

new_Mat = Matrix(ZZ, 16, 16)
for i in range(4):
    for j in range(4):
        coeffs = []
        for k in range(16):
            coeffs.append(int(T(enc[i, j].coefficient(Mgens[k])).integer_representation()))
        #print(len(coeffs))
        new_Mat[4*i + j, :] = vector(ZZ, coeffs)
print(new_Mat)

We symbolically represent a message m as a vector $(m_0, m_1, m_2, \cdots, m_{15})$ that are elements of $\mathbb F_{2^8}$ and run it through our toy AES encryption to show that it can just be represented as $E(x) = A*x$ where

$$x = (m_0, m_1, m_2, \cdots, m_{15})\\
A = \begin{bmatrix}
2 & 0 & 0 & 0 & 3 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 1 & 0 & 0 & 0\\
0 & 1 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 2 & 0 & 0 & 0 & 3 & 0 & 0\\
0 & 0 & 2 & 0 & 0 & 0 & 3 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 1 & 0\\
0 & 0 & 0 & 1 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 2 & 0 & 0 & 0 & 3\\
1 & 0 & 0 & 0 & 2 & 0 & 0 & 0 & 3 & 0 & 0 & 0 & 1 & 0 & 0 & 0\\
0 & 3 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 2 & 0 & 0\\
0 & 0 & 1 & 0 & 0 & 0 & 2 & 0 & 0 & 0 & 3 & 0 & 0 & 0 & 1 & 0\\
0 & 0 & 0 & 3 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 2\\
1 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 2 & 0 & 0 & 0 & 3 & 0 & 0 & 0\\
0 & 2 & 0 & 0 & 0 & 3 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 1 & 0 & 0\\
0 & 0 & 1 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 2 & 0 & 0 & 0 & 3 & 0\\
0 & 0 & 0 & 2 & 0 & 0 & 0 & 3 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 1\\
3 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 2 & 0 & 0 & 0\\
0 & 1 & 0 & 0 & 0 & 2 & 0 & 0 & 0 & 3 & 0 & 0 & 0 & 1 & 0 & 0\\
0 & 0 & 3 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 2 & 0\\
0 & 0 & 0 & 1 & 0 & 0 & 0 & 2 & 0 & 0 & 0 & 3 & 0 & 0 & 0 & 1\\
\end{bmatrix}$$

Note that the entries in $A$ and $x$ are elements of $\mathbb F_{2^8}$, and not $\mathbb F_{2}$, but the matrix $A$ and $x$ can easily be modified to reflect this. Just involves using a simple isomorphism.

### Isomorphism

As a simple example, look at the values in the column vector $A*x$. Each term is of the form
$$a_i*m_i$$

Where $a_i$ and $m_i$ are elements of $\mathbb F_{2^8}$. There is an isomorphism from $\mathbb F_{2^8}$ to $\mathbb F_{2}^{8}$ where we can represent all $256$ elements of $\mathbb F_{2^8}$ as a vector of length 8 where the entries are elements of $\mathbb F_2$, so $0$ or $1$.

This is essentially converting every integer from $0$ to $255$ into a vector of its bits of length 8.

As a example, in AES, we represent $\mathbb F_{2^8}$ as $\mathbb F[x]/(x^8+x^4+x^3+x+1)$, and $x + 1$ would be $3$ represented as an integer, or $(1, 1, 0, 0, 0, 0, 0, 0)$ as a vector.

So using this simple isomorphism, $a_i*m_i$ can be represented as the "product" of 2 vectors of length 8. However, this multiplication operator isn't the same operator we would use to multiply to 2 regular vectors.

We can define this multiplication operator as follows:
$$\begin{align}
a_i*m_i &= \begin{pmatrix}
a_{i, 0} \\
a_{i, 1} \\
\vdots \\
a_{i, 7}
\end{pmatrix}*
\begin{pmatrix}
m_{i, 0} \\
m_{i, 1} \\
\vdots \\
m_{i, 7}
\end{pmatrix}\\
&= \begin{bmatrix}
b_{0} & b_{8} &\cdots & b_{249}\\
b_{1} & b_{9} & \cdots & b_{250}\\
\vdots & \vdots & \ddots & \vdots\\
b_{7} & b_{15} & \cdots & b_{255}
\end{bmatrix}
\begin{pmatrix}
m_{i, 0} \\
m_{i, 1} \\
\vdots \\
m_{i, 7}
\end{pmatrix}\\
&= \begin{pmatrix}
e_{i, 0} \\
e_{i, 1} \\
\vdots \\
e_{i, 7}
\end{pmatrix}\\
&= e_i
\end{align}$$

Because $e_i$ is an element of $\mathbb F_{2^8}$, or a vector of length 8 using our isomorphism, we need to correctly define an $8 \times 8$ matrix so we can use standard matrix multiplication as our "product" operator. Again, the elements of this matrix are elements of $\mathbb F_2$.

To do this, we can just make $(m_{i, 0}, m_{i, 1}, \cdots, m_{i, 7})$ one of the 8 unit vectors, where only one of the entries is $1$. This makes the product one of the columns of the matrix, which we can set equal to the $e_i$ value as a vector.

Now that we have a way of representing $a_i*m_i$ as the product of an $8 \times 8$ matrix with a column vector of length 8, we can take our $A$ matrix, and convert every $a_i$ value in the matrix to an $8 \times 8$ matrix, which will turn $A$ from a $16 \times 16$ to $128 \times 128$ matrix.

Note that if $a_i$ is $0$, then the matrix is just all $0$'s.

The code below accomplishes this.

new_Mat2 = Matrix(GF(2), 128, 128)
for i in range(16):
    for j in range(16):
        coeffs = []
        for k in range(8):
            x = T.fetch_int(2**k)*T.fetch_int(new_Mat[i, j])
            #print(x)
            x_vec = list(map(int, bin(int(x.integer_representation()))[2:].zfill(8)[::-1]))
            coeffs.append(x_vec)
        new_Mat2[8*i:8*i+8, 8*j:8*j+8] = Matrix(GF(2), 8, 8, coeffs).transpose()
print(new_Mat2)

## Adding SubBytes and AddRoundKey

Now we're going to make our toy AES implementation exactly like regular AES and add the SubBytes and AddRoundKey steps back into the encryption process.

### Checking for a Affine SBOX

The SubBytes step is the non-linear part of AES, and is what makes it resistent to linear attacks. The standard SBOX of AES was designed with this in mind. In this question, however, the author used a different SBOX, so let's check if it's still non-linear.

It's very easy to check this, we simplify have to check if the following relation holds for all $i, j \in \{0, 1, 2, \cdots, 255\}$

$$\text{SBOX[i} \oplus \text{j} \oplus \text{0]} = \text{SBOX[i]} \oplus \text{SBOX[j]} \oplus \text{SBOX[0]}$$

boxxed = [105, 121, 73, 89, 41, 57, 9, 25, 233, 249, 201, 217, 169, 185, 137, 153, 104, 120, 72, 88, 40, 56, 8, 24, 232, 248, 200, 216, 168, 184, 136, 152, 107, 123, 75, 91, 43, 59, 11, 27, 235, 251, 203, 219, 171, 187, 139, 155, 106, 122, 74, 90, 42, 58, 10, 26, 234, 250, 202, 218, 170, 186, 138, 154, 109, 125, 77, 93, 45, 61, 13, 29, 237, 253, 205, 221, 173, 189, 141, 157, 108, 124, 76, 92, 44, 60, 12, 28, 236, 252, 204, 220, 172, 188, 140, 156, 111, 127, 79, 95, 47, 63, 15, 31, 239, 255, 207, 223, 175, 191, 143, 159, 110, 126, 78, 94, 46, 62, 14, 30, 238, 254, 206, 222, 174, 190, 142, 158, 97, 113, 65, 81, 33, 49, 1, 17, 225, 241, 193, 209, 161, 177, 129, 145, 96, 112, 64, 80, 32, 48, 0, 16, 224, 240, 192, 208, 160, 176, 128, 144, 99, 115, 67, 83, 35, 51, 3, 19, 227, 243, 195, 211, 163, 179, 131, 147, 98, 114, 66, 82, 34, 50, 2, 18, 226, 242, 194, 210, 162, 178, 130, 146, 101, 117, 69, 85, 37, 53, 5, 21, 229, 245, 197, 213, 165, 181, 133, 149, 100, 116, 68, 84, 36, 52, 4, 20, 228, 244, 196, 212, 164, 180, 132, 148, 103, 119, 71, 87, 39, 55, 7, 23, 231, 247, 199, 215, 167, 183, 135, 151, 102, 118, 70, 86, 38, 54, 6, 22, 230, 246, 198, 214, 166, 182, 134, 150]
zero = boxxed[0]
for i in range(256):
    for j in range(256):
        if boxxed[i ^^ j] != boxxed[i] ^^ boxxed[j] ^^ zero:
            print("nonlinear")

We don't see "nonlinear" printed out, so this proves that the entire SBOX is affine. Which means that the SBOX can be modeled as the affine function $S(x) = A*x + b$.

We can recover $b$ as $S(0)$ and each column of $A$ can be determined from $S(2^i) + b$, by encrypting a unit vector.

It's important to note here that our SBOX $S(x)$ is a function where $x, a, b \in \mathbb F_{2}^8$

b_subbytes = vector(GF(2), map(int, bin(boxxed[0])[2:].zfill(8)))
A_subbytes = matrix(GF(2), 8, 8)
for i in range(8):
    A_subbytes[7 - i, :] = vector(GF(2), map(int, bin(boxxed[2**i] ^^ boxxed[0])[2:].zfill(8)))

print(b_subbytes)
print(A_subbytes)
    
for i in range(256):
    tmpp = A_subbytes*vector(GF(2), map(int, bin(i)[2:].zfill(8))) + b_subbytes
    assert (int(''.join([str(int(x)) for x in list(tmpp)]), 2) == boxxed[i])

### Creating a Symbolic AES

Now that we can represent the SubBytes step using a vector space in $\mathbb F_2$, we can rewrite the entire AES process to take in a vector of length 128 (where are all of the entires are elements of $\mathbb F_2$) and output a vector of same length.

Standard AES operates off of a $4 \times 4$ matrix of bytes, like below, where our input is 16 bytes: $m_0, m_1, \cdots, m_{15}$.

$$
\begin{bmatrix}
m_{0} & m_{4} & m_{8} & m_{12}\\
m_{1} & m_{5} & m_{9} & m_{13}\\
m_{2} & m_{6} & m_{10} & m_{14}\\
m_{3} & m_{7} & m_{11} & m_{15}
\end{bmatrix}
$$
We can modify this to go from a $4 \times 4$ matrix to a vector of length 128 by taking the bits of every number, so our vector looks like this (using a byte, bit ordering):

$$(m_{0, 7}, m_{0, 6}, \cdots, m_{0, 0}, m_{1, 7}, m_{1, 6}, \cdots, m_{16, 7}, m_{16, 6}, \cdots, m_{16, 0})$$

The round key used in the AddRoundKey step is also 16 bytes, so we can apply the same transformation above.

ShiftRows operates on the byte level, so we just group 8 of the entries of the vector as a "byte" and apply the same transformation.

MixColumns is a little more complicated. We can fix the MixColumns step to change the $4 \times 4$ matrix with entries $\in \mathbb F_2^8$ into a $32 \times 32$ matrix with entries $\in \mathbb F_2$ using the isomorphism we explained earlier.

Below code is an AES class that accomplishes the above.

import os

m_vars = []
for i in range(16):
    for j in range(7, -1, -1):
        m_vars.append(f'm{i}{j}')

PR = PolynomialRing(GF(2), m_vars)
P.<x> = PolynomialRing(GF(2))
T.<z> = GF(2^8, modulus=x^8 + x^4 + x^3 + x + 1)
Mgens = PR.gens()

xtime = lambda a: (((a << 1) ^^ 0x1B) & 0xFF) if (a & 0x80) else (a << 1)


Rcon = (
    0x00, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40,
    0x80, 0x1B, 0x36, 0x6C, 0xD8, 0xAB, 0x4D, 0x9A,
    0x2F, 0x5E, 0xBC, 0x63, 0xC6, 0x97, 0x35, 0x6A,
    0xD4, 0xB3, 0x7D, 0xFA, 0xEF, 0xC5, 0x91, 0x39,
)

def text2matrix(text):
    text = int(text.hex(), 16)
    matrix = []
    for i in range(16):
        byte = (text >> (8 * (15 - i))) & 0xFF
        if i % 4 == 0:
            matrix.append([byte])
        else:
            matrix[i // 4].append(byte)
    return matrix

class AES_symbolic:
    def __init__(self, master_key, Sbox):
        self.Sbox = Sbox
        self.change_key(master_key)

    def change_key(self, master_key):
        self.round_keys = text2matrix(master_key)
        for i in range(4, 4 * 11):
            self.round_keys.append([])
            if i % 4 == 0:
                byte = self.round_keys[i - 4][0]        \
                     ^^ self.Sbox[self.round_keys[i - 1][1]]  \
                     ^^ Rcon[i // 4]
                self.round_keys[i].append(byte)

                for j in range(1, 4):
                    byte = self.round_keys[i - 4][j]    \
                         ^^ self.Sbox[self.round_keys[i - 1][(j + 1) % 4]]
                    self.round_keys[i].append(byte)
            else:
                for j in range(4):
                    byte = self.round_keys[i - 4][j]    \
                         ^^ self.round_keys[i - 1][j]
                    self.round_keys[i].append(byte)
    
    def encrypt(self, plaintext):
        self.plain_state = plaintext
        
        self.plain_state = self.__add_round_key(self.plain_state, self.round_keys[:4])
        
        for i in range(1, 10):
            self.plain_state = self.__sub_bytes(self.plain_state)
            self.plain_state = self.__shift_rows(self.plain_state)
            self.plain_state = self.__mix_columns(self.plain_state)
            self.plain_state = self.__add_round_key(self.plain_state, self.round_keys[4 * i : 4 * (i + 1)])
        
        self.plain_state = self.__sub_bytes(self.plain_state)
        self.plain_state = self.__shift_rows(self.plain_state)
        self.plain_state = self.__add_round_key(self.plain_state, self.round_keys[40:])

        return self.plain_state
    
    def __add_round_key(self, M, k):
        k_vec = []
        for i in range(4):
            for j in range(4):
                k_vec.append(k[i][j])
        bit_string = ''.join(bin(k_val)[2:].zfill(8) for k_val in k_vec)
        K = vector(GF(2), map(int, bit_string))
        return M + K
    
    def __sub_bytes(self, M):
        MM = list(M)
        new_M = []
        for i in range(16):
            tmp = A_subbytes*vector(PR, MM[8*i:8*(i + 1)]) + b_subbytes
            new_M.extend(list(tmp))
        return vector(PR, new_M)
    

    def __shift_rows(self, M):
        MM = list(M)
        ss = [MM[8*i:8*(i + 1)] for i in range(16)]
        s = [[ss[4*i + j] for j in range(4)] for i in range(4)]
        s[0][1], s[1][1], s[2][1], s[3][1] = s[1][1], s[2][1], s[3][1], s[0][1]
        s[0][2], s[1][2], s[2][2], s[3][2] = s[2][2], s[3][2], s[0][2], s[1][2]
        s[0][3], s[1][3], s[2][3], s[3][3] = s[3][3], s[0][3], s[1][3], s[2][3]
        m_new = []
        for i in range(4):
            for j in range(4):
                m_new.extend(s[i][j])
        return vector(PR, m_new)


    def __mix_columns(self, M):
        MM = list(M)
        ss = [MM[32*i:32*(i + 1)] for i in range(4)]        
        S = Matrix(T, [
            [T.fetch_int(2), T.fetch_int(3), T.fetch_int(1), T.fetch_int(1)],
            [T.fetch_int(1), T.fetch_int(2), T.fetch_int(3), T.fetch_int(1)],
            [T.fetch_int(1), T.fetch_int(1), T.fetch_int(2), T.fetch_int(3)],
            [T.fetch_int(3), T.fetch_int(1), T.fetch_int(1), T.fetch_int(2)],
        ])
        new_Mat2 = Matrix(GF(2), 32, 32)
        for i in range(4):
            for j in range(4):
                coeffs = []
                for k in range(8):
                    x = T.fetch_int(2**k)*S[i, j]
                    x_vec = list(map(int, bin(int(x.integer_representation()))[2:].zfill(8)))
                    coeffs.append(x_vec)
                new_Mat2[8*i:8*i+8, 8*j:8*j+8] = Matrix(GF(2), 8, 8, coeffs[::-1]).transpose()
        tmp2 = new_Mat2*(Matrix(PR, 4, 32, ss).transpose())
        m_new = []
        for col in tmp2.columns():
            m_new.extend(list(col))
        return vector(PR, m_new)

key = os.urandom(16)
cipher_symbolic = AES_symbolic(key, Sbox=boxxed)
plaintext = vector(PR, Mgens)
print(plaintext)
print("------------------------")
enc = cipher_symbolic.encrypt(plaintext)
print(enc)
print("------------------------")

A_aes = Matrix(GF(2), 128, 128)
b_aes = []
for i in range(128):
    b_aes.append(GF(2)(enc[i].constant_coefficient()))
    coeffs = []
    for k in range(128):
        coeffs.append(GF(2)(enc[i].coefficient(Mgens[k])))
    A_aes[i, :] = vector(GF(2), coeffs)
b_aes = vector(GF(2), b_aes)
print(A_aes)
print("------------------------")
print(b_aes)

The `enc` vector (not printed for clarity) is length 128 and its entries contain some linear combination of the variables we defined ($m_{0, 7}, m_{0, 6}, \cdots, m_{16, 7}, m_{16, 6}, \cdots, m_{16, 0}$) and a constant.

Similarly to what we did in the beginning with a $16 \times 16$ matrix, we can extract a $128 \times 128$ matrix $A$ (`A_aes` in the code above) and a constant vector $b$ (`b_aes`) from `enc`.

The $128 \times 128$ matrix represents the part of `enc` that depends only on the plaintext.

The constant vector represents the "affine" part of the AES encryption process, that's solely dependent on the SubBytes and AddRoundKey step.

This is because these are the only 2 steps that introduce a constant value into the encryption algorithm, while the other steps are merely multiplying our input vector by something.

As a small note, specifically the "affine" part of the SubBytes step directly affects the constant vector, while the "linear" part affects the $128 \times 128$ matrix.

### Putting it all together

So essentially, we have shown that with an affine SBOX, we can represent the AES encryption process as $$E(x) = A*x + b$$

where $x, b \in \mathbb F_2^{128}$, and the $A$ matrix is the linear part of AES (which is ShiftRows, MixColumns, and the linear part of SubBytes), while $b$ is the affine part of AES (which is AddRoundKey and the affine part of SubBytes).

Note here that $A$ isn't dependent on the key at all, and thus we can compute it using any key we want, since that only affects the constant vector.

To prove this we will take the original code, generate a bunch of random strings, and compare the encryptions using the original code and our new encryption function to ensure they're the same.

import aeschall
from Crypto.Util.number import *

def bytes_to_vec(s):
    return vector(GF(2), map(int, bin(bytes_to_long(s))[2:].zfill(128)))

def vec_to_bytes(v):
    return long_to_bytes(int(''.join(map(str, list(v))), 2)).rjust(16, b'\x00')

cipher = aeschall.AES(key, Sbox=boxxed)
for _ in range(10000):
    test_msg = os.urandom(16)
    enc1 = cipher.encrypt(test_msg)
    enc2 = A_aes*bytes_to_vec(test_msg) + b_aes
    assert enc1 == vec_to_bytes(enc2)

## Recovering the Flag

We are given:

1. Ciphertext of the flag
2. Plaintext of the flag is: "Here is your flag: " + flag

We already have the $A$ matrix, described above, and we can use any key to recover that.

To get $b$, which is dependent on the key and other constants, we need a (plaintext, ciphertext) pair.

Luckily the problem is using ECB mode, so we take the first 16 bytes of known plaintext and 16 bytes of ciphertext as our pair to recover $b$.

Once we recover $A$ and $b$ in $E(x) = A*x + b$, we just need to apply the inverse operation, decryption, to recover $x$, or the flag: $D(x) = A^{-1}*(x + b)$, which is correct since $D(E(x)) = D(A*x + b) = x$.

Code below.

flag_ct = bytes.fromhex("725af38e9584f694638a7323e44749c5ba1e175e61f1bd7cf356da50e7c182cf7ed5ea6e12294f697f3b59b125a3940bc86ca5cfad39b4da4be547dcafbbb17b")
known_pt = b"Here is your flag: "[:16]
known_ct = flag_ct[:16]

b_flag = bytes_to_vec(known_ct) - A_aes*bytes_to_vec(known_pt)

A_aesinv = A_aes.inverse()

flag = b''
for i in range(1, len(flag_ct)//16):
    flag += vec_to_bytes(A_aesinv*(bytes_to_vec(flag_ct[16*i:16*(i+1)]) + b_flag))

print(flag)

## Flag

sun{a3$_r34lly_n33ds_sub5tituti0n!}