#include <stdint.h>
#include <string.h>

#include "parameters.h"
#include "tweakey.h"


#define LANE_BITS  64
#define LANE_BYTES (LANE_BITS/8)
#define LANES_NB   (TWEAKEY_BYTES/LANE_BYTES)


void tweakey_state_init(
    uint8_t TK[TWEAKEY_BYTES],
    const uint8_t key[KEY_BYTES],
    const uint8_t tweak[TWEAK_BYTES]
)
{
    memcpy(TK,             tweak, TWEAK_BYTES);
    memcpy(TK+TWEAK_BYTES, key,   KEY_BYTES);
}


void tweakey_state_extract(
    const uint8_t TK[TWEAKEY_BYTES],
    uint8_t round_constant,
    uint8_t round_tweakey[ROUND_TWEAKEY_BYTES]
)
{
    memset(round_tweakey, 0, ROUND_TWEAKEY_BYTES);

    for (size_t j=0; j<LANES_NB; j++)
    {
        const uint8_t *TKj = TK + j*LANE_BYTES;

        for (size_t k=0; k<LANE_BYTES; k++)
        {
            round_tweakey[k] ^= TKj[k];
        }
    }

    round_tweakey[0] ^= round_constant;
}


static void _multiply_M(const uint8_t old[LANE_BYTES], uint8_t new[LANE_BYTES])
{
    new[7] = old[6];
    new[6] = old[5];
    new[5] = old[4] ^ old[5]<<3;
    new[4] = old[3] ^ old[4]>>3;
    new[3] = old[2];
    new[2] = old[1] ^ old[6]<<2;
    new[1] = old[0];
    new[0] = old[7];
}

static const uint8_t M_1[256] = {
    0x00, 0x08, 0x10, 0x18, 0x20, 0x28, 0x30, 0x38,
    0x41, 0x49, 0x51, 0x59, 0x61, 0x69, 0x71, 0x79,
    0x82, 0x8a, 0x92, 0x9a, 0xa2, 0xaa, 0xb2, 0xba,
    0xc3, 0xcb, 0xd3, 0xdb, 0xe3, 0xeb, 0xf3, 0xfb,
    0x04, 0x0c, 0x14, 0x1c, 0x24, 0x2c, 0x34, 0x3c,
    0x45, 0x4d, 0x55, 0x5d, 0x65, 0x6d, 0x75, 0x7d,
    0x86, 0x8e, 0x96, 0x9e, 0xa6, 0xae, 0xb6, 0xbe,
    0xc7, 0xcf, 0xd7, 0xdf, 0xe7, 0xef, 0xf7, 0xff,
    0x08, 0x00, 0x18, 0x10, 0x28, 0x20, 0x38, 0x30,
    0x49, 0x41, 0x59, 0x51, 0x69, 0x61, 0x79, 0x71,
    0x8a, 0x82, 0x9a, 0x92, 0xaa, 0xa2, 0xba, 0xb2,
    0xcb, 0xc3, 0xdb, 0xd3, 0xeb, 0xe3, 0xfb, 0xf3,
    0x0c, 0x04, 0x1c, 0x14, 0x2c, 0x24, 0x3c, 0x34,
    0x4d, 0x45, 0x5d, 0x55, 0x6d, 0x65, 0x7d, 0x75,
    0x8e, 0x86, 0x9e, 0x96, 0xae, 0xa6, 0xbe, 0xb6,
    0xcf, 0xc7, 0xdf, 0xd7, 0xef, 0xe7, 0xff, 0xf7,
    0x10, 0x18, 0x00, 0x08, 0x30, 0x38, 0x20, 0x28,
    0x51, 0x59, 0x41, 0x49, 0x71, 0x79, 0x61, 0x69,
    0x92, 0x9a, 0x82, 0x8a, 0xb2, 0xba, 0xa2, 0xaa,
    0xd3, 0xdb, 0xc3, 0xcb, 0xf3, 0xfb, 0xe3, 0xeb,
    0x14, 0x1c, 0x04, 0x0c, 0x34, 0x3c, 0x24, 0x2c,
    0x55, 0x5d, 0x45, 0x4d, 0x75, 0x7d, 0x65, 0x6d,
    0x96, 0x9e, 0x86, 0x8e, 0xb6, 0xbe, 0xa6, 0xae,
    0xd7, 0xdf, 0xc7, 0xcf, 0xf7, 0xff, 0xe7, 0xef,
    0x18, 0x10, 0x08, 0x00, 0x38, 0x30, 0x28, 0x20,
    0x59, 0x51, 0x49, 0x41, 0x79, 0x71, 0x69, 0x61,
    0x9a, 0x92, 0x8a, 0x82, 0xba, 0xb2, 0xaa, 0xa2,
    0xdb, 0xd3, 0xcb, 0xc3, 0xfb, 0xf3, 0xeb, 0xe3,
    0x1c, 0x14, 0x0c, 0x04, 0x3c, 0x34, 0x2c, 0x24,
    0x5d, 0x55, 0x4d, 0x45, 0x7d, 0x75, 0x6d, 0x65,
    0x9e, 0x96, 0x8e, 0x86, 0xbe, 0xb6, 0xae, 0xa6,
    0xdf, 0xd7, 0xcf, 0xc7, 0xff, 0xf7, 0xef, 0xe7
};

