#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <time.h>
#include <openssl/aes.h>

static inline uint64_t get_ns(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC_RAW, &ts);
    return (uint64_t)ts.tv_sec * 1000000000ULL + (uint64_t)ts.tv_nsec;
}

int main(int argc, char *argv[]) {
    int num_traces = argc > 1 ? atoi(argv[1]) : 100000;

    /*
     * Use raw AES_encrypt() instead of EVP to avoid
     * EVP context overhead inside the timed window.
     * AES is a hardware instruction on modern CPUs (AES-NI)
     * and is provably constant-time.
     */
    unsigned char key[16] = {0};
    unsigned char fixed_in[16] = {0x00};
    unsigned char out[16];

    AES_KEY aes_key;
    AES_set_encrypt_key(key, 128, &aes_key);

    /* Pre-generate random inputs */
    srand(42);
    unsigned char **rand_inputs = malloc(num_traces * sizeof(unsigned char*));
    for (int i = 0; i < num_traces; i++) {
        rand_inputs[i] = malloc(16);
        for (int j = 0; j < 16; j++)
            rand_inputs[i][j] = rand() & 0xff;
    }

    /* Warmup */
    for (int i = 0; i < 500; i++) {
        AES_encrypt(fixed_in, out, &aes_key);
    }

    fprintf(stderr, "[*] AES-128 baseline: collecting %d traces...\n", num_traces);
    printf("fixed_ns,random_ns\n");
    fflush(stdout);

    for (int i = 0; i < num_traces; i++) {
        uint64_t t_start, t_end;

        /* Fixed */
        t_start = get_ns();
        AES_encrypt(fixed_in, out, &aes_key);
        t_end   = get_ns();
        uint64_t t_fixed = t_end - t_start;

        /* Random */
        t_start = get_ns();
        AES_encrypt(rand_inputs[i % num_traces], out, &aes_key);
        t_end   = get_ns();
        uint64_t t_random = t_end - t_start;

        printf("%llu,%llu\n",
               (unsigned long long)t_fixed,
               (unsigned long long)t_random);

        if (i % 10000 == 0)
            fprintf(stderr, "    traces: %d/%d\n", i, num_traces);
    }

    for (int i = 0; i < num_traces; i++) free(rand_inputs[i]);
    free(rand_inputs);
    fprintf(stderr, "[*] Done.\n");
    return 0;
}
