In this lab, we will implement Trust On First Use in a manner similar to that of SSH. We will base our implementation around PyCryptodome - you can find the relevant documentation here.

Use these commands to generate a private key and a Diffie-Hellman parameter file:

openssl genrsa -out private.pem 2048
openssl dhparam -out dhparam.pem 2048

The dhparam file will be used as a hardcoded .pem that contains the necessary Diffie-Hellman parameters to generate our DH keys. These values are generally decided upon by convention and are hardcoded - check RFC 7919. An example which can be found there is ffdhe2048.

To better understand its structure, you can use the following commands:

openssl dhparam -in dhparam.pem -text -noout
openssl asn1parse < dhparam.pem

Create a Python environment or use an existing one, then install the required packages:

python3 -m venv create env
source ./env/bin/activate
pip install --upgrade pip
pip install pycryptodome pyasn1 pyasn1-modules

We will use the pyasn1 library to parse the .pem file's binarized DER format more conveniently.

Starting from these files, solve the TODO 1 series in dhe_server.py and dhe_client.py.

On the server side, you should:

• Send the RSA public key to the client.
• Generate a DH key pair (use the data hardcoded in dhparam.pem) and send the public key to the client.
• Generate a signature over the RSA and DH public keys (concatenate and hash them, then sign using the RSA private key) and send it to the client. Use this as your guide.
• Receive the client's DH public key and generate the shared DH secret.
• Derive a symmetric key from the shared secret. Use HKDF for this.

On the client side, you should:

• Receive the RSA public key from the server.
• Receive the server's DH public key.
• Receive the server's signature over the RSA and DH public keys and verify it using the server's public key.
• Generate a DH key pair and send the public key to the server.
• Compute the shared DH secret.
• Derive a symmetric key from the shared secret.

Generating a signature over the RSA and DH public keys is a way to authenticate the remote host. If the client successfully verifies this signature using the server's public key, then the server is authenticated unless the public key itself has been replaced by the attacker in a man in the middle attack. We will look at a way to (mostly) solve this issue in the next task.

The symmetric key derived from the shared secret will be used to encrypt the communication between the client and server. Although we're stopping the tasks here, this key would be the one that you would use to encrypt and decrypt the communication between the client and server. If you want, you can check the PyCryptodome documentation for more information on how to use AES. To see what ciphers SSH uses, run the following command:

ssh -Q cipher

Now start solving the TODO 2 series by implementing Trust On First Use in dhe_client.py. Do it as follows:

• Store the public key of the server in a file named known_hosts in the following format (the same way SSH does it):

hostname1 public_key1
hostname2 public_key2
...

