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

typedef uint8_t block_t[4][4];

typedef struct aes_ctx {
  uint8_t e_key[4 * (14 + 1)][4];
  uint32_t Nk;
  uint32_t Nr;
} aes_ctx;

// multiply by x modulo x^8 + x^4 + x^3 + x + 1
static inline uint8_t xtime(uint8_t b) {
  return (b << 1) ^ ((b & 0x80) ? 0x1b : 0);
}

// GF(2^8) generic multiplication, double-and-add
static inline uint8_t multiply(uint8_t a, uint8_t b) {
  uint8_t c = 0;
  for(size_t i = 0; i < 8; ++i) {
    if((b >> i) & 1)
      c ^= a;
    a = xtime(a);
  }
  return c;
}


static inline uint8_t rotate_left(uint8_t x, size_t c) {
  return (x << c % 8) | (x >> (8 - c) % 8);
}

static inline uint8_t SubByte(const uint8_t x0) {
  const uint8_t  x1 = multiply( x0,  x0); // x^2
  const uint8_t  x2 = multiply( x1,  x0); // x^3
  const uint8_t  x3 = multiply( x2,  x2); // x^6
  const uint8_t  x4 = multiply( x3,  x3); // x^12
  const uint8_t  x5 = multiply( x4,  x2); // x^15
  const uint8_t  x6 = multiply( x5,  x5); // x^30
  const uint8_t  x7 = multiply( x6,  x6); // x^60
  const uint8_t  x8 = multiply( x7,  x2); // x^63
  const uint8_t  x9 = multiply( x8,  x8); // x^126
  const uint8_t x10 = multiply( x9,  x0); // x^127
  const uint8_t x11 = multiply(x10, x10); // x^254 = x^-1
  return x11 ^ rotate_left(x11, 1) ^ rotate_left(x11, 2) ^ 
         rotate_left(x11, 3) ^ rotate_left(x11, 4) ^ 0x63;
}

static inline void RotWord(uint8_t w[4]) {
  const uint8_t t = w[0];
  w[0] = w[1];
  w[1] = w[2];
  w[2] = w[3];
  w[3] = t;
}

static inline void SubWord(uint8_t w[4]) {
  for(size_t i = 0; i < 4; ++i) 
    w[i] = SubByte(w[i]);
}

static inline void AddRoundKey(block_t block, const uint8_t k[][4]) {
  for(size_t i = 0; i < 4; ++i) 
    for(size_t j = 0; j < 4; ++j)
      block[i][j] ^= k[i][j];
}

static inline void SubBytes(block_t block) {
  for(size_t i = 0; i < 4; ++i) 
    for(size_t j = 0; j < 4; ++j)
      block[i][j] = SubByte(block[i][j]);
}

static inline void ShiftRows(block_t block) {
  for(size_t j = 0; j < 4; ++j) {
    uint8_t row[4];
    for(size_t i = 0; i < 4; ++i)
      row[i] = block[i][j];
    for(size_t i = 0; i < 4; ++i)
      block[i][j] = row[(i+j)%4];
  }
}

static inline void MixColumns(block_t block) {
  for(size_t i = 0; i < 4; ++i) {
    const uint8_t c0 = block[i][0];
    const uint8_t c1 = block[i][1];
    const uint8_t c2 = block[i][2];
    const uint8_t c3 = block[i][3];
    block[i][0] = xtime(c0 ^ c1) ^ c1 ^ c2 ^ c3;
    block[i][1] = c0 ^ xtime(c1 ^ c2) ^ c2 ^ c3;
    block[i][2] = c0 ^ c1 ^ xtime(c2 ^ c3) ^ c3;
    block[i][3] = xtime(c0 ^ c3) ^ c0 ^ c1 ^ c2;
  }
}

void aes_set_key(aes_ctx * ctx, uint8_t const * k, size_t keybits) {
  switch(keybits) {
  case 128:
    ctx->Nr = 10;
    ctx->Nk = 4;
    break;
  case 192:
    ctx->Nr = 12;
    ctx->Nk = 6;
    break;
  case 256:
    ctx->Nr = 14;
    ctx->Nk = 8;
    break;
  default:
    abort();
  }
  uint8_t rcon = 1;
  for(size_t i = 0; i < ctx->Nk; ++i)
    for(size_t j = 0; j < 4; ++j)
      ctx->e_key[i][j] = k[4*i+j];
  for(size_t i = ctx->Nk; i < 4 * (ctx->Nr + 1); ++i) {
    uint8_t temp[4];
    for(size_t j = 0; j < 4; ++j)
      temp[j] = ctx->e_key[i-1][j];
    if(i % ctx->Nk == 0) {
      RotWord(temp);
      SubWord(temp);
      temp[0] ^= rcon;
      rcon = xtime(rcon);
    } else if (ctx->Nk > 6 && i % ctx->Nk == 4) {
      SubWord(temp);
    }
    for(size_t j = 0; j < 4; ++j) 
      ctx->e_key[i][j] = ctx->e_key[i-ctx->Nk][j] ^ temp[j];
  }
}

