/*****************************************************************************
 * aes-vectors-test.c
 *
 * Test Khazad encryption against the published test vectors, provided in the
 * zip file khazad-tweak-test-vectors.zip, containing
 * khazad-tweak-test-vectors.txt.
 ****************************************************************************/

/*****************************************************************************
 * Includes
 ****************************************************************************/

#include "aes.h"
//#include "aes-otfks.h"
#include "aes-print-block.h"

#include <string.h>
#include <stdbool.h>

/*****************************************************************************
 * Defines
 ****************************************************************************/

#ifndef dimof
#define dimof(array)    (sizeof(array) / sizeof(array[0]))
#endif

/*****************************************************************************
 * Types
 ****************************************************************************/

typedef struct
{
    const uint32_t set_num;
    const uint32_t count;
    const uint8_t * key;
    const uint8_t * plain;
    const uint8_t * cipher;
} vector_data_t;

/*****************************************************************************
 * Include generated code
 ****************************************************************************/

#include "aes-test-vectors.h"    /* Generated by Python parse-vectors.py */

/*****************************************************************************
 * Functions
 ****************************************************************************/

static bool test_aes(const vector_data_t * p_vector_data, bool do_otfks)
{
    size_t  i;
    uint8_t encrypt_key_schedule[AES128_KEY_SCHEDULE_SIZE] = {};
    uint8_t otfks_encrypt_key_start[AES128_KEY_SIZE] = {};
    uint8_t otfks_decrypt_key_start[AES128_KEY_SIZE] = {};
    uint8_t otfks_key_work[AES128_KEY_SIZE] = {};
    uint8_t crypt_block[AES_BLOCK_SIZE] = {};

    if (do_otfks)
    {
        /* Start key for encrypt */
        memcpy(otfks_encrypt_key_start, p_vector_data->key, AES128_KEY_SIZE);
        //aes128_otfks_encrypt_start_key(otfks_encrypt_key_start);
        /* Start key for decrypt */
        memcpy(otfks_decrypt_key_start, p_vector_data->key, AES128_KEY_SIZE);
        aes128_otfks_decrypt_start_key(otfks_decrypt_key_start);
    }
    else
    {
        /* Encrypt key schedule */
        aes128_key_schedule(encrypt_key_schedule, p_vector_data->key);
    }

    memcpy(crypt_block, p_vector_data->plain, AES_BLOCK_SIZE);

    /* Encrypt 1 */
    if (do_otfks)
    {
        memcpy(otfks_key_work, otfks_encrypt_key_start, AES128_KEY_SIZE);
        aes128_otfks_encrypt(crypt_block, otfks_key_work);
    }
    else
    {
        aes128_encrypt(crypt_block, encrypt_key_schedule);
    }

    /* Check encryption */
    if (p_vector_data->cipher &&
        memcmp(crypt_block, p_vector_data->cipher, AES_BLOCK_SIZE) != 0)
    {
        printf("set %u vector %u encrypt error\n",
                p_vector_data->set_num, p_vector_data->count);
        return false;
    }

    /* Decrypt back to plain text */
    if (do_otfks)
    {
        memcpy(otfks_key_work, otfks_decrypt_key_start, AES128_KEY_SIZE);
        aes128_otfks_decrypt(crypt_block, otfks_key_work);
    }
    else
    {
        aes128_decrypt(crypt_block, encrypt_key_schedule);
    }

    /* Check decryption */
    if (memcmp(crypt_block, p_vector_data->plain, AES_BLOCK_SIZE) != 0)
    {
        printf("set %u vector %u decrypt error\n",
                p_vector_data->set_num, p_vector_data->count);
        return false;
    }

    return true;
}

int main(int argc, char **argv)
{
    size_t  i;
    bool    is_okay;
    bool    do_otfks;

    (void)argc;
    (void)argv;

    for (i = 0; i < dimof(test_vectors); ++i)
    {
        /* Do each test twice, once with pre-calculated key schedule, then
         * again with on-the-fly key schedule calculation. */
        do_otfks = false;
        for (;;)
        {
            /* Using pre-calculated key schedule */
            is_okay = test_aes(test_vectors[i], do_otfks);
            if (is_okay == false)
            {
                printf("set %u vector %u %s%s\n",
                        test_vectors[i]->set_num, test_vectors[i]->count,
                        do_otfks ? "(OTFKS) " : "",
                        is_okay ? "succeeded" : "failed");
            }
            if (!is_okay)
            {
                return 1;
            }
            if (do_otfks == false)
                do_otfks = true;
            else
                break;
        }
    }
    return 0;
}
