/* This file is in the public domain. */

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

#include "algs.h"
#include "panic.h"
#include "verbose.h"
#include "alg-util.h"

typedef struct rj_state RJ_STATE;

#define MAXBLOCK 32 /* max supported blocksize, in bytes */

struct rj_state {
  void *handle;
  int block;
  unsigned char iv[MAXBLOCK];
  int direction;
#define RJ_DIR_ENC 1
#define RJ_DIR_DEC 2
  } ;

static void *enc_rj_cbc_init(int Nk, int Nb, int dir, const void *key, const void *iv)
{
 RJ_STATE *s;
 void *h;

 s = malloc(sizeof(RJ_STATE));
 if (s == 0) return(0);
 h = rijndael_init(Nk,Nb);
 if (h == 0)
  { free(s);
    return(0);
  }
 rijndael_set_key(h,key);
 s->handle = h;
 s->block = Nb * 4;
 bcopy(iv,&s->iv[0],Nb*4);
 s->direction = dir;
 return(s);
}

#define RJ_CBC(Nk,Nb) \
	static void *enc_rj_k##Nk##b##Nb##_cbc_init_enc(const void *key, const void *iv)\
	{ return(enc_rj_cbc_init(Nk,Nb,RJ_DIR_ENC,key,iv)); }\
	static void *enc_rj_k##Nk##b##Nb##_cbc_init_dec(const void *key, const void *iv)\
	{ return(enc_rj_cbc_init(Nk,Nb,RJ_DIR_DEC,key,iv)); }

RJ_CBC(4,4)
RJ_CBC(4,6)
RJ_CBC(4,8)
RJ_CBC(6,4)
RJ_CBC(6,6)
RJ_CBC(6,8)
RJ_CBC(8,4)
RJ_CBC(8,6)
RJ_CBC(8,8)

#undef RJ_CBC

static void enc_rj_cbc_process(void *state, const void *in, void *out, int datalen)
{
 RJ_STATE *s;
 const unsigned char *ip;
 unsigned char *op;

 s = state;
 ip = in;
 op = out;
 if (datalen % s->block) panic("misaligned length");
 switch (s->direction)
  { case RJ_DIR_ENC:
       while (datalen > 0)
	{ xor_block(&s->iv[0],ip,&s->iv[0],s->block);
	  rijndael_encrypt(s->handle,&s->iv[0],&s->iv[0]);
	  bcopy(&s->iv[0],op,s->block);
	  ip += s->block;
	  op += s->block;
	  datalen -= s->block;
	}
       break;
    case RJ_DIR_DEC:
       while (datalen > 0)
	{ unsigned char ib[MAXBLOCK];
	  bcopy(ip,&ib[0],s->block);
	  rijndael_decrypt(s->handle,ip,op);
	  xor_block(&s->iv[0],op,op,s->block);
	  bcopy(ib,&s->iv[0],s->block);
	  ip += s->block;
	  op += s->block;
	  datalen -= s->block;
	}
       break;
    default:
       panic("bad direction");
       break;
  }
}

static void *enc_rj_ctr_init(int Nk, int Nb, const void *key, const void *iv)
{
 RJ_STATE *s;
 void *h;

 s = malloc(sizeof(RJ_STATE));
 if (s == 0) return(0);
 h = rijndael_init(Nk,Nb);
 if (h == 0)
  { free(s);
    return(0);
  }
 rijndael_set_key(h,key);
 s->handle = h;
 s->block = Nb * 4;
 bcopy(iv,&s->iv[0],Nb*4);
 return(s);
}

#define RJ_CTR(Nk,Nb) \
	static void *enc_rj_k##Nk##b##Nb##_ctr_init(const void *key, const void *iv)\
	{ return(enc_rj_ctr_init(Nk,Nb,key,iv)); }

RJ_CTR(4,4)
RJ_CTR(4,6)
RJ_CTR(4,8)
RJ_CTR(6,4)
RJ_CTR(6,6)
RJ_CTR(6,8)
RJ_CTR(8,4)
RJ_CTR(8,6)
RJ_CTR(8,8)

#undef RJ_CTR

