// I'm aware that this file is written horribly and dooesn't have inverse functions

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

#define LROT8(v, b) ((v) << (b) | (v) >> (8 - (b)))
#define LROT32(v, b) ((v) << (b) | (v) >> (32 - (b)))

char sbox[256], inv_sbox[256];

struct polydiv_result
{
  uint16_t q;
  uint16_t r;
};

struct polyeeuclid_result
{
  uint16_t u;
  uint16_t v;
  uint16_t g;
};

static uint16_t msb(uint16_t x)
{
  x |= x >> 1;
  x |= x >> 2;
  x |= x >> 4;
  x |= x >> 8;
  x ^= (x >> 1);
  return x;
}

static uint16_t poly_mult(uint16_t multiplicand, uint16_t multiplier)
{
  uint16_t accum = 0;
  int ii;
  for (ii = 0; ii < 16; ii++)
  {
    if (multiplier & 1)
      accum ^= multiplicand;
    multiplier >>= 1;
    multiplicand <<= 1;
  }
  return accum;
}

static void poly_div(uint16_t dividend, uint16_t divisor, struct polydiv_result *out)
{
  uint16_t quotient = 0;
  if (divisor == 0)
    return;
  while (msb(dividend) >= msb(divisor))
  {
    uint16_t intermittent = msb(dividend) / msb(divisor);
    quotient ^= intermittent;
    dividend ^= divisor * intermittent;
  }
  out->q = quotient;
  out->r = dividend;
}

static uint16_t poly_mult_mod(uint16_t a, uint16_t b, uint16_t m)
{
  struct polydiv_result qr;
  poly_div(poly_mult(a, b), m, &qr);
  return qr.r;
}

static void poly_eeuclid(uint16_t a, uint16_t b, struct polyeeuclid_result *out)
{
  struct polydiv_result qr;
  uint16_t x0 = 1, y0 = 0, x1 = 0, y1 = 1;
  poly_div(a, b, &qr);
  while (qr.r)
  {
    uint16_t xn = x0 ^ poly_mult(qr.q, x1), yn = y0 ^ poly_mult(qr.q, y1);
    a = b;
    b = qr.r;
    poly_div(a, b, &qr);
    x0 = x1;
    y0 = y1;
    x1 = xn;
    y1 = yn;
  }

  out->g = b;
  out->u = x1;
  out->v = y1;
}

//void print_poly(uint16_t p)
//{
//  int ii;
//  for (ii = 0; ii < 16; ii++)
//  {
//    printf((p & 32768) ? "1" : "0");
//    p <<= 1;
//  }
//  printf("\n");
//}

static uint8_t inv_Rijndael(uint8_t p)
{
  if (!p)
    return 0;
  struct polyeeuclid_result guv;
  poly_eeuclid(p, 0x11b, &guv);
  return (uint8_t) guv.u;
}

static uint8_t Nyberg(uint8_t i)
{
  return i ^ LROT8(i, 1) ^ LROT8(i, 2) ^ LROT8(i, 3) ^ LROT8(i, 4) ^ 0x63;
}

static uint8_t sbox_transform(uint8_t i)
{
  return Nyberg(inv_Rijndael(i));
}

void initializeSBox()
{
  int ii;
  for (ii = 0; ii < 256; ii++)
  {
    uint8_t in = (uint8_t) ii;
    uint8_t out = sbox_transform(ii);
    sbox[in] = out;
    inv_sbox[out] = in;
  }
}

static uint32_t wordSBox(uint32_t word)
{
  return
    (sbox[(uint8_t) word]) |
    (sbox[(uint8_t) (word >> 8)] << 8) |
    (sbox[(uint8_t) (word >> 16)] << 16) |
    (sbox[(uint8_t) (word >> 24)] << 24);
}

void expandKey(uint32_t short_key[8], uint32_t expanded_key[60])
{
  // TODO: Make the round constants constant
  uint8_t rcs[8] = {1};
  int ii;
  for (ii = 1; ii < 8; ii++)
  {
    if (rcs[ii - 1] < 0x80)
    {
      rcs[ii] = rcs[ii - 1] << 1;
    } else
    {
      rcs[ii] = 0x1b ^ (rcs[ii - 1] << 1);
    }
  }

  memcpy(expanded_key, short_key, 8 * sizeof(uint32_t));

  int rci = 0;

  for (ii = 8; ii < 60; ii++)
  {
    if (ii % 8 == 0)
    {
      expanded_key[ii] = expanded_key[ii - 8] ^ wordSBox(LROT32(expanded_key[ii - 1], 8)) ^ (rcs[rci] << 24);
      rci++;
    } else if (ii % 8 == 4)
    {
      expanded_key[ii] = expanded_key[ii - 8] ^ wordSBox(expanded_key[ii - 1]);
    } else
    {
      expanded_key[ii] = expanded_key[ii - 8] ^ expanded_key[ii - 1];
    }
  }
}

