#include <errno.h>
#include <stdlib.h>
#include <strings.h>

#include <rijndael.h>

#include "tables.h"

typedef struct ctx CTX;

struct ctx {
  int Nk;
  int Nb;
  int Nr;
  int C[3];
  int cn[3];
  int cq[3];
  unsigned char (*S)[4];
  unsigned char (*W)[4];
  } ;

static int std_c[3][3]
 = { { 1, 2, 3 },
     { 1, 2, 3 },
     { 1, 3, 4 } };
static int std_rounds[3][3]
 = { { 10, 12, 14 },
     { 12, 12, 14 },
     { 14, 14, 14 } };

#ifdef DEBUG_OUTPUT
#include <stdio.h>
#include <stdarg.h>
static void dump_m(unsigned char *, int, const char *, ...)
	__attribute__((__format__(__printf__,3,4)));
static void dump_m(unsigned char *m, int w, const char *fmt, ...)
{
 va_list ap;
 int r;
 int c;

 va_start(ap,fmt);
 vprintf(fmt,ap);
 va_end(ap);
 printf("\n");
 for (r=0;r<4;r++)
  { for (c=0;c<w;c++) printf(" %02x",m[r+(4*c)]);
    printf("\n");
  }
}
#else
#define dump_m(x...) (void)0
#endif

static inline unsigned char M(unsigned char a, unsigned char b)
{
 return(rijndael__mult[a][b]);
}

static inline unsigned int gcd(unsigned int a, unsigned int b)
{
 unsigned int t;

 while (a)
  { t = b % a;
    b = a;
    a = t;
  }
 return(b);
}

void *rijndael_init(int keylen, int blocklen)
{
 if ( (keylen < 4) || (blocklen < 4) ||
      (keylen > 8) || (blocklen > 8) ||
      (keylen & 1) || (blocklen & 1) )
  { errno = EINVAL;
    return(0);
  }
 return(rijndael_init_special(keylen,blocklen,&std_c[(blocklen-4)>>1][0],std_rounds[(keylen-4)>>1][(blocklen-4)>>1]));
}

void *rijndael_init_special(int keylen, int blocklen, int *c, int rounds)
{
 int i;
 CTX *x;

 if ((keylen < 4) || (blocklen < 4) || (rounds < 1))
  { errno = EINVAL;
    return(0);
  }
 for (i=3-1;i>=0;i--)
  { if ((c[i] < 0) || (c[i] >= blocklen))
     { errno = EINVAL;
       return(0);
     }
  }
 x = malloc(sizeof(CTX));
 if (x == 0) return(0);
 x->S = malloc(sizeof(*x->S) *
		( blocklen			/* for S */
		+ (blocklen*(rounds+1))		/* for W */
		));
 if (x->S == 0)
  { i = errno;
    free(x);
    errno = i;
    return(0);
  }
 x->W = x->S + blocklen;
 x->Nk = keylen;
 x->Nb = blocklen;
 x->Nr = rounds;
 x->C[0] = c[0];
 x->C[1] = c[1];
 x->C[2] = c[2];
 for (i=0;i<3;i++)
  { x->cn[i] = gcd(c[i],blocklen);
    x->cq[i] = blocklen / x->cn[i];
  }
 return(x);
}

void rijndael_done(void *xv)
{
 CTX *x;

 x = xv;
 bzero(x->S,x->Nb*sizeof(*x->S));
 bzero(x->W,x->Nb*(x->Nr+1)*sizeof(*x->W));
 free(x->S);
 bzero(x,sizeof(*x));
 free(x);
}

#define ByteSub(x) rijndael__sbox[(unsigned char)(x)]
#define InvByteSub(x) rijndael__isbox[(unsigned char)(x)]

static void get_W(CTX *x)
{
 int i;
 int n;
 unsigned char t[4];
 int j;
 unsigned char rcon;
 int Nk;

 Nk = x->Nk;
 n = x->Nb * (x->Nr + 1);
 j = 0;
 rcon = 0x01;
 for (i=Nk;i<n;i++)
  { if (j == 0)
     { t[0] = ByteSub(x->W[i-1][1]) ^ rcon;
       t[1] = ByteSub(x->W[i-1][2]);
       t[2] = ByteSub(x->W[i-1][3]);
       t[3] = ByteSub(x->W[i-1][0]);
     }
    else if ((x->Nk > 6) && (j == 4))
     { t[0] = ByteSub(x->W[i-1][0]);
       t[1] = ByteSub(x->W[i-1][1]);
       t[2] = ByteSub(x->W[i-1][2]);
       t[3] = ByteSub(x->W[i-1][3]);
     }
    else
     { bcopy(&x->W[i-1][0],&t[0],4);
     }
    x->W[i][0] = x->W[i-Nk][0] ^ t[0];
    x->W[i][1] = x->W[i-Nk][1] ^ t[1];
    x->W[i][2] = x->W[i-Nk][2] ^ t[2];
    x->W[i][3] = x->W[i-Nk][3] ^ t[3];
    if (++j >= Nk)
     { j = 0;
       rcon = M(2,rcon);
     }
  }
}

