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 try to replace the key exchange used above with post-quantum cryptography. For this you can use different libraries.

One option 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 – you may do something similar to get a starting point for the server):

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