/* nanoka-imm.c - Random block cipher - Don Yang (uguu.org)

   01/06/05
                                                                           */
/*@ -unqualifiedtrans -usedef @*/

#include<assert.h>
#include<stdio.h>
#include<stdlib.h>
#include<string.h>
#include<time.h>
#include"nanoka-core.h"


#ifdef BIG_ENDIAN
   static void Bswap(unsigned int *d, unsigned int blockcount);
#endif

static void BlockRead(FILE *infile, unsigned int *d, unsigned int blocksize);
static void BlockWrite(FILE *outfile, unsigned int *d, unsigned int blocksize);
static int Decrypt(Cipher *c, FILE *infile, FILE *outfile);
static void DecryptBlock(Cipher *c, unsigned int *d);
static int Encrypt(Cipher *c, FILE *infile, FILE *outfile);
static void EncryptBlock(Cipher *c, unsigned int *d);
static unsigned int RoundFunc(Cipher *c, unsigned int k, unsigned int d);
static unsigned int RoundFunc0(Cipher *c, int i, int *x,
                               unsigned int k, unsigned int d);

/******************************************************************** main */
int main(int argc, char **argv)
{
   FILE *infile, *outfile;
   Cipher *c;
   int e;

   if( argc < 5 )
      return printf("%s <key> <encrypt|decrypt> <infile> <outfile>\n", *argv);

   if( (infile = fopen(*++argv, "rb")) == NULL )
      return printf("Can not open %s\n", *argv);
   c = GenerateCipher(infile);
   (void)fclose(infile);
   if( c == NULL )
      return puts("Out of memory");

   argv++;
   e = (**argv == 'e' || **argv == 'E') ? 1 : 0;

   if( (infile = fopen(*++argv, "rb")) == NULL )
   {
      FreeCipher(c);
      return printf("Can not open %s\n", *argv);
   }
   if( (outfile = fopen(*++argv, "wb+")) == NULL )
   {
      FreeCipher(c);
      (void)fclose(infile);
      return printf("Can not create %s\n", *argv);
   }

   e = (e != 0) ? Encrypt(c, infile, outfile) : Decrypt(c, infile, outfile);
   if( e != 0 )
      (void)puts("Out of memory");

   FreeCipher(c);
   (void)fclose(infile);
   (void)fclose(outfile);
   return 0;
}


/******************************************************************* Bswap */
#ifdef BIG_ENDIAN
static void Bswap(unsigned int *d, unsigned int blockcount)
{
   unsigned int i;
   for(i = 0; i < blockcount; i++)
   {
      *d = ((*d >> 24) & 0x000000ff)
         | ((*d >>  8) & 0x0000ff00)
         | ((*d <<  8) & 0x00ff0000)
         | ((*d << 24) & 0xff000000);
      d++;
   }
}
#endif

/*************************************************************** BlockRead */
static void BlockRead(FILE *infile, unsigned int *d, unsigned int blocksize)
{
   (void)fread(d, (size_t)blocksize, 1, infile);
   #ifdef BIG_ENDIAN
      Bswap(d, blocksize >> 2);
   #endif
}

/************************************************************** BlockWrite */
static void BlockWrite(FILE *outfile, unsigned int *d, unsigned int blocksize)
{
   #ifdef BIG_ENDIAN
      unsigned int t, i;
      for(i = 0; i < (blockcount >> 2); i++)
      {
         t = *(d++);
         Bswap(&t, 1);
         (void)fwrite(&t, 4, 1, outfile);
      }
   #else
      (void)fwrite(d, (size_t)blocksize, 1, outfile);
   #endif
}