void rijndael_set_key(void *xv, const void *key)
{
 CTX *x;

 x = xv;
 /* Convenient, how arrays are defined. */
 bcopy(key,x->W,x->Nk*4);
 get_W(x);
}

static void shift_row(unsigned char (*S)[4], int Nb, int o, int cn, int cq, int n)
{
 unsigned char t;
 int i;
 int j;
 int y;
 int x;

 for (i=cn-1;i>=0;i--)
  { t = S[i][o];
    x = i;
    for (j=cq-1;j>0;j--)
     { y = x;
       if ((x+=n) >= Nb) x -= Nb;
       S[y][o] = S[x][o];
     }
    S[x][o] = t;
  }
}

static void inv_shift_row(unsigned char (*S)[4], int Nb, int o, int cn, int cq, int n)
{
 unsigned char t;
 int i;
 int j;
 int y;
 int x;

 for (i=cn-1;i>=0;i--)
  { t = S[i][o];
    x = i;
    for (j=cq-1;j>0;j--)
     { y = x;
       if ((x-=n) < 0) x += Nb;
       S[y][o] = S[x][o];
     }
    S[x][o] = t;
  }
}

static inline void xor2N(unsigned char *d, const unsigned char *s, int n)
{
 for (;n>0;n--) *d++ ^= *s++;
}

static inline void ByteSubN(unsigned char *d, int n)
{
 for (;n>0;n--,d++) *d = ByteSub(*d);
}

static inline void InvByteSubN(unsigned char *d, int n)
{
 for (;n>0;n--,d++) *d = InvByteSub(*d);
}

static inline void MixColumn(unsigned char *c)
{
 unsigned char r[4];

 r[0] = M(2,c[0]) ^ M(3,c[1]) ^ c[2] ^ c[3];
 r[1] = c[0] ^ M(2,c[1]) ^ M(3,c[2]) ^ c[3];
 r[2] = c[0] ^ c[1] ^ M(2,c[2]) ^ M(3,c[3]);
 c[3] = M(3,c[0]) ^ c[1] ^ c[2] ^ M(2,c[3]);
 c[2] = r[2];
 c[1] = r[1];
 c[0] = r[0];
}

static inline void InvMixColumn(unsigned char *c)
{
 unsigned char r[4];

 r[0] = M(0x0e,c[0]) ^ M(0x0b,c[1]) ^ M(0x0d,c[2]) ^ M(0x09,c[3]);
 r[1] = M(0x09,c[0]) ^ M(0x0e,c[1]) ^ M(0x0b,c[2]) ^ M(0x0d,c[3]);
 r[2] = M(0x0d,c[0]) ^ M(0x09,c[1]) ^ M(0x0e,c[2]) ^ M(0x0b,c[3]);
 c[3] = M(0x0b,c[0]) ^ M(0x0d,c[1]) ^ M(0x09,c[2]) ^ M(0x0e,c[3]);
 c[2] = r[2];
 c[1] = r[1];
 c[0] = r[0];
}

