Advanced Encryption Standard¶

Links:

Python Implementation¶

The implementation is written for visualizing & demonstrating the steps of AES algorithm, so it doesn't completely focus on efficiency.

Note: The algorithm uses Electronic Code Book (ECB) mode to process longer data for the sake of simplicity in demonstration. However, ECB mode is insecure as it is vulnerable to frequency analysis, so this encryption mode must be avoided.

The symbol "▼" signifies expandable sections. Click to expand/collapse.

from numpy import array, zeros, arange, roll
from numpy.core.fromnumeric import shape
from IPython.core.display import display, HTML
from secrets import token_bytes
from tqdm.notebook import tqdm
from PIL import Image

class AdvancedEncryptionStandard:

    _substitution_box = array(
    [
        0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76,
        0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0,
        0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15,
        0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a, 0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75,
        0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84,
        0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b, 0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf,
        0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85, 0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8,
        0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5, 0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2,
        0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73,
        0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88, 0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb,
        0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79,
        0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08,
        0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a,
        0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e,
        0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf,
        0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16
    ], dtype='B')

    _substitution_box_inv = zeros(256, dtype='B')
    _substitution_box_inv[_substitution_box] = arange(256, dtype='B')

    _Rcon = array([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], dtype='B')
    
    _c = array([[2, 3, 1, 1],
                [1, 2, 3, 1],
                [1, 1, 2, 3],
                [3, 1, 1, 2]], dtype='B')
    
    _c_inv = array([[0x0e, 0x0b, 0x0d, 0x09],
                    [0x09, 0x0e, 0x0b, 0x0d],
                    [0x0d, 0x09, 0x0e, 0x0b],
                    [0x0b, 0x0d, 0x09, 0x0e]], dtype='B')
    
    _key_to_rounds = {16:10, 24:12, 32:14}

    _formatters = {
            'HEX' : lambda x: hex(x)[2:].zfill(2), 
            'BIN' : lambda x: bin(x)[2:].zfill(8), 
            'DEC' : lambda x: str(x).zfill(3)
        }

    def __init__(self):
        self.html_output = ''

    def _get_html(self, body):
        return '''
        <html>
        <head>
        <style>
            *{
                font-size: 15px;
            }
            div.i{
                font-style: italic;
                text-decoration: underline;
                margin-top: 24px;
                margin-top: 5px;
                font-weight: bold;
                font-size: 17px;
            }
            .op{
                vertical-align:middle;
            }
            caption {
                font-size: 10px;
                opacity: 0.6;
            }
            table {
                border-collapse: collapse;
                margin: 3px;
                display: inline;
            }
            th, td {
                border: 1px solid;
                padding: 3px;
                margin: 0;
                text-align: center;
            }
            summary *{
                display: inline-block;
            }
            details{
                padding-left: 16px;
                margin-top: 5px;
            }
        </style>
        </head>
        <body>'''+body+'</body></html>'
    
    def _html_key_table(self, key):
        return f'''
            <h5>Given Key</h5>
            <table>
                <tr>
                    <th>ASCII</th>
                    {''.join('<td>'+chr(k)+'</td>' for k in key)}
                </tr>
                <tr>
                    <th>HEX</th>
                    {''.join('<td>'+self._formatters['HEX'](k)+'</td>' for k in key)}
                </tr>
                <tr>
                    <th>BIN</th>
                    {''.join('<td>'+self._formatters['BIN'](k)+'</td>' for k in key)}
                </tr>
                <tr>
                    <th>DEC</th>
                    {''.join('<td>'+self._formatters['DEC'](k)+'</td>' for k in key)}
                </tr>
            </table>
            '''

    def arr_to_html(self, arr, caption, fmttr = lambda x:x):
        html=f'<table><caption>{caption}</caption>'
        for row in arr:
            html+='<tr>'
            for val in row:
                html+=f'<td>{fmttr(val)}</td>'
            html+='</tr>'
        return html+'</table>'

    def state_matrix_to_html(self, data):
        html = ''
        for c, f in self._formatters.items():
            html+=self.arr_to_html(data, c, f)
        return html
        
    def pad(self, data):
        p = 16 - len(data)%16
        if p==16: return data
        return data+p.to_bytes(p, 'big')

    def unpad(self, data):
        p = int(data[-1])
        if p>16: return data
        return data.rstrip(p.to_bytes(p, 'big'))

    def galois_multiply(self, a, b):
        p = 0
        for _ in range(8):
            p^= -(b&1)&a
            mask = -((a>>7)&1)
            a = (((a<<1)%(1<<9))^(0b100011011 & mask))
            b>>=1
        return p

    def key_expansion(self, key):
        rounds = (self._key_to_rounds[len(key)]+1)*4
        w = [None]*rounds
        for i in range(4):
            w[i] = array([key[4*i+j] for j in range(4)], dtype='B')
        for i in range(4, rounds):
            temp = w[i-1].copy()
            if i%4==0:
                temp = roll(temp, -1)
                temp[:4] = self._substitution_box[temp[:4]]
                temp[0] = temp[0]^self._Rcon[i//4-1]
            w[i] = w[i-4]^temp
        w = array(w, dtype='B')

        html = '<div class="i">Key Expansion</div><table><tr><th></th>'
        for i in range(rounds):
            html+=f'<th>w<sub>{i}</sub></th>'
        for i in range(4):
            html+=f'<tr><th>B<sub>{i}</sub></th>'
            for j in range(rounds): 
                html+=f'<td>{hex((w[j][i]))[2:].zfill(2)}</td>'
            html+='</tr>'
        return w, html+'</table>'
    
    def substitute_bytes(self, data, sbox):
        for i in range(4):
            for j in range(4):
                data[i, j] = sbox[data[i, j]]
        self.html_output+='<div class="i">Substitute Bytes</div>'+self.state_matrix_to_html(data)

    def mix_columns(self, data, c):
        for i in range(4):
            column = data[:, i].copy()
            for j in range(4):
                res = 0
                for k in range(4):
                    res^=self.galois_multiply(c[j, k], column[k])
                data[j, i] = res
        self.html_output+='<div class="i">Mix Columns</div>'+self.state_matrix_to_html(data)

    def shift_rows(self, data, inverse=False):
        v = 1 if inverse else -1
        for i in range(len(data)):
            data[i] = roll(data[i], v*i)
        self.html_output+='<div class="i">Shift Rows</div>'+self.state_matrix_to_html(data)

    def add_round_key(self, data, round_key):
        b = 'BIN'
        self.html_output+='<div class="i">Add Round Key</div>'+self.arr_to_html(data, b, self._formatters[b])+'<span class="op">⊕</span>'+self.arr_to_html(round_key.T, b, self._formatters[b])+'<span class="op">=</span>'
        data^=round_key.T
        self.html_output+=self.arr_to_html(data, b, self._formatters[b])+'<p>State Matrix</p>'+self.state_matrix_to_html(data)

    def _encrypt(self, message, round_keys, rounds):
        self.html_output = ''
        data = array(message, dtype='B').reshape(4, 4).T
        self.html_output+='<p>State Matrix</p>'+self.state_matrix_to_html(data)
        self.add_round_key(data, round_keys[0:4])

        for i in range(1, rounds):
            self.html_output+=f'<details><summary>▼ Round {i}</summary>'
            self.substitute_bytes(data, self._substitution_box)
            self.shift_rows(data)
            self.mix_columns(data, self._c)
            self.add_round_key(data, round_keys[4*i:4*(i+1)])
            self.html_output+='</details>'

        self.html_output+=f'<details><summary>▼ Round {i+1}</summary>'
        self.substitute_bytes(data, self._substitution_box)
        self.shift_rows(data)
        self.add_round_key(data, round_keys[-4:])
        self.html_output+='</details>'

        return bytes(data.T.reshape(1, -1))

    def _decrypt(self, data, round_keys, rounds):
        self.html_output = ''

        data = array(data, dtype='B').reshape(4, 4).T

        self.html_output+='<p>State Matrix</p>'+self.state_matrix_to_html(data)

        self.add_round_key(data, round_keys[-4:])

        for i in range(rounds, 1, -1):
            self.html_output+=f'<details><summary>▼ Round {rounds-i+1}</summary>'
            self.shift_rows(data, inverse=True)
            self.substitute_bytes(data, self._substitution_box_inv)
            self.add_round_key(data, round_keys[4*(i-1):4*i])
            self.mix_columns(data, self._c_inv)
            self.html_output+='</details>'

        self.html_output+=f'<details><summary>▼ Round {rounds-i+2}</summary>'
        self.shift_rows(data, inverse=True)
        self.substitute_bytes(data, self._substitution_box_inv)
        self.add_round_key(data, round_keys[:4])
        self.html_output+='</details>'

        return bytes(data.T.reshape(1, -1))

    def encrypt(self, data, key=None, round_keys=None, rounds=None, key_info_expanded = False):
        html_body = f'<p>Given Plain Text: <b>{data}</b></p>'
        data = bytearray(data)
        data = self.pad(data)
        if key is not None:
            html_body+=f'<details {"open" if key_info_expanded else ""}><summary>▼ Key Information</summary>'+self._html_key_table(key)
            round_keys, html = self.key_expansion(key)
            html_body+=html+'</details>'
            rounds = self._key_to_rounds[len(key)]
        cipher_text = bytearray()
        if len(data)==16:
            html_body+='<details open><summary>▼ AES Encryption</summary>'
            cipher_text = self._encrypt(data, round_keys, rounds)
            html_body += self.html_output
        else:
            html_body+='<p>Using Electronic Code Book (ECB) mode, </p><details open><summary>▼ AES Encryption</summary>'
            for i in tqdm(range(0, len(data), 16), desc='Blocks Processed'):
                message = data[i:i+16]
                html_body+=f'<details><summary>▼ Block {i//16 + 1} | Plaintext: <b>{bytes(message)}</b></summary>'
                cipher_text += self._encrypt(message, round_keys, rounds)
                html_body += self.html_output+'</details>'
        cipher_text = bytes(cipher_text)
        html_body+= f'</details><p>Cipher Text: <b>{cipher_text}</b></p>'
        return cipher_text, HTML(self._get_html(html_body))
    
    def decrypt(self, data, key=None, round_keys=None, rounds=None, key_info_expanded = False):
        html_body = f'<p>Given Cipher Text: <b>{data}</b></p>'
        data = bytearray(data)
        if key is not None:
            html_body+=f'<details {"open" if key_info_expanded else ""}><summary>▼ Key Information</summary>'+self._html_key_table(key)
            round_keys, html = self.key_expansion(key)
            html_body+=html+'</details>'
            rounds = self._key_to_rounds[len(key)]
        decipher_text = bytearray()
        if len(data)==16:
            html_body+='<details open><summary>▼ AES Decryption</summary>'
            decipher_text = self._decrypt(data, round_keys, rounds)
            html_body += self.html_output
        else:
            html_body+='<p>Using Electronic Code Book (ECB) mode, </p><details open><summary>▼ AES Decryption</summary>'
            for i in tqdm(range(0, len(data), 16), desc='Blocks Processed'):
                message = data[i:i+16]
                html_body+=f'<details><summary>▼ Block {i//16 + 1} | Ciphertext: <b>{bytes(message)}</b></summary>'
                decipher_text += self._decrypt(message, round_keys, rounds)
                html_body += self.html_output+'</details>'

        decipher_text = bytes(self.unpad(decipher_text))
        html_body+= f'</details><p>Deciphered Text: <b>{decipher_text}</b></p>'
        return decipher_text, HTML(self._get_html(html_body))

key = 'asdfghjklqwertyu'.encode('ascii')
message = 'Hi! I am Samuel.'

aes = AdvancedEncryptionStandard()

print('AES Encryption:')
cipher, html = aes.encrypt(message.encode('ascii'), key, key_info_expanded = True)
display(html)

print('\n\nAES Decryption:')
decipher, html = aes.decrypt(cipher, key)
display(html)

AES Encryption:


AES Decryption:

key = token_bytes(16)
message = 'Advanced Encryption Standard is cool!'

aes = AdvancedEncryptionStandard()

print('AES Encryption:')
cipher, html = aes.encrypt(message.encode('ascii'), key)
display(html)

print('\n\nAES Decryption:')
decipher, html = aes.decrypt(cipher, key)
display(html)

AES Encryption:


AES Decryption:

image = Image.open('shapes.jpeg')
print('Original Image:')
display(image)
img_arr = array(image)
img_shape = img_arr.shape

def to_image(data, shape):
    return Image.fromarray(array(bytearray(data)).reshape(shape))

Original Image:

key = 'This is the key!'.encode('ascii')
message = img_arr.flatten()

aes = AdvancedEncryptionStandard()

print('AES Encryption:')
cipher, html = aes.encrypt(message, key)
display(to_image(cipher, img_shape))

print('\n\nAES Decryption:')
decipher, html = aes.decrypt(cipher, key)
display(to_image(decipher, img_shape))

AES Encryption:


AES Decryption:

# 24-bit key
key = 'abcdefghijklmnopqrstuvwx'.encode('ascii')
message = 'Hi! I am Samuel.'

aes = AdvancedEncryptionStandard()

print('AES Encryption:')
cipher, html = aes.encrypt(message.encode('ascii'), key, key_info_expanded = True)
display(html)

print('\n\nAES Decryption:')
decipher, html = aes.decrypt(cipher, key)
display(html)

AES Encryption:


AES Decryption:

ASCII	a	s	d	f	g	h	j	k	l	q	w	e	r	t	y	u
HEX	61	73	64	66	67	68	6a	6b	6c	71	77	65	72	74	79	75
BIN	01100001	01110011	01100100	01100110	01100111	01101000	01101010	01101011	01101100	01110001	01110111	01100101	01110010	01110100	01111001	01110101
DEC	097	115	100	102	103	104	106	107	108	113	119	101	114	116	121	117

01001000	01001001	00100000	01110101
01101001	00100000	01010011	01100101
00100001	01100001	01100001	01101100
00100000	01101101	01101101	00101110

072	073	032	117
105	032	083	101
033	097	097	108
032	109	109	046

01001000	01001001	00100000	01110101
01101001	00100000	01010011	01100101
00100001	01100001	01100001	01101100
00100000	01101101	01101101	00101110

01100001	01100111	01101100	01110010
01110011	01101000	01110001	01110100
01100100	01101010	01110111	01111001
01100110	01101011	01100101	01110101

	w₀	w₁	w₂	w₃	w₄	w₅	w₆	w₇	w₈	w₉	w₁₀	w₁₁	w₁₂	w₁₃	w₁₄	w₁₅	w₁₆	w₁₇	w₁₈	w₁₉	w₂₀	w₂₁	w₂₂	w₂₃	w₂₄	w₂₅	w₂₆	w₂₇	w₂₈	w₂₉	w₃₀	w₃₁	w₃₂	w₃₃	w₃₄	w₃₅	w₃₆	w₃₇	w₃₈	w₃₉	w₄₀	w₄₁	w₄₂	w₄₃
B₀	61	67	6c	72	f3	94	f8	8a	30	a4	5c	d6	1e	ba	e6	30	6e	d4	32	02	9f	4b	79	7b	35	7e	07	7c	99	e7	e0	9c	6b	8c	6c	f0	d4	58	34	c4	96	ce	fa	3e
B₁	73	68	71	74	c5	ad	dc	a8	9b	36	ea	42	54	62	88	ca	49	2b	a3	69	21	0a	a9	c0	93	99	30	f0	67	fe	ce	3e	2b	d5	1b	25	fe	2b	30	15	5e	75	45	50
B₂	64	6a	77	79	f9	93	e4	9d	b5	26	c2	5f	65	43	81	de	eb	a8	29	f7	48	e0	c9	3e	ad	4d	84	ba	2e	63	e7	5d	6c	0f	e8	b5	15	1a	f2	47	31	2b	d9	9e
B₃	66	6b	65	75	26	4d	28	5d	58	15	3d	60	ae	bb	86	e6	aa	11	97	71	dd	cc	5b	2a	fc	30	6b	41	ec	dc	b7	f6	32	ee	59	af	be	50	09	a6	a2	f2	fb	5d

00101001	00101110	01001100	00000111
00011010	01001000	00100010	00010001
01000101	00001011	00010110	00010101
01000110	00000110	00001000	01011011

10100101	00110001	00101001	11000101
10100010	01010010	10010011	10000010
01101110	00101011	01000111	01011001
01011010	01101111	00110000	00111001

11011001	11001111	11001110	01110111
11110001	10111101	11101011	11010111
00110010	11111110	11000110	01100001
10010011	00101101	01001001	10111101

11110011	10010100	11111000	10001010
11000101	10101101	11011100	10101000
11111001	10010011	11100100	10011101
00100110	01001101	00101000	01011101

00101010	01011011	00110110	11111101
00110100	00010000	00110111	01111111
11001011	01101101	00100010	11111100
10110101	01100000	01100001	11100000

11100101	00111001	00000101	01010100
00011000	11001010	10011010	11010010
00011111	00111100	10010011	10110000
11010101	11010000	11101111	11100001

11100110	10100010	10101000	01010011
00100101	00001000	01001011	11001111
00101010	10111100	10000010	00011110
10110100	11010000	01111001	00011101

00110000	10100100	01011100	11010110
10011011	00110110	11101010	01000010
10110101	00100110	11000010	01011111
01011000	00010101	00111101	01100000

11010110	00000110	11110100	10000101
10111110	00111110	10100001	10001101
10011111	10011010	01000000	01000001
11101100	11000101	01000100	01111101

DEC
041	046	076	007
026	072	034	017
069	011	022	021
070	006	008	091

DEC
165	049	041	197
162	082	147	130
110	043	071	089
090	111	048	057

DEC
165	049	041	197
082	147	130	162
071	089	110	043
057	090	111	048

DEC
217	207	206	119
241	189	235	215
050	254	198	097
147	045	073	189

11110110	01101111	10111111	10010111
10101110	10110010	00110010	01011101
11011011	10111000	00001001	10000011
11001110	10100110	00011011	11111111

11001100	11000101	11111111	01111111
01101101	01011011	11010101	00011000
01001100	00001001	10111110	01111111
01011111	10000111	00001011	10000010

00011110	10111010	11100110	00110000
01010100	01100010	10001000	11001010
01100101	01000011	10000001	11011110
10101110	10111011	10000110	11100110

11010010	01111111	00011001	01001111
00111001	00111001	01011101	11010010
00101001	01001010	00111111	10100001
11110001	00111100	10001101	01100100

10110101	11010010	11010100	10000100
00010010	00010010	01001100	10110101
10100101	11010110	01110101	00110010
10100001	11101011	01011101	01000011

00011111	00101100	00001011	10101100
00000100	10010110	00011001	11110101
01100011	10101010	00111111	00110001
10001111	01011011	00101101	10011000

11000000	01110001	00101011	10010001
11110010	10010000	11010100	11100110
11111011	10101100	01110101	11000111
01110011	00111001	11011000	01000110

00000011	00110001	10100101	01000000
00100010	11100011	11010011	01011001
01110000	10100101	01101011	01010011
00110010	01100110	00010010	01011101

10011100	01111010	11011100	00111011
00000011	11101001	01111010	10011001
00111000	01000101	10100010	01101101
11101111	10101010	01001001	01110111

11011110	11011010	10000110	11100010
01111011	00011110	11011010	11101110
00000111	01101110	00111010	00111100
11011111	10101100	00111011	11110101

01001010	00111001	10010101	00000111
01011001	11101110	11100100	10011101
10110000	00000010	10001001	00001000
10101100	00110110	00111011	01011110

00110101	01111110	00000111	01111100
10010011	10011001	00110000	11110000
10101101	01001101	10000100	10111010
11111100	00110000	01101011	01000001

01111111	01000111	10010010	01111011
11001010	01110111	11010100	01101101
00011101	01001111	00001101	10110010
01010000	00000110	01010000	00011111

11010010	10100000	01001111	00100001
01110100	11110101	01001000	00111100
10100100	10000100	11010111	00110111
01010011	01101111	01010011	11000000

10101100	11100111	00010001	00001001
10000001	00111010	10101111	00001101
11001001	01110011	10010001	10110011
11010100	00100010	10010111	00110101

00110101	00000000	11110001	10010101
11100110	11000100	01100001	00110011
11100111	00010000	01110110	11101110
00111000	11111110	00100000	11000011

01101011	10001100	01101100	11110000
00101011	11010101	00011011	00100101
01101100	00001111	11101000	10110101
00110010	11101110	01011001	10101111

10010110	11001110	11111010	00111110
01011110	01110101	01000101	01010000
00110001	00101011	11011001	10011110
10100010	11110010	11111011	01011101

10100100	11001101	10001100	11110111
11111110	01000000	11100010	00001111
00010111	10100011	00011100	01110110
10000101	11001111	11011111	00000111

00001100	10110110	00011000	10000101
01001010	00011001	10000101	00110110
11111111	01111101	10110001	00100101
00110010	01100011	11101110	11001111