• If the client already has a public key for the given IP of the server, check that the public key of the server matches the one that is stored (if it's a first connection - i.e., First Use - just store the public key and Trust the host).
• If it matches, print a “connection established” message and proceed to use that key for verification of the signature over the DH share of the server.
• If it doesn't match, print a suggestive error message and exit.

This is very similar to what SSH does when connecting to a server using a pair of public/private keys and is known as Trust On First Use (TOFU) authentication.

Now let's update the key exchange to use post-quantum cryptography.

For this you can either modify the work you did above to use a post-quantum key exchange method such as ML-KEM in Python. To do this in Python you can use a library such as ml-kem. You can install this as follows:

pip install ml-kem

Below is an example of how to use this library in a simple client-server scenario. Note: this example is provided by Gemini AI, so use it with caution and double-check it: </code> # Install the library first: pip install ml-kem

from mlkem.ml_kem import ML_KEM from mlkem.parameter_set import ML_KEM_768 # Recommended security level import secrets

# — Server (Alice) Side —

def server_keygen():

  """Alice generates her ML-KEM key pair."""
  # Initialize ML-KEM with the desired security level
  ml_kem = ML_KEM(parameters=ML_KEM_768, randomness=secrets.token_bytes)
  
  # Generate the encapsulation key (ek, public) and decapsulation key (dk, private)
  ek, dk = ml_kem.key_gen()
  
  print("Alice: Generated Public Key (ek) and Private Key (dk).")
  return ek, dk, ml_kem

def server_decapsulate(dk, c, ml_kem):

  """Alice decapsulates the ciphertext to get the shared secret."""
  try:
      K_prime = ml_kem.decaps(dk, c)
      print("Alice: Successfully decapsulated the Shared Secret (K').")
      return K_prime
  except ValueError as e:
      print(f"Alice: Decapsulation failed! {e}")
      return None

# — Client (Bob) Side —

def client_encapsulate(ek):

  """Bob encapsulates a shared secret using Alice's public key."""
  ml_kem = ML_KEM(parameters=ML_KEM_768, randomness=secrets.token_bytes)
  
  # Encapsulate to get the shared secret (K) and the ciphertext (c)
  K, c = ml_kem.encaps(ek)
  
  print("Bob: Encapsulated a Shared Secret (K) and created Ciphertext (c).")
  return K, c

# — Communication Flow Simulation —

# 1. Server (Alice) Key Generation ek_server, dk_server, ml_kem_instance = server_keygen()

# 2. Public Key Transmission (ek_server is sent to the client) print(“\n— Network Transmission: ek sent to Bob —”)

# 3. Client (Bob) Encapsulation K_client, c_client = client_encapsulate(ek_server)

# 4. Ciphertext Transmission (c_client is sent back to the server) print(“\n— Network Transmission: c sent to Alice —”)

# 5. Server (Alice) Decapsulation K_server = server_decapsulate(dk_server, c_client, ml_kem_instance)

# 6. Verification print(“\n— Verification —”) if K_client is not None and K_server is not None:

  if K_client == K_server:
      print("Success! Alice's and Bob's shared secrets match.")
      # The shared secret can now be used as an AES key, e.g., K_client
      # The shared secret is bytes:
      # print(f"Shared Secret: {K_client.hex()}")
  else:
      print("Failure! Shared secrets do not match.")

else:

  print("Failure in key exchange process.")

</code>

Otherwise, you can start from the Diffie-Hellman key exchange lab we did in OpenSSL/C. You may start from this code, that provides a working solution for the Diffie-Hellman lab in OpenSSL.

If you go for the C implementation in OpenSSL, one option to include post-quantum encryption is using liboqs. To install it on Ubuntu, use the command below:

sudo apt install astyle cmake gcc ninja-build libssl-dev python3-pytest python3-pytest-xdist unzip xsltproc doxygen graphviz python3-yaml valgrind

If using this library, you may use this example to implement your own version of key exchange and test it in the scenarios above (key exchange and key exchange + TOFU). You may use any resources or available KEMs.

Alternatively, you can also use OpenSSL directly (starting with version 3.5). You can find some documentation for ML-KEM here and here.

Below is an example code obtained through Gemini AI (not fully tested, but perhaps a good starting point):

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/err.h>
#include <openssl/provider.h>

// --- HARDCODED ML-KEM-768 Public Key (Replace with a real key) ---
// A real ML-KEM-768 public key is 608 bytes long.
// This is a placeholder for demonstration.
unsigned char server_public_key_bytes[] = {
    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
    // ... 592 more bytes of a real ML-KEM-768 public key would go here ...
    0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
    // Note: This placeholder is NOT a valid ML-KEM-768 key and will likely fail.
    // Replace with a 608-byte key generated via 'openssl genpkey' for real use.
};
size_t public_key_len = 608; // ML-KEM-768 Public Key size

// Error handling macro
#define HANDLE_ERROR(msg) \
    { \
        fprintf(stderr, "%s failed.\n", msg); \
        ERR_print_errors_fp(stderr); \
        goto cleanup; \
    }

int main(void) {
    EVP_PKEY_CTX *kctx = NULL;
    EVP_PKEY *peer_pub_key = NULL;
    OSSL_PARAM params[2];
    unsigned char *shared_secret = NULL;
    unsigned char *ciphertext = NULL;
    size_t shared_secret_len = 0;
    size_t ciphertext_len = 0;
    int ret = 0;

    printf("Starting ML-KEM Client Key Encapsulation (ML-KEM-768).\n");

    // 1. Load the OpenSSL default provider (ML-KEM is in the default provider since 3.5)
    if (!OSSL_PROVIDER_load(NULL, "default")) {
        HANDLE_ERROR("Failed to load default provider");
    }

    // 2. Import the raw public key bytes into an EVP_PKEY structure
    printf("2. Importing ML-KEM-768 Public Key...\n");

    // a. Create parameter array to describe the key components
    params[0] = OSSL_PARAM_construct_octet_string("pub", server_public_key_bytes, public_key_len);
    params[1] = OSSL_PARAM_construct_end();

    // b. Import the key. ML-KEM-768 is named "ML-KEM-768"
    peer_pub_key = EVP_PKEY_new_from_params(NULL, params, "ML-KEM-768");
    if (!peer_pub_key) {
        // NOTE: This will fail if the hardcoded bytes are not a valid key.
        HANDLE_ERROR("EVP_PKEY_new_from_params failed (Public Key Import)");
    }

    // 3. Initialize the Key Encapsulation Context
    kctx = EVP_PKEY_CTX_new_from_pkey(NULL, peer_pub_key, NULL);
    if (!kctx) {
        HANDLE_ERROR("EVP_PKEY_CTX_new_from_pkey failed");
    }

    if (EVP_PKEY_encapsulate_init(kctx) <= 0) {
        HANDLE_ERROR("EVP_PKEY_encapsulate_init failed");
    }

    // 4. Determine buffer sizes
    printf("4. Determining buffer sizes...\n");
    if (EVP_PKEY_encapsulate(kctx, NULL, &shared_secret_len, NULL, &ciphertext_len) <= 0) {
        HANDLE_ERROR("EVP_PKEY_encapsulate (size determination) failed");
    }
    
    printf("   - Shared Secret Length: %zu bytes\n", shared_secret_len);
    printf("   - Ciphertext Length: %zu bytes\n", ciphertext_len);

    // 5. Allocate buffers
    shared_secret = OPENSSL_malloc(shared_secret_len);
    ciphertext = OPENSSL_malloc(ciphertext_len);

    if (!shared_secret || !ciphertext) {
        HANDLE_ERROR("Memory allocation failed");
    }

    // 6. Perform the Encapsulation (Client's Key Exchange Step)
    printf("6. Performing Key Encapsulation...\n");
    if (EVP_PKEY_encapsulate(kctx, shared_secret, &shared_secret_len, ciphertext, &ciphertext_len) <= 0) {
        HANDLE_ERROR("EVP_PKEY_encapsulate failed");
    }

    printf("✅ Encapsulation successful!\n");
    printf("   The client has generated a **Shared Secret** and the **Ciphertext**.\n");
    
    printf("   - Shared Secret (First 16 bytes): ");
    for (size_t i = 0; i < (shared_secret_len > 16 ? 16 : shared_secret_len); i++) {
        printf("%02x", shared_secret[i]);
    }
    printf("...\n");

    printf("   - Ciphertext (First 16 bytes, to be sent to Server): ");
    for (size_t i = 0; i < (ciphertext_len > 16 ? 16 : ciphertext_len); i++) {
        printf("%02x", ciphertext[i]);
    }
    printf("...\n");

    ret = 0; // Success
    
cleanup:
    // Free all allocated resources
    EVP_PKEY_CTX_free(kctx);
    EVP_PKEY_free(peer_pub_key);
    OPENSSL_free(shared_secret);
    OPENSSL_free(ciphertext);
    OSSL_PROVIDER_unload(OSSL_PROVIDER_load(NULL, "default")); // Unload provider

    if (ret != 0) {
        return EXIT_FAILURE;
    }
    return EXIT_SUCCESS;
}

Some instructions to use this code (also from Gemini AI):

• Save the code as mlkem_client.c.
• Generate a real ML-KEM-768 public key (if you didn't do this in the previous step) by running the following commands:

openssl genpkey -algorithm ML-KEM-768 -out server_private.pem
openssl pkey -in server_private.pem -pubout -out server_public.pem

• Extract the raw public key bytes to replace the placeholder in the C code. A common way to get the raw bytes is to convert the PEM file to a DER format and extract the key component. For ML-KEM, the raw key is often found within the ASN.1 structure.

• Compile the code (adjusting paths to your OpenSSL 3.5 installation), using a command like this (check again also the DH lab for more details):

gcc -o mlkem_client mlkem_client.c -I/path/to/openssl/include -L/path/to/openssl/lib -lssl -lcrypto

The following GeminiAI-generated scripts (might have some bugs, please double-check, they are only given as reference/helper tools) generate an ML-KEM-768 key pair and, importantly, forces OpenSSL to save the private key in the FIPS 203 'dk' format (the one the C code expects for decapsulation) rather than the default 'seed' format, which is not what the raw C code imports:

#!/bin/bash
KEY_ALG="ML-KEM-768"
PRIV_FILE="server_private_dk.pem"
PUB_FILE="server_public_ek.pem"
RAW_PRIV_FILE="server_private_raw.bin"
RAW_PUB_FILE="server_public_raw.bin"

echo "--- 1. Generating $KEY_ALG Key Pair (forcing dk-only private key) ---"

# Generate the private key, using '-provparam ml-kem.retain_seed=no'
# to ensure it saves the full FIPS 203 'dk' decapsulation key,
# NOT just the seed.
openssl genpkey -algorithm "$KEY_ALG" \
    -provparam ml-kem.retain_seed=no \
    -out "$PRIV_FILE"

# Extract the public key from the private key file
openssl pkey -in "$PRIV_FILE" -pubout -out "$PUB_FILE"

echo "Keys generated: $PRIV_FILE and $PUB_FILE"
echo "----------------------------------------------------------------"

The ML-KEM public key is in the SubjectPublicKeyInfo structure in the PEM file. We use openssl pkey to output the key in DER format and then remove the ASN.1 header/wrapper to get the raw bytes:

echo "--- 2. Extracting RAW Public Key (608 bytes) ---"

# The ML-KEM public key is 608 bytes long.
# The 'pubout' option converts to DER, and 'tail' removes the ASN.1 wrapper.
# NOTE: The ASN.1 wrapper size may vary by OpenSSL version/config. We must
# determine the exact offset to trim off the SubjectPublicKeyInfo header.
# For standard ML-KEM-768 PKCS#8 output, the header is typically 24 bytes.
HEADER_BYTES=24 # Common offset for ML-KEM-768 PKCS#8 SubjectPublicKeyInfo header

openssl pkey -in "$PUB_FILE" -pubin -outform DER -out "$RAW_PUB_FILE"
# Trim the ASN.1 header to get the raw 608-byte key
# (Skip first $HEADER_BYTES, save the next 608 bytes)
dd if="$RAW_PUB_FILE" of="client_public_key_raw_final.bin" bs=1 skip=$HEADER_BYTES count=608 status=none

echo "Raw public key saved to client_public_key_raw_final.bin"
# Print the hex array for C code
echo "Public Key Hex Array:"
xxd -p "client_public_key_raw_final.bin" | tr -d '\n' | sed 's/../&, 0x/g' | sed 's/^, 0x/0x/' | sed 's/, $//'
echo
echo "----------------------------------------------------------------"

Similarly, the private key (the FIPS 203 'dk' component) is inside the PKCS#8 private key structure. We again convert to DER and trim the ASN.1 wrappers:

echo "--- 3. Extracting RAW Private Key (1184 bytes) ---"

# The ML-KEM-768 private key (dk) is 1184 bytes long.
# We convert to DER and again skip the ASN.1 header bytes.
# For PKCS#8 'dk' private key, the header is typically 50 bytes.
HEADER_BYTES=50 # Common offset for ML-KEM-768 PKCS#8 PrivateKeyInfo header

openssl pkey -in "$PRIV_FILE" -outform DER -out "$RAW_PRIV_FILE"
# Trim the ASN.1 header to get the raw 1184-byte key
# (Skip first $HEADER_BYTES, save the next 1184 bytes)
dd if="$RAW_PRIV_FILE" of="server_private_key_raw_final.bin" bs=1 skip=$HEADER_BYTES count=1184 status=none

echo "Raw private key saved to server_private_key_raw_final.bin"
# Print the hex array for C code
echo "Private Key Hex Array:"
xxd -p "server_private_key_raw_final.bin" | tr -d '\n' | sed 's/../&, 0x/g' | sed 's/^, 0x/0x/' | sed 's/, $//'
echo
echo "----------------------------------------------------------------"

You can now use the hex output from these scripts to reliably replace the placeholder arrays in your C client and server code. See also this video about Post-Quantum Cryptography support in OpenSSL.

Below is also an example code from Gemini AI as starting point for the server, to pair with the code above for the client. Note that these examples use hardcoded values for keys and do not imply any communication. Therefore, make sure you use them for a proper key exchange that uses communication between parties to exchange the public keys and ciphertexts, in a similar manner to what you did for the DH lab.

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/err.h>
#include <openssl/provider.h>

// --- HARDCODED ML-KEM-768 Private Key (Replace with a real key) ---
// A real ML-KEM-768 private key is 1184 bytes long.
unsigned char server_private_key_bytes[] = {
    0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2a, 0x2b, 0x2c, 0x2d, 0x2e, 0x2f,
    // ... 1168 more bytes of a real ML-KEM-768 private key would go here ...
    0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3a, 0x3b, 0x3c, 0x3d, 0x3e, 0x3f,
};
size_t private_key_len = 1184; // ML-KEM-768 Private Key size

// --- HARDCODED Ciphertext (Generated by the Client - Replace with real data) ---
// A real ML-KEM-768 ciphertext is 1088 bytes long.
unsigned char client_ciphertext_bytes[] = {
    0xa0, 0xa1, 0xa2, 0xa3, 0xa4, 0xa5, 0xa6, 0xa7, 0xa8, 0xa9, 0xaa, 0xab, 0xac, 0xad, 0xae, 0xaf,
    // ... 1072 more bytes of a real ML-KEM-768 ciphertext would go here ...
    0xb0, 0xb1, 0xb2, 0xb3, 0xb4, 0xb5, 0xb6, 0xb7, 0xb8, 0xb9, 0xba, 0xbb, 0xbc, 0xbd, 0xbe, 0xbf,
};
size_t ciphertext_len = 1088; // ML-KEM-768 Ciphertext size

// --- Hardcoded Shared Secret Size (32 bytes for ML-KEM-768) ---
size_t shared_secret_len = 32;

// Error handling macro
#define HANDLE_ERROR(msg) \
    { \
        fprintf(stderr, "%s failed.\n", msg); \
        ERR_print_errors_fp(stderr); \
        goto cleanup; \
    }

int main(void) {
    EVP_PKEY_CTX *kctx = NULL;
    EVP_PKEY *server_priv_key = NULL;
    OSSL_PARAM params[2];
    unsigned char *recovered_secret = NULL;
    int ret = 0;

    printf("Starting ML-KEM Server Key Decapsulation (ML-KEM-768).\n");

    // 1. Load the OpenSSL default provider
    if (!OSSL_PROVIDER_load(NULL, "default")) {
        HANDLE_ERROR("Failed to load default provider");
    }

    // 2. Import the raw private key bytes into an EVP_PKEY structure
    printf("2. Importing ML-KEM-768 Private Key...\n");

    // a. Create parameter array to describe the key components
    // The ML-KEM private key is referred to as "priv" (or 'dk' in FIPS 203 terms)
    params[0] = OSSL_PARAM_construct_octet_string("priv", server_private_key_bytes, private_key_len);
    params[1] = OSSL_PARAM_construct_end();

    // b. Import the key.
    server_priv_key = EVP_PKEY_new_from_params(NULL, params, "ML-KEM-768");
    if (!server_priv_key) {
        HANDLE_ERROR("EVP_PKEY_new_from_params failed (Private Key Import)");
    }

    // 3. Initialize the Key Decapsulation Context
    kctx = EVP_PKEY_CTX_new_from_pkey(NULL, server_priv_key, NULL);
    if (!kctx) {
        HANDLE_ERROR("EVP_PKEY_CTX_new_from_pkey failed");
    }

    if (EVP_PKEY_decapsulate_init(kctx) <= 0) {
        HANDLE_ERROR("EVP_PKEY_decapsulate_init failed");
    }

    // 4. Allocate buffer for the shared secret
    recovered_secret = OPENSSL_malloc(shared_secret_len);
    if (!recovered_secret) {
        HANDLE_ERROR("Memory allocation failed");
    }

    // 5. Perform the Decapsulation (Server's Key Exchange Step)
    printf("5. Performing Key Decapsulation...\n");

    // The ciphertext is passed as the input buffer (client_ciphertext_bytes).
    // The recovered shared secret is written to the output buffer (recovered_secret).
    if (EVP_PKEY_decapsulate(kctx, recovered_secret, &shared_secret_len, client_ciphertext_bytes, ciphertext_len) <= 0) {
        HANDLE_ERROR("EVP_PKEY_decapsulate failed");
    }

    printf("✅ Decapsulation successful!\n");
    printf("   The server has recovered the **Shared Secret**.\n");
    
    printf("   - Recovered Shared Secret (First 16 bytes): ");
    for (size_t i = 0; i < (shared_secret_len > 16 ? 16 : shared_secret_len); i++) {
        printf("%02x", recovered_secret[i]);
    }
    printf("...\n");

    printf("\n**NOTE: For a successful key exchange, the hash of this secret must match the hash of the client's secret.**\n");

    ret = 0; // Success
    
cleanup:
    // Free all allocated resources
    EVP_PKEY_CTX_free(kctx);
    EVP_PKEY_free(server_priv_key);
    OPENSSL_free(recovered_secret);
    OSSL_PROVIDER_unload(OSSL_PROVIDER_load(NULL, "default")); // Unload provider

    if (ret != 0) {
        return EXIT_FAILURE;
    }
    return EXIT_SUCCESS;
}