static void _multiply_M2(const uint8_t old[LANE_BYTES], uint8_t new[LANE_BYTES])
{
    new[7] = old[5];
    new[6] = old[4]    ^ old[5]<<3;
    new[5] = old[3]    ^ M_1[old[4]] ^ old[5]<<6;
    new[4] = old[2]    ^ old[3]>>3   ^ old[4]>>6;
    new[3] = old[6]<<2 ^ old[1];
    new[2] = old[5]<<2 ^ old[0];
    new[1] = old[7];
    new[0] = old[6];
}

static void _multiply_M3(const uint8_t old[LANE_BYTES], uint8_t new[LANE_BYTES])
{
    uint8_t tmp[LANE_BYTES];
    memcpy(tmp, old, LANE_BYTES);

    _multiply_M2(old, tmp);
    _multiply_M(tmp, new);
}

static void _multiply_MR(const uint8_t old[LANE_BYTES], uint8_t new[LANE_BYTES])
{
    new[0] = old[1];
    new[1] = old[2];
    new[2] = old[3] ^ old[4]<<3;
    new[3] = old[4];
    new[4] = old[5] ^ old[6]>>3;
    new[5] = old[6] ^ old[3]>>2;
    new[6] = old[7];
    new[7] = old[0];
}

static void _multiply_MR2(const uint8_t old[LANE_BYTES], uint8_t new[LANE_BYTES])
{
    uint8_t tmp[LANE_BYTES];
    memcpy(tmp, old, LANE_BYTES);

    _multiply_MR(old, tmp);
    _multiply_MR(tmp, new);
}

static void _multiply_MR3(const uint8_t old[LANE_BYTES], uint8_t new[LANE_BYTES])
{
    uint8_t tmp[LANE_BYTES];
    memcpy(tmp, old, LANE_BYTES);

    _multiply_MR2(old, tmp);
    _multiply_MR(tmp, new);
}

typedef void (*matrix_multiplication)(const uint8_t old[LANE_BYTES], uint8_t new[LANE_BYTES]);

static const matrix_multiplication ALPHAS[6] = {
    _multiply_M,
    _multiply_M2,
    _multiply_M3,
    _multiply_MR,
    _multiply_MR2,
    _multiply_MR3
};


void tweakey_state_update(uint8_t TK[TWEAKEY_BYTES])
{
    /* Skip lane 0, as it is multiplied by the identity matrix. */

    for (size_t j=1; j<LANES_NB; j++)
    {
        uint8_t *TKj = TK + j*LANE_BYTES;

        uint8_t TKj_old[LANE_BYTES];
        memcpy(TKj_old, TKj, LANE_BYTES);

        ALPHAS[j-1](TKj_old, TKj);
    }
}