Mechanic II - CryptoCTF

In this challenge, we need to analyze an oracle involving PQC. More specifically, involving MLKEM-1024. We are given a file mechanic_ii.py. First, it generates 1337 MLKEM key pairs:

kem = MLKEM_1024()
c, n = 0, 1337
KEY_PAIR = [kem.keygen() for _ in range(n)]
SKEYS = [KEY_PAIR[_][1] for _ in range(n)]
r = randint(0, n - 1)
cipher, shasec = kem.encaps(KEY_PAIR[r][0])
secret = hashlib.sha3_256(shasec + hashlib.sha3_256(shasec + str(r).encode()).digest()).hexdigest()

It picks one of those pairs at random, computes the shared secret, and then computes a secret string based on the shared secret and the random chosen index r. The main loop of the challenge goes like this:

while True:
    if c > 3 * n:
        die(border, 'The server is need to rest :/')
    pr(f"{border} Options: \n{border}\t[D]ecrypt cipher \n{border}\t[R]andomize a secret key! \n{border}\t[S]ubmit the secret \n{border}\t[Q]uit")
    ans = sc().decode().strip().lower()
    c += 1
    if ans == 'd':
        pr(border, 'Please select an ID: ')
        _id = sc().decode().strip()
        try:
            _id = int(_id)
            _shasec = kem.decaps(SKEYS[_id], cipher)
        except:
            die(border, 'Your input ID is invalid! Bye!!')
        pr(border, f'{_shasec = }')
    elif ans == 'r':
        pr(border, 'Please select an ID: ')
        _id = sc().decode().strip()
        try:
            _id = int(_id)
            _skey = SKEYS[_id][:-32] + os.urandom(32)
        except:
            die(border, 'Your input ID is invalid! Bye!!')
        SKEYS.append(_skey)
    elif ans == 's':
        pr(border, 'Please send the secret: ')
        _secret = sc().decode().strip()
        if _secret == secret:
            die(border, f'Congrats, you got the flag: {flag}')
        else:
            die(border, 'Your secret is incorrect! Bye!!')
    elif ans == 'q':
        die(border, "Quitting...")
    else:
        die(border, "Bye...")

If we want the flag, we need to know the secret string calulated before, and we only have one chance to do so. If we fail, the connection is closed and cannot try again.

ML-KEM

Let’s do a recap on ML-KEM. It is a post-quantum key encapsulation mechanism (KEM) whose purpose is to securely establish shared secret keys, even against quantum adversaries. It is supposed to replace classical key exchange (like RSA or ECDH) in a quantum-safe world.

We are not going to dive in the internals of ML-KEM, but just keep in mind that it relies on the hardness of the Module Learning With Errors (MLWE) problem.

This scheme works like this:

1. Alice generates her key, consisting of a public and a private key
2. Alice sends her public to Bob
3. Bob uses Alice’s public key to generate a shared secret and a ciphertext. This process is called encapsulation
4. Alice receives the ciphertext and recovers the same shared secret using her private key. This process is called decapsulation

Back to the challenge

In python, when using the quantcrypt library, the public and private keys of ML-KEM are represented like bytes:

sage: from quantcrypt.kem import MLKEM_1024
....: kem = MLKEM_1024()
....: key = kem.keygen()
....: type(key[0]), len(key[0])
(<class 'bytes'>, 1568)
sage: type(key[1]), len(key[1])
(<class 'bytes'>, 3168)

In the loop of the challenge, we can do any of these actions (up to $3\cdot 1337 = 4011$ times):

1. Decpasulate the ciphertext using the private key at the index of our choosing
2. Generate a new secret key in the array SKEYS, by changing the last 32 bytes of an already-existing secret key
3. Get the flag

If we look at the ML-KEM standard, published by the NIST, we can see how the key is constructed in the Algorithm 16. More specifically, we can see that the private key is generated as follows:

dk = (dk_PKE || ek || H(ek) || z)

where ek is the encapsulation key, H(ek) is some hash function of the public key, dk_PKE is the key for decapsulation and z is a 32-bytes random string. What the challenge does is changing this last part of the key (the z component). Let’s look more carefully how this value is used in the decapsulation Algorithm 18

𝑚′ ← K-PKE.Decrypt(dkPKE, 𝑐) # Decrypt ciphertext
(𝐾′, 𝑟′) ← G(𝑚′‖ℎ)̄
𝐾̄  ← J(𝑧‖𝑐)
𝑐′ ← K-PKE.Encrypt(ekPKE, 𝑚′, 𝑟′) # re-encrypt using the derived randomness 𝑟′
if 𝑐 ≠ 𝑐′ then
    𝐾′ ← 𝐾̄ # if ciphertexts do not match, “implicitly reject”
end if
return 𝐾′

As we can see, the z value is used only when the decapsulation algorithm fails. That is, when the ciphertext has not really been generated with that key. Say that we have a key pair, $K$, whose public key has been used to generate a ciphertext, $c$ and let’s look at two different cases. First:

1. We decapsulate the ciphertext $c$ using the private key component of $K$, getting a shared secret $S$
2. We change the last 32 bytes of the private key
3. We decapsulate the same ciphertext again. As the ciphertext has actually been generated with the encapsulation key of $K$, the $z$ value is not used, and we get the same shared secret $S$

Second case:

1. We decapsulate the ciphertext $c$ using the private key component of another key $K’$, getting a shared secret $S_1$
2. We change the last 32 bytes of the private key of $K'$
3. We decapsulate the same ciphertext again. As the ciphertext has not been generated with the encapsulation key of $K$, the $z$ value is used, and we get the another whole different shared secret $S_2$, with $S_1 \neq S_2$.

Thus, we have an oracle: changing the last 32 bytes of the key used to generate the shared secret won’t affect the shared secret. Changin the last 32 bytes of a key that has not been used will give us a completely different shared secret. Knowing this, we can now easily solve the challenge with this code:

def hash(shared, r):
    return hashlib.sha3_256(shared + hashlib.sha3_256(shared + str(r).encode()).digest()).hexdigest()

r = process("./mechanic_ii.py")

for _ in range(1337):
    r.sendlineafter(b"[Q]uit\n", b"r")
    r.recvline()
    r.sendline(str(_).encode())

shareds = []
for _ in range(1337):
    r.sendlineafter(b"[Q]uit\n", b"d")
    r.recvline()
    r.sendline(str(_).encode())
    r.recvuntil(b"_shasec = b")
    line = r.recvline().strip()[1:-1].decode("unicode_escape").encode("latin1")
    shareds.append(line)


for _ in range(1337):
    r.sendlineafter(b"[Q]uit\n", b"d")
    r.recvline()
    r.sendline(str(1337+_).encode())
    r.recvuntil(b"_shasec = b")
    shared = r.recvline().strip()[1:-1].decode("unicode_escape").encode("latin1")
    if shareds[_] == shared:
        r.sendlineafter(b"[Q]uit\n", b"s")
        hashed = hash(shared, _)
        r.recvline()
        r.sendline(hashed.encode())
        r.interactive()

$ python get_flag.py
[+] Starting local process './mechanic_ii.py': pid 1137859
[*] Switching to interactive mode
┃ Congrats, you got the flag: CCTF{fake_flag_4_testing}
[*] Got EOF while reading in interactive