#include "wvds_crypto.h"
#include <string.h>
 
int encrypt_message(
    const uint8_t* plaintext, size_t pt_len,
    const char* aad, size_t aad_len,
    uint8_t* ciphertext, size_t* ct_len,
    uint8_t nonce[12], uint8_t tag[16]
) {
    uint8_t request[65536];
    size_t request_len = sizeof(request);
 
    // 1. Izgradi zahtjev
    int rc = wvds_build_aes_encrypt_request(
        request, &request_len,
        1,                      // Key-ID (mora postojati)
        aad, aad_len,          // Additional Authenticated Data
        plaintext, pt_len      // Podaci za šifriranje
    );
    if (rc != 0) return rc;
 
    // 2. Pošalji zahtjev Crypto Serviceu
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    // 3. Parsiraj odgovor
    rc = wvds_parse_aes_encrypt_response(
        g_shared_memory, g_response_len,
        nonce, tag,            // Bit će popunjeni
        ciphertext, ct_len     // Šifrirani tekst + duljina
    );
 
    return rc;
}

int decrypt_message(
    const uint8_t* ciphertext, size_t ct_len,
    const uint8_t nonce[12], const uint8_t tag[16],
    const char* aad, size_t aad_len,
    uint8_t* plaintext, size_t* pt_len
) {
    uint8_t request[65536];
    size_t request_len = sizeof(request);
 
    // 1. Izgradi zahtjev
    int rc = wvds_build_aes_decrypt_request(
        request, &request_len,
        1,                      // Key-ID
        nonce, tag,            // Primljeno iz Encrypt
        aad, aad_len,          // Mora biti identično!
        ciphertext, ct_len
    );
    if (rc != 0) return rc;
 
    // 2. Pošalji zahtjev
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    // 3. Parsiraj odgovor
    rc = wvds_parse_aes_decrypt_response(
        g_shared_memory, g_response_len,
        plaintext, pt_len
    );
 
    // rc == 6 znači: Provjera oznake nije uspjela!
    return rc;
}

int sign_message(
    const uint8_t* message, size_t msg_len,
    uint8_t* signature, size_t* sig_len
) {
    uint8_t request[65536];
    size_t request_len = sizeof(request);
 
    int rc = wvds_build_mldsa_sign_request(
        request, &request_len,
        1,                      // Private Key-ID
        message, msg_len
    );
    if (rc != 0) return rc;
 
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    rc = wvds_parse_mldsa_sign_response(
        g_shared_memory, g_response_len,
        signature, sig_len      // ML-DSA-65: 3293 bajta
    );
 
    return rc;
}

int verify_signature(
    const uint8_t* message, size_t msg_len,
    const uint8_t* signature, size_t sig_len,
    int* valid                  // 1 = važeći, 0 = nevažeći
) {
    uint8_t request[65536];
    size_t request_len = sizeof(request);
 
    int rc = wvds_build_mldsa_verify_request(
        request, &request_len,
        2,                      // Public Key-ID
        message, msg_len,
        signature, sig_len
    );
    if (rc != 0) return rc;
 
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    rc = wvds_parse_mldsa_verify_response(
        g_shared_memory, g_response_len,
        valid
    );
 
    return rc;
}

// Serverska strana: KeyGen + Decaps
int server_key_exchange(
    uint8_t public_key[1184],   // ML-KEM-768 Public Key
    size_t* pk_len
) {
    uint8_t request[256];
    size_t request_len = sizeof(request);
 
    // KeyGen Request
    int rc = wvds_build_mlkem_keygen_request(
        request, &request_len,
        3                       // Key-ID za novi ključni par
    );
    if (rc != 0) return rc;
 
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    // Izvuci javni ključ
    rc = wvds_parse_mlkem_keygen_response(
        g_shared_memory, g_response_len,
        public_key, pk_len      // Privatni ključ ostaje u servisu
    );
 
    return rc;
}
 
int server_decaps(
    const uint8_t* ciphertext, size_t ct_len,
    uint8_t shared_secret[32]
) {
    uint8_t request[2048];
    size_t request_len = sizeof(request);
 
    int rc = wvds_build_mlkem_decaps_request(
        request, &request_len,
        3,                      // Key-ID iz KeyGen
        ciphertext, ct_len
    );
    if (rc != 0) return rc;
 
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    rc = wvds_parse_mlkem_decaps_response(
        g_shared_memory, g_response_len,
        shared_secret           // 32-bajtna dijeljena tajna
    );
 
    return rc;
}
 
// Klijentska strana: Encaps
int client_encaps(
    const uint8_t* public_key, size_t pk_len,
    uint8_t* ciphertext, size_t* ct_len,
    uint8_t shared_secret[32]
) {
    uint8_t request[2048];
    size_t request_len = sizeof(request);
 
    int rc = wvds_build_mlkem_encaps_request(
        request, &request_len,
        public_key, pk_len
    );
    if (rc != 0) return rc;
 
    memcpy(g_shared_memory, request, request_len);
    l4_ipc_send(g_crypto_ep, L4_TIMEOUT_NEVER);
    l4_ipc_receive(g_crypto_ep, L4_TIMEOUT_NEVER);
 
    rc = wvds_parse_mlkem_encaps_response(
        g_shared_memory, g_response_len,
        ciphertext, ct_len,     // Za server
        shared_secret           // Identično serveru nakon Decaps
    );
 
    return rc;
}

typedef struct {
    uint32_t sensor_id;
    uint32_t timestamp;
    float temperature;
    float humidity;
} SensorData;
 
int send_sensor_data(const SensorData* data) {
    // 1. Serijaliziraj
    uint8_t plaintext[sizeof(SensorData)];
    memcpy(plaintext, data, sizeof(SensorData));
 
    // 2. AAD: Sensor-ID + Timestamp (autenticirano, ne šifrirano)
    char aad[64];
    snprintf(aad, sizeof(aad), "sensor:%u:ts:%u",
             data->sensor_id, data->timestamp);
 
    // 3. Šifriraj
    uint8_t ciphertext[1024], nonce[12], tag[16];
    size_t ct_len;
 
    int rc = encrypt_message(
        plaintext, sizeof(plaintext),
        aad, strlen(aad),
        ciphertext, &ct_len,
        nonce, tag
    );
    if (rc != 0) return rc;
 
    // 4. Potpiši
    uint8_t to_sign[2048];
    size_t to_sign_len = 0;
    memcpy(to_sign + to_sign_len, nonce, 12); to_sign_len += 12;
    memcpy(to_sign + to_sign_len, tag, 16); to_sign_len += 16;
    memcpy(to_sign + to_sign_len, ciphertext, ct_len); to_sign_len += ct_len;
 
    uint8_t signature[4096];
    size_t sig_len;
 
    rc = sign_message(to_sign, to_sign_len, signature, &sig_len);
    if (rc != 0) return rc;
 
    // 5. Sastavi paket i pošalji
    // ... (mrežni kod) ...
 
    return 0;
}

< Instalacija | Dalje: Specifikacija protokola >