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.

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

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;
}