/***************************************************************** Encrypt */
static int Decrypt(Cipher *c, FILE *infile, FILE *outfile)
{
   unsigned int *d[2], *p, filesize, bytesize, i, b;

   bytesize = (unsigned int)(c->blocksize * sizeof(unsigned int));
   if( (d[0] = (unsigned int*)malloc((size_t)bytesize)) == NULL )
      return 1;
   if( (d[1] = (unsigned int*)malloc((size_t)bytesize)) == NULL )
   {
      free(d[0]);
      return 1;
   }
   if( (p = (unsigned int*)malloc((size_t)bytesize)) == NULL )
   {
      free(d[0]);
      free(d[1]);
      return 1;
   }

   /* Get file size */
   (void)fseek(infile, 0, SEEK_END);
   filesize = (unsigned int)ftell(infile);
   (void)fseek(infile, 0, SEEK_SET);
   if( filesize < bytesize || (filesize % bytesize) != 0 )
   {
      (void)puts("Decrypt error");
      free(d[0]);
      free(d[1]);
      free(p);
      return 0;
   }

   /* Read initialization vector */
   BlockRead(infile, d[0], bytesize);
   DecryptBlock(c, d[b = 0]);
   memcpy(p, d[0], (size_t)bytesize);

   /* Check file size */
   if( p[c->blocksize - 1] <= filesize - bytesize * 2 ||
       p[c->blocksize - 1] > filesize - bytesize )
   {
      (void)puts("Decrypt error");
      free(d[0]);
      free(d[1]);
      free(p);
      return 0;
   }
   filesize = p[c->blocksize - 1];

   /* Read padding */
   if( (filesize % bytesize) != 0 )
   {
      BlockRead(infile, d[1], bytesize);
      memcpy(p, d[b = 1], (size_t)bytesize);
      DecryptBlock(c, p);
      for(i = 0; (int)i < c->blocksize; i++)
         p[i] ^= d[0][i];
      i = filesize % bytesize;
      #ifdef BIG_ENDIAN
         Bswap(p, c->blocksize);
      #endif
      (void)fwrite((unsigned char*)p + bytesize - i, (size_t)i, 1, outfile);
      filesize -= i;
   }

   /* Decrypt contents */
   assert((filesize % bytesize) == 0);
   for(; filesize > 0; filesize -= bytesize)
   {
      BlockRead(infile, d[b ^= 1], bytesize);
      memcpy(p, d[b], (size_t)bytesize);
      DecryptBlock(c, p);
      for(i = 0; (int)i < c->blocksize; i++)
         p[i] ^= d[b ^ 1][i];
      (void)fwrite(p, (size_t)bytesize, 1, outfile);
   }

   free(d[0]);
   free(d[1]);
   free(p);
   return 0;
}

/************************************************************ DecryptBlock */
static void DecryptBlock(Cipher *c, unsigned int *d)
{
   unsigned int *k, x;
   Round *r;
   int i, j;

   k = c->k_table + c->rounds * c->iter - 1;
   for(i = 0; i < c->iter; i++)
   {
      r = c->r + c->rounds - 1;
      for(j = 0; j < c->rounds; j++)
      {
         if( r->xor != 0 )
            x = d[r->in] ^ RoundFunc(c, *k--, d[r->out]);
         else
            x = d[r->in] - RoundFunc(c, *k--, d[r->out]);
         d[r->in] = d[r->out];
         d[r--->out] = x;
      }
   }
}

