Building AES-128 from the ground up with python
In my Individual Development Plan as a security professional at Wavy Global, I have a goal to become a pro at cryptography. Someone that knows it well enough to identify problematic setups and come up with PoCs that breaks those setups.
My journey starts with the cryptopals challenges, they are designed to set you up to speed with cryptography and teach you some attacks along the way. One of the most interesting things I had to do is to re implement AES on ECB mode from the ground up. Since there isn't many articles about how to do it in python, I decided to make one.
First of all, how does it work
AES-128 is a block cypher and as the name says, it operates on blocks of 128 bits (16 bytes). Another important notion of AES is that it treats the 16 byte blocks of 4 bytes by 4 bytes. This means that at any point from now on, we must be able to imagine any sequence of 16 bytes in the following disposition:
C0 C1 C3 C4 ╔═══════╦═══════╦═══════╦═══════╗
R0 ║ B0 ║ B1 ║ B2 ║ B3 ║
╠═══════╬═══════╬═══════╬═══════╣
[B0,B1,B2,B3,B4,B5,B6, R1 ║ B4 ║ B5 ║ B6 ║ B7 ║
B7,B8,B9,B10,B11,B12, = ╠═══════╬═══════╬═══════╬═══════╣
B13,B14,B15] R2 ║ B8 ║ B9 ║ B10 ║ B11 ║
╠═══════╬═══════╬═══════╬═══════╣
R3 ║ B12 ║ B13 ║ B14 ║ B15 ║
╚═══════╩═══════╩═══════╩═══════╝
This already let us draw some code
The algorithm go through multiple rounds of substitution and permutation for each block, then concatenate everything. There are multiple modes of operation (you can look at them all here), in this article we are going to focus on the ECB mode (the simplest one).
To implement this “Encrypt” black box we need to understand two core concepts that lives inside this box, the XOR, and the S-BOX.
Good Old XOR
You probably already know this, but it is always good to talk about it. XOR is an operation between two bit and it follows the following truth table
╔═════╦═════╦═════╗
║ X ║ Y ║ Z ║
╠═════╬═════╬═════╣
║ 0 ║ 0 ║ 0 ║
╠═════╬═════╬═════╣
║ 1 ║ 0 ║ 1 ║ X xor Y = Z
╠═════╬═════╬═════╣
║ 0 ║ 1 ║ 1 ║
╠═════╬═════╬═════╣
║ 1 ║ 1 ║ 0 ║
╚═════╩═════╩═════╝We can expand this concept to bytes (you will hear this as a bit wise xor) like this
This can be achieved using the ^ operator , which is the bit-wise xor operator in python. A code in python that xor two numbers would be like this:
a = 1
b = 2
c = a ^ b
print(c)
# 3It is also important to notice that to reverse a xor you just need to apply it again, because
a XOR \x00 = a
b XOR b = \x00then
a XOR b = cand
c XOR b = a
Cryptography S-Box
S-Box are lookup tables for substitution, let me give you a simple example of an identity S-Box
This identity S-Box works for pairs of bit, then it consider the leftmost bit as Rows and the rightmost bit as Columns. Lets try the pair b’10', the leftmost bit is the Rows (in this case 1), and the rightmost is Columns (in this case 0), now looking up into the S-Box we get b’10'.
A good S-Box have to attend some cryptographic criteria, such as size, non linearity, and a be well distributed. A bad S-Box on the other hand can weaken a lot an encryption (you can learn more about sbox design and differential cryptanalysis in this article). AES uses a S-Box called the Rijndael S-box, and since AES is a symmetric encryption algorithm there is also a Reverse Rijndael S-Box for decryption.
Finally we are good to go
Now that we have basic knowledge around XOR and S-Boxes we can break apart the inner of AES
The AES consist of four basic operations that are repeated over N rounds. These three operations are ADDING, SUBSTITUTING, SHIFTING, and MIXING.
This cycle of ADD, SUBSTITUTE, SHIFT, and MIX will repeat for 9 times for 128 bit keys, 11 times for 192 bit keys, and 13 for 256 keys. There is also an initial and final round that we will cover latter.
Implementing the S-BOX
Since KEY_EXPANSION and SUBSTITUTE uses the S-Box, we will begin by implementing it. This is a simple lookup table, so we can just make two matrix and a function that access a position. The way to map a byte to this S-Box is to take the fist most significant nibble as the row, and the least significant nibble as the columns
with the above code, we should get the following results
a = 0x01
print(hex(a))
# 0x1b = lookup(a)
print(hex(b))
# 0x9c = reverse_lookup(b)
print(hex(c))
# 0x1
Implementing the Key Expansion
When performing the multiple rounds AES utilizes an expanded key to improve the security of the algorithm. The expansion is well defined in Wikipedia (I must confess that Wikipedia is the only source that did not confused me :\ ). From the definition we can elaborate a simple code to generate it:
Rotate Rows
At a certain point we will need to put our block in the form of a matrix, rotate the Nth row N times. So the 0th row is rotated 0 times, and so on. This can be achieved with the following code:
The only important thing to notice here is that
row = [1,2,3,4]
rot = rotate_row_left(row, 2)
print(rot)
# [3,4,1,2]rot_back = rotate_row_left(rot, -2)
print(rot_back)
# [1,2,3,4]
So in order to reverse the rotation, we just need to rotate to the same amount multiplied by -1
Mix Columns
To the mix columns step, AES uses a matrix multiplication. The only twists here is that in this matrix multiplication, instead of adding, we are going to XOR the values and, instead of multiplying (not totally trivial).
The matrix that we are going to multiply or column against is
╔ ╗
║ 2 3 1 1 ║
║ 1 2 3 1 ║
║ 1 1 2 3 ║
║ 3 1 1 2 ║
╚ ╝Luckily, there is a better way to perform this multiplication. Since this is a constant matrix, some really cool folks already chewed this down to us (look here)
So the whole operation becomes this in code
To be able to reverse this matrix operation, we would need to multiply by the inverse of this matrix, which is not fun at all. What I’ll do instead is to exploit a cool feature of this particular matrix which is M⁴ = I, let me break that down.
for every matrix M X M^-1 = Ifor this matrix in particular M^4 = Iso for this matrix in particular M X M X M X M = Iso for this matrix in particular M X (M X M X M) = Iso for this matrix in particular M X M X M = M^-1
Funny right? This awesome property allow us to be lazy and do a code like this
a = [212, 212, 212, 213]
mixed = mix_columns(a)
print(mixed)
# [213, 213, 215, 214]unmixed = mix_columns(mixed)
unmixed = mix_columns(unmixed)
unmixed = mix_columns(unmixed)
print(unmixed)
# [212, 212, 212, 213]
Add Sub Key
Add sub key is the easiest part, it is just a xor byte by byte of the array
Since this is a xor, to undo this operation you just need to perform it again.
Initial and Final Rounds
The initial and final rounds are just simplification of a general round.
The first round is just a Add Sub Key with the first 16 bytes of the key (A.K.A the key you provided)
The final round is a sequence of
- Sub Bytes (sbox)
- Shift Row
- Add Sub Key
Now that we have implemented all functions, we just need to put them in order to implement encryption and decryption .
Encryption
Decryption
Conclusion and final thoughts
Now this was a fun ride. Implementing the algorithm gave me a better understanding of how this all works and certainly moved me in the direction of becoming a pro in cryptography. Of course the process itself was not as easy as it looks like, I’ve spent some nights learning AES and trying to chew it in a way that is easy enough to put on a medium article.
One thing that I need to point out is to NEVER use this or any hobby made cryptography code in production. Always use well known and tested libraries, otherwise you are prone to side channel attacks.
I hope you enjoyed it as much as I did. If you have any questions, just leave it down here
Some more info
[14:40 GMT–3 16/04/2021] — I forgot to include one piece of code. The extract key for round is defined in the following way:
def extract_key_for_round(expanded_key, round):
return [row[round*4: round*4 + 4] for row in expanded_key]It takes the expanded key and extracts the segment needed for the current round