void aes_encrypt(aes_ctx * ctx, uint8_t * output, uint8_t const * input) {
  block_t block;
  for(size_t i = 0; i < 4; ++i)
    for(size_t j = 0; j < 4; ++j)
      block[i][j] = input[i * 4 + j];
  AddRoundKey(block, &ctx->e_key[0]);
  for(size_t i = 1; i <= ctx->Nr - 1; ++i) {
    SubBytes(block);
    ShiftRows(block);
    MixColumns(block);
    AddRoundKey(block, &ctx->e_key[4 * i]);
  }
  SubBytes(block);
  ShiftRows(block);
  // No MixColumns
  AddRoundKey(block, &ctx->e_key[4 * ctx->Nr]);
  for(size_t i = 0; i < 4; ++i)
    for(size_t j = 0; j < 4; ++j)
      output[i * 4 + j] = block[i][j];
}

typedef struct rc4_ctx {
	uint8_t S[256];
	uint8_t i;
	uint8_t j;
} rc4_ctx;

typedef struct jipsam2_ctx {
	rc4_ctx first;
	rc4_ctx last;
} jipsam2_ctx;

#define SWAP(a, b) { uint8_t t = a; a = b; b = t; }

void jipsam2_key_init(jipsam2_ctx * ctx, uint8_t * out) {
	rc4_ctx * rc4 = &ctx->first;
	for(size_t n = 0; n < 32; ++n) {
		rc4->i += 1;
		rc4->j += rc4->S[rc4->i];
		SWAP(rc4->S[rc4->i], rc4->S[rc4->j]);
		out[n] = rc4->S[(rc4->S[rc4->i] + rc4->S[rc4->j])%256];
	}
}

void jipsam2_key_last(jipsam2_ctx * ctx, uint8_t * out) {
	rc4_ctx * rc4 = &ctx->last;
	for(size_t n = 0; n < 32; ++n) {
		rc4->i += 1;
		rc4->j += rc4->S[rc4->i];
		SWAP(rc4->S[rc4->i], rc4->S[rc4->j]);
		out[n] = rc4->S[(rc4->S[rc4->i] + rc4->S[rc4->j])%256];
	}
}

void jipsam2_set_key(jipsam2_ctx * ctx, uint8_t const * k, size_t klen) {
	uint8_t discard[64];
	memset(ctx, 0, sizeof(*ctx));
	ctx->first.i = ctx->last.i = 0;
	ctx->first.j = ctx->last.j = 0;
	for(size_t i = 0; i < 256; ++i) {
		ctx->first.S[i] = ctx->last.S[i] = i;
	}

  size_t j = 2;
  size_t key_index = 0;
	for(size_t i = 0; i < 256; ++i) {
		j = (j + ctx->first.S[i] + k[key_index]) % 256;
		SWAP(ctx->first.S[i], ctx->first.S[j]);
    key_index += 1;
    if(key_index >= klen / 2)
      key_index = 0;
    key_index += 1;
	}

  // Reusing j from the previous state
  key_index = klen / 2;
	for(size_t i = 0; i < 256; ++i) {
		j = (j + ctx->last.S[i] + k[key_index]) % 256;
		SWAP(ctx->last.S[i], ctx->last.S[j]);
    key_index += 2;
    if(key_index >= klen)
      key_index = klen / 2;
	}
	jipsam2_key_init(ctx, discard);
	jipsam2_key_last(ctx, discard);
}

void jipsam2_encrypt(jipsam2_ctx * ctx, uint8_t * output, const uint8_t * input) {
	uint8_t xor_key[32], aes_key[32];
	aes_ctx aes_ks;
	jipsam2_key_init(ctx, aes_key);
	jipsam2_key_last(ctx, xor_key);
	aes_set_key(&aes_ks, aes_key, 256);
	aes_encrypt(&aes_ks, output, input);
	for(size_t i = 0; i < 32; ++i)
		output[i] ^= xor_key[i];
}

#include <stdio.h>

// No test vectors, we'll make our own
int main() {
  uint8_t k[32];
  uint8_t b[16];

  for(size_t i = 0; i < 16; ++i)
    b[i] = i * 0x11;

  for(size_t i = 0; i < 32; ++i)
    k[i] = i;

  jipsam2_ctx ctx;
  jipsam2_set_key(&ctx, k, 32);

  // Key schedule
  for(size_t i = 0; i < sizeof(ctx); ++i)
    printf("%02x%c", ((unsigned char *)&ctx)[i], i % 16 == 15 ? '\n' : ' ');
  printf("\n");

  jipsam2_encrypt(&ctx, b, b);

  // Encryption
  for(size_t i = 0; i < 16; ++i)
    printf("%02x ", b[i]);
  printf("\n");
  return 0;
}