/***************************************************************** Encrypt */
static int Encrypt(Cipher *c, FILE *infile, FILE *outfile)
{
   unsigned int *d[2], filesize, bytesize, i, b;

   bytesize = (unsigned int)(c->blocksize * sizeof(unsigned int));
   if( (d[0] = (unsigned int*)malloc((size_t)bytesize)) == NULL )
      return 1;
   if( (d[1] = (unsigned int*)malloc((size_t)bytesize)) == NULL )
   {
      free(d[0]);
      return 1;
   }

   /* Generate initialization vector */
   srand((unsigned)time(NULL) ^ *(c->k_table));
   for(i = 0; (int)i < c->blocksize - 1; d[0][i++] = (unsigned int)rand());
   (void)fseek(infile, 0, SEEK_END);
   d[0][i] = filesize = (unsigned int)ftell(infile);
   (void)fseek(infile, 0, SEEK_SET);

   /* Write initialization vector */
   memcpy(d[1], d[0], (size_t)bytesize);
   EncryptBlock(c, d[1]);
   BlockWrite(outfile, d[1], bytesize);
   b = 0;

   /* Write padding block */
   if( (filesize % bytesize) != 0 )
   {
      for(i = 0; (int)i < c->blocksize; d[1][i++] = (unsigned int)rand());
      i = filesize % bytesize;
      (void)fread((unsigned char*)d[1] + bytesize - i, (size_t)i, 1, infile);
      #ifdef BIG_ENDIAN
         Bswap(d[1], c->blocksize);
      #endif
      for(i = 0; (int)i < c->blocksize; i++)
         d[1][i] ^= d[0][i];
      EncryptBlock(c, d[b = 1]);
      BlockWrite(outfile, d[1], bytesize);
      filesize -= filesize % bytesize;
   }

   /* Encrypt contents */
   assert((filesize % bytesize) == 0);
   for(; filesize > 0; filesize -= bytesize)
   {
      (void)fread(d[b ^= 1], (size_t)bytesize, 1, infile);
      for(i = 0; (int)i < c->blocksize; i++)
         d[b][i] ^= d[b ^ 1][i];
      EncryptBlock(c, d[b]);
      BlockWrite(outfile, d[b], bytesize);
   }

   free(d[0]);
   free(d[1]);
   return 0;
}

/************************************************************ EncryptBlock */
static void EncryptBlock(Cipher *c, unsigned int *d)
{
   unsigned int *k, x;
   Round *r;
   int i, j;

   k = c->k_table;
   for(i = 0; i < c->iter; i++)
   {
      r = c->r;
      for(j = 0; j < c->rounds; j++)
      {
         if( r->xor != 0 )
            x = d[r->out] ^ RoundFunc(c, *k++, d[r->in]);
         else
            x = d[r->out] + RoundFunc(c, *k++, d[r->in]);
         d[r->out] = d[r->in];
         d[r++->in] = x;
      }
   }
}

/*************************************************************** RoundFunc */
static unsigned int RoundFunc(Cipher *c, unsigned int k, unsigned int d)
{
   int x;
   x = 0;
   return k ^ RoundFunc0(c, 0, &x, k, d);
}
static unsigned int RoundFunc0(
   Cipher *c, int i, int *x, unsigned int k, unsigned int d)
{
   #define b5(_)        (unsigned int)((_) & 31)
   #define ror(a, b)    (((a) << b5(b)) | ((a) >> (unsigned)(31 - b5(b))))
   unsigned int r, a;

   if( i >= c->funcsize - 1 )
      return (c->f[(*x)++].term != 0) ? d : k;

   /* Here the operations must be on separate lines
      to prevent optimizer from reversing evaluation order */
   switch( c->f[i].op )
   {
      case OP_XOR:
         r = RoundFunc0(c, i * 2 + 1, x, k, d);
         r ^= RoundFunc0(c, i * 2 + 2, x, k, d);
         break;
      case OP_MUL:
         r = RoundFunc0(c, i * 2 + 1, x, k, d);
         r *= RoundFunc0(c, i * 2 + 2, x, k, d);
         r++;
         break;
      case OP_ADD:
         r = RoundFunc0(c, i * 2 + 1, x, k, d);
         r += RoundFunc0(c, i * 2 + 2, x, k, d);
         break;
      case OP_SUB:
         r = RoundFunc0(c, i * 2 + 1, x, k, d);
         r -= RoundFunc0(c, i * 2 + 2, x, k, d);
         break;
      case OP_TABLE:
         r = RoundFunc0(c, i * 2 + 1, x, k, d);
         r ^= c->s_table[b5(RoundFunc0(c, i * 2 + 2, x, k, d))];
         break;
      case OP_ROR:
         r = RoundFunc0(c, i * 2 + 1, x, k, d);
         a = RoundFunc0(c, i * 2 + 2, x, k, d);
         r = ror(r, a);
         break;
      default:
         assert(0);
         r = ~0U;
         break;
   }
   return r;
}