[Python] How can I correctly implement AES-256-GCM from scratch in Python?

I am sorry if this is the wrong way to do it, but I am using Windows 11 and I plan to implement the code in C++ and as you might know building C++ libraries on Windows is very complicated and so I want to make sure it uses the least amount of dependencies possible.

I am behind the Great Firewall of China, and I am a liberal, I bypass the GFW to access the international internet, for more context see this.

I have decided to implement a load-balancing ShadowSocks5 proxy in C++ from scratch, it should catch all network requests and send them through a ShadowSocks5 tunnel via Ethernet and Wi-Fi alternatively to remote proxy servers, and it should receive traffic from the remote proxy servers, and it should have its own DNS resolving solution.

This is a programming challenge, a learning project, and a practical project all in one.

I decided to start from the easiest problem, the encryption I need to use is AES-256-GCM. GCM stands for Galois/Counter_Mode, I haven't figured out how to implement it from reading the Wikipedia article. But the Wikipedia article on Advanced Encryption Standard is very helpful and it is the one of the primary reference I used in implementing this. Another Wikipedia article I referenced is AES key schedule. I got the values for SBOX and RSBOX from this article, the algorithm for the inverse of Galois multiplication is based on this Stack Overflow answer.

Now, here is the implementation, I wrote it all by myself, I spent two days on it:

import json


with open("D:/AES_256.json", "r") as f:
AES_256 = json.load(f)

MAX_256 = (1 << 256) - 1
SBOX = AES_256["SBOX"]
RSBOX = AES_256["RSBOX"]
RCON = AES_256["RCON"]
OCTUPLE = (
(0, 4),
(4, 8),
(8, 12),
(12, 16),
(16, 20),
(20, 24),
(24, 28),
(28, 32),
)

SEXAGESY = (
(0, 1, 2, 3),
(4, 5, 6, 7),
(8, 9, 10, 11),
(12, 13, 14, 15),
(16, 17, 18, 19),
(20, 21, 22, 23),
(24, 25, 26, 27),
(28, 29, 30, 31),
(32, 33, 34, 35),
(36, 37, 38, 39),
(40, 41, 42, 43),
(44, 45, 46, 47),
(48, 49, 50, 51),
(52, 53, 54, 55),
(56, 57, 58, 59),
)

HEXMAT = (0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15)
SHIFT = (0, 1, 2, 3, 5, 6, 7, 4, 10, 11, 8, 9, 15, 12, 13, 14)

COLMAT = (
((2, 0), (1, 4), (1, 8), (3, 12)),
((2, 1), (1, 5), (1, 9), (3, 13)),
((2, 2), (1, 6), (1, 10), (3, 14)),
((2, 3), (1, 7), (1, 11), (3, 15)),
((3, 0), (2, 4), (1, 8), (1, 12)),
((3, 1), (2, 5), (1, 9), (1, 13)),
((3, 2), (2, 6), (1, 10), (1, 14)),
((3, 3), (2, 7), (1, 11), (1, 15)),
((1, 0), (3, 4), (2, 8), (1, 12)),
((1, 1), (3, 5), (2, 9), (1, 13)),
((1, 2), (3, 6), (2, 10), (1, 14)),
((1, 3), (3, 7), (2, 11), (1, 15)),
((1, 0), (1, 4), (3, 8), (2, 12)),
((1, 1), (1, 5), (3, 9), (2, 13)),
((1, 2), (1, 6), (3, 10), (2, 14)),
((1, 3), (1, 7), (3, 11), (2, 15)),
)


def add_round_key(state: list, key: list) -> list:
return [a ^ b for a, b in zip(state, key)]


def state_matrix(data: list) -> list:
return [data for i in HEXMAT]


def sub_bytes(state: list) -> list:
return [SBOX for i in state]


def shift_rows(state: list) -> list:
return [state for i in SHIFT]


def rot8(byte: int, x: int) -> int:
x &= 7
return (byte << x | byte >> (8 - x)) & 0xFF


def quadword(quad: bytes) -> int:
a, b, c, d = quad
return (a << 24) | (b << 16) | (c << 8) | d


def rot_word(word: int) -> int:
return (word << 8 | word >> 24) & 0xFFFFFFFF


