# Help with breaking AES reduced to 2 rounds

I was studying Rijndael AES proposal, and its cryptanalysis. I started with writing code to break AES reduced to 2 rounds. I am using 128 bit key size and 128 bit block size. 2 round AES means the following sequence of steps:

$AR \to BS \to SR \to MC \to AR \to BS \to SR \to AR$

where
BS = ByteSub
SR = ShiftRows
MC = MixedColumn

One of the way of breaking the algorithm can be generating inputs following a particular pattern and analyzing the output. Following are patterns I consider:

$P:$ All values are distinct
$Z:$ Sum of all values in zero
$X:$ Any other pattern

I give 256 blocks of input to the algorithm at a time. The blocks follow the following pattern:

$$\begin{matrix} P & P & P & P \\ P & P & P & P \\ P & P & P & P \\ P & P & P & P \\ \end{matrix}$$

i.e. all the elements at same position in 256 blocks are distinct. I will denote this matrix as $[P]$ for simplicity.

After applying the above AES steps following is the fate of these matrices:

$[P] -AR \to [P] -BS \to [P] -SR \to [P] -MC \to [Z] -AR \to [Z] -BS,SR,AR \to [X]$

The last $BS$ operation destroys the structure of the matrix.

We can consider the last 4 steps $AR, BS, SR, AR$ to be a map (say $f$) from 256 inputs at each matrix position to outputs and try to obtain this map using the fact that the input has pattern $Z$. There will be 16 such maps. We can solve it in following way:

Lets consider we want to find the map $f_{00}$ i.e. for the first element in the matrix. We know what all outputs we are getting for the 256 inputs we have taken. Lets say they are:

$u_i, 0 \le i < 256$

Lets denote:
$x_{u_i} = f_{00}^{-1}(u_i)$

Since the input has pattern $Z$, we can write the following equation:

$$\sum_{i=0}^{255} x_{u_i} = 0$$

This is a liner equation in 256 variables. We can take another set of 256 blocks of input of pattern $[P]$ and generate another equation. We can generate 256 of these equations and solve them to find $f_{00}$. Same we can do for other positions as well. Once we get all the maps. We can easily get the key of encryption.

However my problem is this: No matter how many equations I generate, I always get at max 247 equations that are linearly independent. This means they don't have a unique solution. But according to me they should get a unique solution. Where am I going wrong? Why only 247 equations?

You may be overthinking this a bit. Just pick any byte in the ciphertext, and for each possible value of the last RoundKey in that byte, do the following: decrypt just that single byte of the 256 ciphertexts (i.e. compute $AR-SR^{-1}-BS^{-1}$ for that single selected byte using the candidate RoundKey byte in the $AR$ step), and then check to see if the corresponding byte in the State at the beginning of Round 2 all sum to zero. For the correct RoundKey value, the byte will sum to zero. Incorrect RoundKey values will likely not sum to zero. In that manner, you can sieve out most of the incorrect values for each of the RoundKey bytes for the last RoundKey, and can then exhaustively check the remaining possible Keys (i.e. all combinations of RoundKey byte values that were not sieved out) by using the Key Schedule in reverse and checking to see if a candidate Key correctly decrypts all 256 ciphertexts.

Edit to add: By the way, there are more efficient attacks on 2 round AES than the one you outline (which needs 256 chosen plaintexts). For example, here is one that only needs a handful of chosen plaintexts:

Request $N$ chosen plaintexts that differ from each other only in four bytes, such that if any pair of plaintexts is xored together you get a result like so: $$\begin{matrix} W & 0 & 0 & 0 \\ X & 0 & 0 & 0 \\ Y & 0 & 0 & 0 \\ Z & 0 & 0 & 0 \\ \end{matrix}$$ Where $0$ denotes a zero byte, and $W, X, Y$ and $Z$ denote the 'active' bytes with some non-zero difference. With $N$ such chosen plaintexts, there will be ${N \choose 2}$ pairs of plaintexts that have differentially 'active' bytes in the above pattern. This efficient 'packing' of pairs will enable us to attack 2 round AES with only a handful of plaintexts (a dozen is likely to be more than enough).

Now observe from the difference distribution table of the AES S-box that for any given input difference, about half of the possible output differences have a probability of zero (they wont occur given that input difference). This will enable us to sieve out the incorrect keys very efficiently using an impossible differential attack.

Pick a pair of plaintexts and focus on each active byte, one at a time. From the pair of plaintexts we are given an input difference for the focus active byte, and we can use the S-box difference distribution table to figure out the impossible outputs that cannot be attained from that input difference.

For each 'impossible' output difference $I$ for the focus byte after the first SubBytes step, apply ShiftRows and MixColumns to $I$ (notice that for any of the active bytes, after ShiftRows the active byte is mixed with inactive 'zero' bytes). This will give us a set of impossible differences for four different bytes going into the second round (so from all four focus bytes we can derive a set of impossible differences for all 16 bytes of the State for our pair of plaintexts).

Now select any of the corresponding bytes of the ciphertext pair, and decipher that byte under all possible 256 values for the last RoundKey byte (i.e. given some guess for the last RoundKey byte, compute $AR-SR^{-1}-BS^{-1}$ for both ciphertext bytes). For each guess of the RoundKey byte, we can xor the 'deciphered' byte together to get a difference in that byte just before the 2nd round. We can then compare the 'deciphered' difference under that RoundKey guess against the 'impossible' differences that we computed from the plaintext side. If a RoundKey byte guess leads to one of the impossible differences, then we can rule that guess out as not being correct.

We should be able to rule out about a quarter of the possible RoundKey values for any given byte with each plaintext pair, and for each subsequent pair that we test we can be more efficient in the key-guessing phase by not bothering to test the RoundKey values that we have already ruled out. In this way, we can sieve out the incorrect RoundKeys very efficiently.

If you actually let a single byte take on all values, and let the rest all be the same constant value, the attack should work. This so-called square attack was first tried on the cipher Square by its designers, who were Daemen, Knudsen and Rijmen. In the way you are applying it, it seems like you're getting stuck in some subspace. Informally, the AES is quite uniform in how it processes input bytes and you are presenting it with an "equidistributed" input set when all bytes take on all possible values, instead of choosing a "distinguished" byte, as in the square attack.