static void enc_rj_ctr_process(void *state, const void *in, void *out, int datalen)
{
 RJ_STATE *s;
 const unsigned char *ip;
 unsigned char *op;
 unsigned char *x;
 unsigned char xorbuf[32];
 int left;
 int i;

 s = state;
 ip = in;
 op = out;
 x = &s->iv[0];
 if (datalen % s->block) panic("misaligned length");
 if (VERB(CRYPTO))
  { int i;
    verb(CRYPTO,"rj ctr %p\n in",(void *)s);
    for (i=0;i<datalen;i++) verb(CRYPTO,"%s%02x",(i%s->block)?"":" ",ip[i]);
    verb(CRYPTO,"\n  x ");
    for (i=0;i<s->block;i++) verb(CRYPTO,"%02x",x[i]);
    verb(CRYPTO,"\n");
  }
 for (left=datalen;left>0;left-=s->block)
  { rijndael_encrypt(s->handle,&x[0],&xorbuf[0]);
    xor_block(&xorbuf[0],ip,op,s->block);
    i = s->block - 1;
    while (1)
     { if (i < 0) break;
       if (x[i] != 0xff)
	{ x[i] ++;
	  break;
	}
       x[i] = 0;
       i --;
     }
    ip += s->block;
    op += s->block;
  }
 if (VERB(CRYPTO))
  { int i;
    verb(CRYPTO,"out");
    for (i=0;i<datalen;i++) verb(CRYPTO,"%s%02x",(i&15)?"":" ",op[i-datalen]);
    verb(CRYPTO,"\n");
  }
}

static void enc_rj_done(void *state)
{
 rijndael_done(((RJ_STATE *)state)->handle);
 free(state);
}

ENCALG encalg_aes128_cbc
 = { "aes128-cbc", 0, 16, 16, 16, 0,
     &enc_rj_k4b4_cbc_init_enc,
     &enc_rj_k4b4_cbc_init_dec,
     &enc_rj_cbc_process,
     &enc_rj_done };

ENCALG encalg_aes192_cbc
 = { "aes192-cbc", 0, 16, 16, 24, 0,
     &enc_rj_k6b4_cbc_init_enc,
     &enc_rj_k6b4_cbc_init_dec,
     &enc_rj_cbc_process,
     &enc_rj_done };

ENCALG encalg_aes256_cbc
 = { "aes256-cbc", 0, 16, 16, 32, 0,
     &enc_rj_k8b4_cbc_init_enc,
     &enc_rj_k8b4_cbc_init_dec,
     &enc_rj_cbc_process,
     &enc_rj_done };

ENCALG encalg_rijndael_cbc_lysator
 = { "rijndael-cbc@lysator.liu.se", "same as aes256-cbc", 16, 16, 32, 0,
     &enc_rj_k8b4_cbc_init_enc,
     &enc_rj_k8b4_cbc_init_dec,
     &enc_rj_cbc_process,
     &enc_rj_done };

ENCALG encalg_aes128_ctr
 = { "aes128-ctr", 0, 16, 16, 16, 0,
     &enc_rj_k4b4_ctr_init,
     &enc_rj_k4b4_ctr_init,
     &enc_rj_ctr_process,
     &enc_rj_done };

ENCALG encalg_aes192_ctr
 = { "aes192-ctr", 0, 16, 16, 24, 0,
     &enc_rj_k6b4_ctr_init,
     &enc_rj_k6b4_ctr_init,
     &enc_rj_ctr_process,
     &enc_rj_done };

ENCALG encalg_aes256_ctr
 = { "aes256-ctr", 0, 16, 16, 32, 0,
     &enc_rj_k8b4_ctr_init,
     &enc_rj_k8b4_ctr_init,
     &enc_rj_ctr_process,
     &enc_rj_done };

#define RJ_ALGS(Nk,Nb,comment_cbc,comment_ctr) \
	ENCALG encalg_rijndael_rodents_k##Nk##b##Nb##_cbc		\
	 = { "rijndael-k" #Nk "b" #Nb "-cbc@rodents.montreal.qc.ca",	\
	     comment_cbc, Nb*4, Nb*4, Nk*4, 0,				\
	     &enc_rj_k##Nk##b##Nb##_cbc_init_enc,			\
	     &enc_rj_k##Nk##b##Nb##_cbc_init_dec,			\
	     &enc_rj_cbc_process,					\
	     &enc_rj_done };						\
	ENCALG encalg_rijndael_rodents_k##Nk##b##Nb##_ctr		\
	 = { "rijndael-k" #Nk "b" #Nb "-ctr@rodents.montreal.qc.ca",	\
	     comment_ctr, Nb*4, Nb*4, Nk*4, 0,				\
	     &enc_rj_k##Nk##b##Nb##_ctr_init,				\
	     &enc_rj_k##Nk##b##Nb##_ctr_init,				\
	     &enc_rj_ctr_process,					\
	     &enc_rj_done };						\

RJ_ALGS(4,4,"same as aes128-cbc","same as aes128-ctr")
RJ_ALGS(4,6,0,0)
RJ_ALGS(4,8,0,0)
RJ_ALGS(6,4,"same as aes192-cbc","same as aes192-ctr")
RJ_ALGS(6,6,0,0)
RJ_ALGS(6,8,0,0)
RJ_ALGS(8,4,"same as aes256-cbc","same as aes256-ctr")
RJ_ALGS(8,6,0,0)
RJ_ALGS(8,8,0,0)

#undef RJ_ALGS