def sub_word(word: int) -> int:
return (
(SBOX[(word >> 24) & 0xFF] << 24)
| (SBOX[(word >> 16) & 0xFF] << 16)
| (SBOX[(word >> 8) & 0xFF] << 8)
| SBOX[word & 0xFF]
)


def galois_mult(x: int, y: int) -> int:
p = 0
while x and y:
if y & 1:
p ^= x

if x & 0x80:
x = (x << 1) ^ 0x11B
else:
x <<= 1

y >>= 1

return p


def galois_inv(x: int) -> int:
if not x:
return 0

p = galois_mult(x, x)
x = galois_mult(p, p)
duo = [p, x]
b = 0
for _ in range(11):
duo = galois_mult(duo, duo[1])
b ^= 1

return duo[0]


def galois_div(x: int, y: int) -> int:
if not y:
raise ZeroDivisionError("division by zero")

return 1 if x == y else galois_mult(x, galois_inv(y))


def mix_columns(state: list) -> list:
result = [0] * 16
for e, row in zip(state, COLMAT):
for mult, i in row:
result ^= galois_mult(e, mult)

return state_matrix(result)


def key_matrix(key: list) -> list:
mat = []
for row in SEXAGESY:
line = []
for col in row:
n = key[col]
line.extend([n >> 24 & 0xFF, n >> 16 & 0xFF, n >> 8 & 0xFF, n & 0xFF])

mat.append(state_matrix(line))

return mat


def derive_key(password: bytes) -> list:
keys = [quadword(password[a:b]) for a, b in OCTUPLE]
result = keys.copy()
last = result[7]
for i in range(8, 60):
if not i & 7:
last = sub_word(rot_word(last)) ^ (RCON[i // 8] << 24)
elif i & 7 == 4:
last = sub_word(last)

key = result[i - 8] ^ last
result.append(key)
last = key

return key_matrix(result)


def aes_256_cipher(data: bytes, password: bytes) -> list:
state = add_round_key(state_matrix(data), password[0])
for i in range(1, 14):
state = add_round_key(mix_columns(shift_rows(sub_bytes(state))), password)

return state_matrix(add_round_key(shift_rows(sub_bytes(state)), password[14]))


def get_padded_data(data: bytes | str) -> bytes:
if isinstance(data, str):
data = data.encode("utf8")

if not isinstance(data, bytes):
raise ValueError("argument data must be bytes or str")

return data + b"\x00" * (16 - len(data) % 16)


def get_key(password: bytes | int | str) -> list:
if isinstance(password, int):
if password < 0 or password > MAX_256:
raise ValueError("argument password must be between 0 and 2^256-1")

password = password.to_bytes(32, "big")

if isinstance(password, str):
password = "".join(i for i in password if i.isalnum()).encode("utf8")
if len(password) > 32:
raise ValueError("argument password must be 32 bytes or less")

if not isinstance(password, bytes):
raise ValueError("argument password must be bytes | int | str")

return derive_key(password.rjust(32, b"\x00"))


def ecb_encrypt(data: bytes | str, password: bytes | str) -> str:
data = get_padded_data(data)
key = get_key(password)
blocks = [aes_256_cipher(data[i : i + 16], key) for i in range(0, len(data), 16)]
return "".join(f"{e:02x}" for block in blocks for e in block)


AES_256.json

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


It doesn't raise exceptions, but I have absolutely no idea what I am doing.

This is an example output:

In [428]: ecb_encrypt('6a84867cd77e12ad07ea1be895c53fa3', '0'*32)
Out[428]: 'b981b1853c16fbb6adc7cf4a01c9c57b94a3e5ce608239660c324b01400ebdd5d45a5452d22fed94b7ca9d916ac47736'


I have used this website to check the correctness of the output, with the same key and plaintext, AES-256-ECB mode, the cipher text in hex is:

e07beff38697f04e7adbc971adc2a9135f60746178fcd0f1b3040e4d15c920ad0318e084e1666e699891c78f8aa98960


It isn't what my code outputs as you can see.

Why isn't my code working properly? And how can I correctly implement AES-256-GCM in Python (so that I can implement it in C++ later)?

Continue reading...