void rijndael_encrypt(const void *xv, const void *in, void *out)
{
 const CTX *x;
 int i;
 int j;
 unsigned char (*Wp)[4];

 x = xv;
 Wp = &x->W[0];
  { unsigned char S[x->Nb][4];
    bcopy(in,&S[0][0],4*x->Nb);
    dump_m(&S[0][0],x->Nb,"Initial input");
    /* Initial AddRoundKey */
    xor2N(&S[0][0],&Wp[0][0],4*x->Nb);
    dump_m(&Wp[0][0],x->Nb,"Initial round key");
    Wp += x->Nb;
    for (i=x->Nr-1;i>0;i--)
     { /* ByteSub */
       dump_m(&S[0][0],x->Nb,"Beginning of round %d",x->Nr-i);
       ByteSubN(&S[0][0],4*x->Nb);
       dump_m(&S[0][0],x->Nb,"After ByteSub");
       /* ShiftRow */
       shift_row(&S[0],x->Nb,1,x->cn[0],x->cq[0],x->C[0]);
       shift_row(&S[0],x->Nb,2,x->cn[1],x->cq[1],x->C[1]);
       shift_row(&S[0],x->Nb,3,x->cn[2],x->cq[2],x->C[2]);
       dump_m(&S[0][0],x->Nb,"After ShiftRow");
       /* MixColumn */
       for (j=x->Nb-1;j>=0;j--) MixColumn(&S[j][0]);
       dump_m(&S[0][0],x->Nb,"After MixColumn");
       /* AddRoundKey */
       xor2N(&S[0][0],&Wp[0][0],4*x->Nb);
       dump_m(&Wp[0][0],x->Nb,"Round key");
       Wp += x->Nb;
     }
    dump_m(&S[0][0],x->Nb,"Beginning of final round");
    /* Final ByteSub */
    ByteSubN(&S[0][0],4*x->Nb);
    dump_m(&S[0][0],x->Nb,"After ByteSub");
    /* Final ShiftRow */
    shift_row(&S[0],x->Nb,1,x->cn[0],x->cq[0],x->C[0]);
    shift_row(&S[0],x->Nb,2,x->cn[1],x->cq[1],x->C[1]);
    shift_row(&S[0],x->Nb,3,x->cn[2],x->cq[2],x->C[2]);
    dump_m(&S[0][0],x->Nb,"After ShiftRow");
    /* No MixColumn in final round */
    /* Final AddRoundKey */
    xor2N(&S[0][0],&Wp[0][0],4*x->Nb);
    dump_m(&Wp[0][0],x->Nb,"Round key");
    dump_m(&S[0][0],x->Nb,"Output");
    bcopy(&S[0][0],out,4*x->Nb);
  }
}

void rijndael_decrypt(const void *xv, const void *in, void *out)
{
 const CTX *x;
 int i;
 int j;
 unsigned char (*Wp)[4];

 x = xv;
 Wp = &x->W[x->Nb*(x->Nr+1)];
  { unsigned char S[x->Nb][4];
    bcopy(in,&S[0][0],4*x->Nb);
    dump_m(&S[0][0],x->Nb,"Initial input");
    /* Initial AddRoundKey */
    Wp -= x->Nb;
    xor2N(&S[0][0],&Wp[0][0],4*x->Nb);
    dump_m(&Wp[0][0],x->Nb,"Initial round key");
    for (i=x->Nr-1;i>0;i--)
     { /* inverse ShiftRow */
       dump_m(&S[0][0],x->Nb,"Beginning of round %d",x->Nr-i);
       inv_shift_row(&S[0],x->Nb,1,x->cn[0],x->cq[0],x->C[0]);
       inv_shift_row(&S[0],x->Nb,2,x->cn[1],x->cq[1],x->C[1]);
       inv_shift_row(&S[0],x->Nb,3,x->cn[2],x->cq[2],x->C[2]);
       dump_m(&S[0][0],x->Nb,"After inverse ShiftRow");
       /* inverse ByteSub */
       InvByteSubN(&S[0][0],4*x->Nb);
       dump_m(&S[0][0],x->Nb,"After inverse ByteSub");
       /* AddRoundKey */
       Wp -= x->Nb;
       xor2N(&S[0][0],&Wp[0][0],4*x->Nb);
       dump_m(&Wp[0][0],x->Nb,"Round key");
       /* inverse MixColumn */
       for (j=x->Nb-1;j>=0;j--) InvMixColumn(&S[j][0]);
       dump_m(&S[0][0],x->Nb,"After inverse MixColumn");
     }
    dump_m(&S[0][0],x->Nb,"Beginning of final round");
    /* Final inverse ShiftRow */
    inv_shift_row(&S[0],x->Nb,1,x->cn[0],x->cq[0],x->C[0]);
    inv_shift_row(&S[0],x->Nb,2,x->cn[1],x->cq[1],x->C[1]);
    inv_shift_row(&S[0],x->Nb,3,x->cn[2],x->cq[2],x->C[2]);
    dump_m(&S[0][0],x->Nb,"After inverse ShiftRow");
    /* Final inverse ByteSub */
    InvByteSubN(&S[0][0],4*x->Nb);
    dump_m(&S[0][0],x->Nb,"After inverse ByteSub");
    /* No MixColumn in final round */
    /* Final AddRoundKey */
    Wp -= x->Nb;
    xor2N(&S[0][0],&Wp[0][0],4*x->Nb);
    dump_m(&Wp[0][0],x->Nb,"Round key");
    dump_m(&S[0][0],x->Nb,"Output");
    bcopy(&S[0][0],out,4*x->Nb);
  }
}