IV. Thuật toán AES - hoàn thiện

1. Xây dựng hàm mã hóa và giải mã

Nhằm thực hiện mã hóa AES chúng ta có $5$ công việc chính: ExpandKey, AddRoundKey, SubBytes, ShiftRows, MixColumns. Đến thời điểm hiện tại chúng ta chỉ còn thiếu việc chuyển thể ExpandKey từ lý thuyết sang lập trình (Do ExpandKey được giới thiệu ở phần $2$ nhưng chưa đủ kiến thức thực hiện). Đối với AES-128, cùng xem lại công thức tính các word trong quá trình mở rộng khóa:

$$W_i = \begin{cases} K_i & \text{nếu } i < 4 \\ W_{i-4} \oplus \operatorname{SubWord}(\operatorname{RotWord}(W_{i-1})) \oplus rcon_{i/4} & \text {nếu } i \ge 4 \text{ và } i \equiv 0 \pmod{4} \\ W_{i-4} \oplus W_{i-1} & \text{Các trường hợp còn lại} \\ \end{cases}$$

Tham khảo hàm expand_key() trong challenge Bringing It All Together

def expand_key(master_key):

    r_con = (
        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,
    )

    # Initialize round keys with raw key material.
    key_columns = bytes2matrix(master_key)
    iteration_size = len(master_key) // 4

    # Each iteration has exactly as many columns as the key material.
    i = 1
    while len(key_columns) < (N_ROUNDS + 1) * 4:
        # Copy previous word.
        word = list(key_columns[-1])

        # Perform schedule_core once every "row".
        if len(key_columns) % iteration_size == 0:
            # Circular shift.
            word.append(word.pop(0))
            # Map to S-BOX.
            word = [s_box[b] for b in word]
            # XOR with first byte of R-CON, since the others bytes of R-CON are 0.
            word[0] ^= r_con[i]
            i += 1

        # XOR with equivalent word from previous iteration.
        word = bytes(i^j for i, j in zip(word, key_columns[-iteration_size]))
        key_columns.append(word)

    # Group key words in 4x4 byte matrices.
    return [key_columns[4*i : 4*(i+1)] for i in range(len(key_columns) // 4)]

Quan sát hàm trên, một số yếu tố cần chú ý:

• Biến r_con khai báo một tuple các $rc_i$ trong Round constants.
• Biến key_columns lưu trữ giá trị secret key chúng ta sẽ thực hiện mở rộng.
• Biến iteration_size chỉ giá trị $N$, trong trường hợp này bằng $4$.
• Chương trình sử dụng vòng lặp while len(key_columns) < (N_ROUNDS + 1) * 4: liên tục kiểm tra điều kiện độ dài khóa mở rộng đã đạt tới $44$ word hay chưa.
• Điều kiện if len(key_columns) % iteration_size == 0: xử lý trường hợp $i \ge 4 \text{ và } i \equiv 0 \pmod{4}$.

Như vậy, chúng ta đã trang bị đầy đủ các hàm cần thiết cho quá trình mã hóa AES-128. Công việc hiện tại trở nên đơn giản hơn bao giờ hết: sắp xếp chúng theo đúng thứ tự.

State table là thành phần trung tâm thay đổi xuyên suốt quá trình mã hóa, chúng ta đặt biến state lưu lại giá trị tại các thời điểm khác nhau của state table, với giá trị ban đầu:

state = bytes2matrix(plaintext)

AddRoundKey là công việc đầu tiên sẽ thực hiện, và sử dụng 4 word đầu tiên từ khóa đã mở rộng round_keys = expand_key(key):

add_round_key(state, round_keys[0])

Ở $9$ round tiếp theo, bốn công việc SubBytes, ShiftRows, MixColumns, AddRoundKey được thực hiện liên tục và lặp lại nên chúng ta sử dụng một vòng lặp if:

for i in range(1, N_ROUNDS):
    sub_bytes(state)
    shift_rows(state)
    mix_columns(state)
    add_round_key(state, round_keys[i])

Round cuối cùng chỉ thực hiện ba công việc SubBytes, ShiftRows, AddRoundKey:

sub_bytes(state)
shift_rows(state)
add_round_key(state, round_keys[10])

Cuối cùng chúng ta có hàm encrypt_aes() hoàn thiện:

def encrypt_aes(key, plaintext):
    round_keys = expand_key(key)

    # Convert plaintext to state matrix
    state = bytes2matrix(plaintext)
    # First AddRoundKey
    add_round_key(state, round_keys[0])

    for i in range(1, N_ROUNDS):
        sub_bytes(state)
        shift_rows(state)
        mix_columns(state)
        add_round_key(state, round_keys[i])
    # Final round
    sub_bytes(state)
    shift_rows(state)
    add_round_key(state, round_keys[10])
    # Convert state matrix to ciphertext
    ciphertext = matrix2bytes(state)
    # print(ciphertext)
    return ciphertext

Ví dụ thực hiện mã hóa chuỗi byte plaintext = b'VIBLOCTF{crypto}' với secret key key = b'\xc3,\\xa6\xb5\x80^\x0c\xdb\x8d\xa5z*\xb6\xfe\'

Đối với hàm giải mã, dễ dàng thay đổi từ hàm mã hóa theo thứ tự các công việc sắp xếp lại:

def decrypt_aes(key, ciphertext):
    round_keys = expand_key(key)

    # Convert ciphertext to state matrix
    state = bytes2matrix(ciphertext)
    # First AddRoundKey
    add_round_key(state, round_keys[10])

    for i in range(N_ROUNDS - 1, 0, -1):
        inv_shift_rows(state)
        inv_sub_bytes(state)
        add_round_key(state, round_keys[i])
        inv_mix_columns(state)
    # Final round
    inv_shift_rows(state)
    inv_sub_bytes(state)
    add_round_key(state, round_keys[0])

    # Convert state matrix to plaintext
    plaintext = matrix2bytes(state)
    return plaintext

Hai hàm mã hóa và giải mã trên thực ra chỉ có thể làm việc với thông điệp có độ dài $16$ ký tự. Câu hỏi đặt ra là đối với các thông điệp dài hơn thì sẽ giải quyết ra sao? Xin dành cho bạn đọc tìm hiểu và trả lời cho câu hỏi này!

2. Module AES trong thư viện Crypto.Cipher

Ở phần trên chúng ta đã xây dựng hoàn thiện hàm mã hóa và giải mã theo một block cho thuật toán AES-128. Lúc này có thể nhiều bạn sẽ có thắc mắc rằng: có phải mỗi lần sử dụng tới thuật toán này, chúng ta đều cần lập trình lại hai hàm này, chẳng lẽ không có thư viện nào hỗ trợ tốt cho thuật toán AES sao? Lý do Viblo cùng các bạn đi qua và phân tích từng giai đoạn giải mã, rồi lập trình từng hàm nhỏ, cuối cùng mới có được một hàm mã hóa / giải mã hoàn thiện là bởi vì chúng ta cần hiểu được bản chất của bài toán, sau đó mới có thể áp dụng tốt một thuật toán mạnh mẽ như AES.

Trả lời cho câu hỏi trên, tất nhiên là có! Trong Python có rất nhiều thư viện hỗ trợ cho thuật toán AES, Viblo muốn giới thiệu tới các bạn module AES trong thư viện Crypto.Cipher:

from Crypto.Cipher import AES

Ví dụ chương trình mã hóa:

from Crypto.Cipher import AES

data = b'VIBLOCTF{crypto}'
key = b'\xc3,\\\xa6\xb5\x80^\x0c\xdb\x8d\xa5z*\xb6\xfe\\'

cipher = AES.new(key, AES.MODE_ECB)
ciphertext = cipher.encrypt(data)
print(ciphertext)

Ví dụ chương trình giải mã:

from Crypto.Cipher import AES

key = b'\xc3,\\\xa6\xb5\x80^\x0c\xdb\x8d\xa5z*\xb6\xfe\\'
ciphertext = b'B\xf5\x9d\xb8\xd5\x84\x9a\x19\xcaS\xf6\xd0!d\xf3p'

cipher = AES.new(key, AES.MODE_ECB)
plaintext = cipher.decrypt(ciphertext)
print(plaintext)

Bạn đọc có thể tham khảo thêm cách thực hiện mã hóa và giải mã trong challenge Block cipher starter.

V. Challenge CTF

Passwords as keys là một challenge CTF mang tính tiêu biểu và khởi đầu cho thuật toán AES. Source code như sau:

from Crypto.Cipher import AES
import hashlib
import random


# /usr/share/dict/words from
# https://gist.githubusercontent.com/wchargin/8927565/raw/d9783627c731268fb2935a731a618aa8e95cf465/words
with open("/usr/share/dict/words") as f:
    words = [w.strip() for w in f.readlines()]
keyword = random.choice(words)

KEY = hashlib.md5(keyword.encode()).digest()
FLAG = ?


@chal.route('/passwords_as_keys/decrypt/<ciphertext>/<password_hash>/')
def decrypt(ciphertext, password_hash):
    ciphertext = bytes.fromhex(ciphertext)
    key = bytes.fromhex(password_hash)

    cipher = AES.new(key, AES.MODE_ECB)
    try:
        decrypted = cipher.decrypt(ciphertext)
    except ValueError as e:
        return {"error": str(e)}

    return {"plaintext": decrypted.hex()}


@chal.route('/passwords_as_keys/encrypt_flag/')
def encrypt_flag():
    cipher = AES.new(KEY, AES.MODE_ECB)
    encrypted = cipher.encrypt(FLAG.encode())

    return {"ciphertext": encrypted.hex()}

Dựa vào cách định nghĩa đối tượng cipher = AES.new(KEY, AES.MODE_ECB) có thể thấy challenge sử dụng thuật toán AES với mode ECB (Electronic Codebook). Tức là mỗi khối dữ liệu đầu vào được mã hóa độc lập và không phụ thuộc vào các khối khác trong văn bản gốc.

Trước hết, chúng ta có thể lấy được encrypt_flag dạng hex tại route /passwords_as_keys/encrypt_flag/. Response ở dạng JSON format nên chúng ta cần xử lý một chút:

import requests
import json

BASE_URL = 'https://aes.cryptohack.org'
encrypt_flag_path = '/passwords_as_keys/encrypt_flag/'

def get_ciphertext(url, path):
    r = requests.get(url = url + path)
    response = r.text.strip()
    data = json.loads(response)
    return data['ciphertext']

ciphertext = get_ciphertext(BASE_URL, encrypt_flag_path)
print(ciphertext)

Để giải mã AES, chúng ta buộc phải có giá trị của KEY. Chú ý cách key được tạo ra:

with open("/usr/share/dict/words") as f:
    words = [w.strip() for w in f.readlines()]
keyword = random.choice(words)

KEY = hashlib.md5(keyword.encode()).digest()

Sau khi chọn ngẫu nhiên một keyword trong danh sách words lấy tại /usr/share/dict/words, chương trình thực hiện mã hóa MD5 word đó làm giá trị của KEY.

Bởi vậy, có thể lựa chọn phương pháp tấn công brute force thử tất cả trường hợp của KEY trong wordlist nhằm tìm ra flag. Chương trình lời giải tham khảo như sau:

import requests
import json
import hashlib
from Crypto.Cipher import AES

BASE_URL = 'https://aes.cryptohack.org'
encrypt_flag_path = '/passwords_as_keys/encrypt_flag/'
wordlist_url = 'https://gist.githubusercontent.com/wchargin/8927565/raw/d9783627c731268fb2935a731a618aa8e95cf465/words'

def get_ciphertext(url, path):
    r = requests.get(url = url + path)
    response = r.text.strip()
    data = json.loads(response)
    return data['ciphertext']

def decrypt_aes(KEY, ciphertext):
    cipher = AES.new(KEY, AES.MODE_ECB)
    flag = cipher.decrypt(ciphertext)
    return flag

def decrypt(ciphertext):
    r = requests.get(wordlist_url)
    wordlist = r.text.strip().split()
    ciphertext = bytes.fromhex(ciphertext)
    for word in wordlist:
        KEY = hashlib.md5(word.encode()).digest()
        flag = decrypt_aes(KEY, ciphertext)
        try:
            print(flag.decode())
        except:
            next

ciphertext = get_ciphertext(BASE_URL, encrypt_flag_path)
decrypt(ciphertext)

Tài liệu tham khảo

• https://en.wikipedia.org/wiki/Advanced_Encryption_Standard
• https://www.geeksforgeeks.org/advanced-encryption-standard-aes/
• https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.197.pdf
• https://cs.ru.nl/~joan/papers/JDA_VRI_Rijndael_2002.pdf
• https://pycryptodome.readthedocs.io/en/latest/src/cipher/aes.html