void AESRound(uint32_t state[4], uint32_t key[60])
{
  int ii, ij;
  for (ii = 0; ii < 4; ii++)
  {
    state[ii] ^= *key;
    key++;
  }
  for (ij = 0; ij < 13; ij++)
  {
    for (ii = 0; ii < 4; ii++)
      state[ii] = wordSBox(state[ii]);
    for (ii = 0; ii < 4; ii++)
      state[ii] = LROT32(state[ii], 8 * ii);
    uint32_t newstate[4] = {0};
    for (ii = 0; ii < 4; ii++)
    {
      int shc = 8 * ii;
      newstate[0] |= (poly_mult_mod((uint8_t) (state[0] >> shc), 2, 0x11b) ^
          poly_mult_mod((uint8_t) (state[1] >> shc), 3, 0x11b) ^ 
          ((uint8_t) (state[2] >> shc)) ^ ((uint8_t) (state[3] >> shc))) << shc;
      newstate[1] |= (poly_mult_mod((uint8_t) (state[1] >> shc), 2, 0x11b) ^
          poly_mult_mod((uint8_t) (state[2] >> shc), 3, 0x11b) ^ 
          ((uint8_t) (state[3] >> shc)) ^ ((uint8_t) (state[0] >> shc))) << shc;
      newstate[2] |= (poly_mult_mod((uint8_t) (state[2] >> shc), 2, 0x11b) ^
          poly_mult_mod((uint8_t) (state[3] >> shc), 3, 0x11b) ^ 
          ((uint8_t) (state[0] >> shc)) ^ ((uint8_t) (state[1] >> shc))) << shc;
      newstate[3] |= (poly_mult_mod((uint8_t) (state[3] >> shc), 2, 0x11b) ^
          poly_mult_mod((uint8_t) (state[0] >> shc), 3, 0x11b) ^ 
          ((uint8_t) (state[1] >> shc)) ^ ((uint8_t) (state[2] >> shc))) << shc;
    }
    memcpy(state, newstate, 4 * sizeof(uint32_t));
    for (ii = 0; ii < 4; ii++)
    {
      state[ii] ^= *key;
      key++;
    }
  }
  for (ii = 0; ii < 4; ii++)
    state[ii] = wordSBox(state[ii]);
  for (ii = 0; ii < 4; ii++)
    state[ii] = LROT32(state[ii], 8 * ii);
  for (ii = 0; ii < 4; ii++)
  {
    state[ii] ^= *key;
    key++;
  }
}

void dataToState(char data[16], uint32_t state[4])
{
  int ii;
  for (ii = 0; ii < 4; ii++)
  {
    state[ii] = data[0] | (data[1] << 8) | (data[2] << 16) | (data[3] << 24);
    data += 4;
  }
}

void stateToData(uint32_t state[4], char data[16])
{
  int ii, ij;;
  for (ii = 0; ii < 4; ii++)
  {
    uint32_t d = state[ii];
    for (ij = 0; ij < 4; ij++)
    {
      *data = d;
      data++;
      d >>= 8;
    }
  }
}

void AESRound_data(char data[16], uint32_t key[60])
{
  uint32_t state[4];
  dataToState(data, state);
  AESRound(state, key);
  stateToData(state, data);
}


//int main()
//{
//  int ii;
//  //for (ii = 0; ii < 256; ii++)
//  //{
//  //  printf("%02x %02x %02x\n", ii & 0xf0, ii & 0x0f, sbox_transform(ii));
//  //}
//
//  uint8_t state[4] = {0xf2, 0x0a, 0x22, 0x5c};
//  uint8_t newstate[4] = {0};
//  int shc = 0;
//  newstate[0] |= (poly_mult_mod((uint8_t) (state[0] >> shc), 2, 0x11b) ^
//      poly_mult_mod((uint8_t) (state[1] >> shc), 3, 0x11b) ^ 
//      ((uint8_t) (state[2] >> shc)) ^ ((uint8_t) (state[3] >> shc))) << shc;
//  newstate[1] |= (poly_mult_mod((uint8_t) (state[1] >> shc), 2, 0x11b) ^
//      poly_mult_mod((uint8_t) (state[2] >> shc), 3, 0x11b) ^ 
//      ((uint8_t) (state[3] >> shc)) ^ ((uint8_t) (state[0] >> shc))) << shc;
//  newstate[2] |= (poly_mult_mod((uint8_t) (state[2] >> shc), 2, 0x11b) ^
//      poly_mult_mod((uint8_t) (state[3] >> shc), 3, 0x11b) ^ 
//      ((uint8_t) (state[0] >> shc)) ^ ((uint8_t) (state[1] >> shc))) << shc;
//  newstate[3] |= (poly_mult_mod((uint8_t) (state[3] >> shc), 2, 0x11b) ^
//      poly_mult_mod((uint8_t) (state[0] >> shc), 3, 0x11b) ^ 
//      ((uint8_t) (state[1] >> shc)) ^ ((uint8_t) (state[2] >> shc))) << shc;
//
//  for (ii = 0; ii < 4; ii++)
//  {
//    printf("%02x ", newstate[ii]);
//  }
//  printf("\n");
//  return 0;
//}
//