// hazuki.cc - Don Yang (uguu.org)
//
// 06/09/12

#include<stdio.h>
#include<string.h>
#include<time.h>
#include<unistd.h>
#include<sys/time.h>

typedef unsigned int Word;
typedef unsigned char Byte;

// Set this to true when compiling on big-endian machines
static const bool swap_byte_order = true;

// Number of encryption rounds
static const int rounds = 64;

// Global encryption states
static Word key[rounds * 4], inverse_key[rounds * 4];
static Byte sbox[16][256], inverse_sbox[16][256];

// Preallocated buffers.  16k blocks (as opposed to anything larger)
// appears to offer the best performance on mingw.
static const int buffer_size = 0x4000;
static Byte buffer[buffer_size + 16];
static Word iv_buffer[buffer_size / 4];

// Random number generator state
static Word ibaa_memory[256], ibaa_accumulator, ibaa_index, ibaa_value;

// Global file handles
static FILE *infile, *outfile;

// Rotate right
static Word rotr(Word x, int n)
{
   return (x >> n) | (x << (32 - n));
}

// Rotate left
static Word rotl(Word x, int n)
{
   return (x << n) | (x >> (32 - n));
}

// Convert byte order for a block of data
static void ConvertByteOrder(Word block_count, Word *data)
{
   if( swap_byte_order )
   {
      for(Word i = 0; i < block_count; i++)
         data[i] = (rotl(data[i], 8) & 0x00ff00ff) | (rotr(data[i], 8) & 0xff00ff00);
   }
}

// Update random number state.  This is the IBAA random number generator, as
// described here: http://burtleburtle.net/bob/rand/isaac.html
//
// This is preferred over ISAAC because the source code to implement it is
// smaller.
static void NextRand()
{
   Word x, y;

   x = ibaa_memory[ibaa_index];
   ibaa_accumulator = (ibaa_accumulator << 19) ^ (ibaa_accumulator >> 13);
   ibaa_memory[ibaa_index] = y =
      ibaa_memory[x & 0xff] + ibaa_accumulator + ibaa_value;
   ibaa_value = ibaa_memory[(y >> 8) & 0xff] + x;
   ibaa_index = (ibaa_index + 1) & 0xff;
}

// Mixing function from Serpent
static void Mix(Word &a, Word &b, Word &c, Word &d)
{
   a = rotl(a, 13);
   c = rotl(c, 3);
   d ^= c ^ (a << 3);
   b ^= a ^ c;
   d = rotl(d, 7);
   b = rotl(b, 1);
   a ^= b ^ d;
   c ^= d ^ (b << 7);
   a = rotl(a, 5);
   c = rotl(c, 22);
}

// Reverse mixing function from Serpent
static void Unmix(Word &a, Word &b, Word &c, Word &d)
{
   c = rotr(c, 22);
   a = rotr(a, 5);
   c ^= d ^ (b << 7);
   a ^= b ^ d;
   d = rotr(d, 7);
   b = rotr(b, 1);
   d ^= c ^ (a << 3);
   b ^= a ^ c;
   c = rotr(c, 3);
   a = rotr(a, 13);
}

// Initialize sboxes
static void InitSbox(int group)
{
   // Sbox parameters.  See sbox_params.c
   static const char params[] =
      "h\32:,<$4~$HhVF+\16tz;Z}8^\1Mt|:0t;^OC|1h|||L|Av>jLt{H*J"
      "H*\nb9]zIcx3hp9n`Nm*?$F-\rzWdN{\4nC&^KEf\33;z0h\0~z25K25"
      "\13Tb{R7\2P8\t04KDNBZo|Lt427J27\n42KT\27\37n'Kb9\rb;\14*"
      "I~PzGPZ!PX\"jg\4.W\16XP;PV\r(ZeRU\16~/.(Xft\22\3h\2=$>BV"
      "aD&%9@.eD,e~u3&[6\10z\rr\32k(X6\n{\rzbmX\"e$J\16|t3&U1 o"
      "JR/9n\t\33,7\13PJ{^\21[Z\23[@\36[p^yB\37[D\32\\P$fj\rzx"
      "\\<vK\2NyMvM\1.%gz!.xTBh\0>,\6Dj[if[;D~BjYj&?B`\16\rR1\5"
      "&1Af[j:_4FiDDhDF9DP>\6@\2aP\6\n4Zv@\0b^}7j\7Mh\4NTr<j\5N"
      "hNkjOkH(KH(\13J}r~>hld<J7|R\31\22R\t\2l\14qXt\35ry=l\nr`"
      "?hR\32\21F\27""7ne<b15*IJ2u>";

   for(int s = 0; s < 16; s++)
   {
      const int offset1 = params[(group * 16 + s) * 3] ^ 81;
      const int offset2 = params[(group * 16 + s) * 3 + 1] ^ 25;
      const int offset3 = params[(group * 16 + s) * 3 + 2] ^ 16;
      for(int i = 0; i < 256; i++)
      {
         int j = ((i + offset1) & 0x7f) | (i & 0x80);
         j = (((j >> 1) + offset2) & 0x7f) | ((j << 7) & 0x80);
         sbox[s][i] = (((j >> 1) + offset3) & 0x7f) | ((j << 7) & 0x80);
      }

      // A desirable quality for sboxes is that any single bit flip in
      // input will cause at least 2 bits to be flipped in the output.
      // The above will get us close, but not quite.  By swapping
      // entries with single bit flipped in output, we can get the
      // desirable quality.  See sbox.c for more details.
      for(int iteration = 0; iteration < 2; iteration++)
      {
         for(int bit = 0; bit < 8; bit++)
         {
            for(int i = 0; i < 256; i++)
            {
               const Byte current = sbox[s][i];
               const Byte neighbor = sbox[s][i ^ (1 << bit)];
               const int difference = current ^ neighbor;

               // Check if "difference" contains only one bit set.  This
               // expression would not work if difference is zero, but this
               // will never happen here since sbox outputs for two
               // different input bytes are guaranteed to be different.
               if( (difference & (difference - 1)) == 0 )
               {
                  int j = i ^ (1 << ((bit + 1) % 8));
                  const Byte tmp = sbox[s][i];
                  sbox[s][i] = sbox[s][j];
                  sbox[s][j] = tmp;
               }
            }
         }
      }
   }
}

// Decrypt a single block
static void Decrypt(const Word *iv, Word *data)
{
   for(int r = 0; r < rounds; r++)
   {
      Unmix(data[0], data[1], data[2], data[3]);
      for(int b = 0; b < 4; b++)
      {
         Word old_block = data[b];
         for(int e = 0; e < 4; e++)
         {
            old_block = rotl(old_block, 8);
            old_block = (old_block & ~0xff) |
                        inverse_sbox[r & 15][old_block & 0xff];
         }
         data[b] = old_block;
         data[b] ^= inverse_key[r * 4 + b];
      }
   }
   for(int b = 0; b < 4; b++)
      data[b] ^= iv[b];
}

// Program entry point
int main(int argc, char **argv)
{
   if( argc < 2 )
      return printf("%s [yyyy-mm-dd] <key> <input> <output>\n", *argv);

   int arg_offset;
   struct tm t;
   time_t s;
   memset(&t, 0, sizeof(t));
   if( sscanf(argv[1], "%d-%d-%d", &t.tm_year, &t.tm_mon, &t.tm_mday) == 3 )
   {
      // Set encryption mode.  We require a target decryption date to
      // be specified for encryption as opposed to using current time,
      // since most operating systems will store the last modification
      // date in the output file's metadata, which will make it easier
      // to guess what the decryption date should be.
      arg_offset = 2;

      // Set target decryption date from command line argument
      t.tm_year -= 1900;
      t.tm_mon--;

      // If an invalid date was specified (e.g. 2012-03-32), mktime
      // will fail.  We will continue to use this invalid timestamp
      // anyways as opposed to switching to decrypt operation,
      // otherwise an incorrectly specified date will cause <input> to
      // be overwritten with <key>.
   }
   else
   {
      // Set decryption mode
      arg_offset = 1;

      // Set decryption timestamp from current time
      s = time(NULL);
      
      // Align timestamp to midnight local time.  This is because
      // mktime will also do that alignment, and it would be desirable
      // to be able to decrypt the file we have just encrypted if the
      // target decryption time was set to today.
      //
      // The default installation of MingW/Cygwin appears to have bad
      // environment settings, such that while encryption uses
      // localtime, decryption uses gmtime, so sometimes decryption
      // will not work for the intended date.  The fix is to set TZ
      // environment string to something sensible, such as "PDT+7".
      struct tm *lt = localtime(&s);
      t.tm_year = lt->tm_year;
      t.tm_mon = lt->tm_mon;
      t.tm_mday = lt->tm_mday;
   }

   // Seed IBAA random number generator from key file.  This is done
   // by filling seed memory with newlines, then replace part of it
   // with key file data.  Because the filler bytes are newlines, final
   // trailing newlines in the key file are insignificant.
   if( (infile = fopen(argv[arg_offset], "rb")) == NULL )
      return printf("Can not initialize key from %s\n", argv[arg_offset]);
   memset(ibaa_memory, 10, 1024);
   fread(ibaa_memory, 1, 1024, infile); 
   fclose(infile);
   ConvertByteOrder(256, ibaa_memory);

   ibaa_index = ibaa_value = 0;
   ibaa_accumulator = 1;

   // Initialize first round of keys from seeded random number generator
   for(int i = 0; i < 4; i++)
   {
      for(int j = 0; j < 8192; j++)
         NextRand();
      key[i] = ibaa_value;
   }

   // Generate the remaining keys by passing them through the
   // encryption function.
   InitSbox(8);
   for(int i = 1; i < rounds; i++)
   {
      for(int j = 0; j < 4; j++)
      {
         Word x = key[(i - 1) * 4 + j];
         for(int k = 0; k < 4; k++)
         {
            x = (x & ~0xff) | sbox[i & 15][x & 0xff];
            x = rotr(x, 8);
         }
         key[i * 4 + j] = x;
      }
      Mix(key[i * 4], key[i * 4 + 1], key[i * 4 + 2], key[i * 4 + 3]);
   }

   arg_offset++;
   if( argc <= arg_offset ||
       (argv[arg_offset][0] == '-' && argv[arg_offset][1] == '\0') )
   {
      infile = stdin;
   }
   else
   {
      if( (infile = fopen(argv[arg_offset], "rb")) == NULL )
      {
         return printf("Can not open %s for reading\n", argv[arg_offset]);
      }
   }
   if( argc <= arg_offset + 1 ||
       (argv[arg_offset + 1][0] == '-' && argv[arg_offset + 1][1] == '\0') )
   {
      outfile = stdout;
   }
   else
   {
      if( (outfile = fopen(argv[arg_offset + 1], "wb+")) == NULL )
      {
         fclose(infile);
         return printf("Can not open %s for writing\n", argv[arg_offset + 1]);
      }
   }

   // Initialize sbox to current moon phase
   const int synodic_period = (int)(29.530589 * 86400);
   InitSbox((((mktime(&t) - 434928) % synodic_period) * 8) / synodic_period);

   if( arg_offset == 3 )
   {
      // Encrypt

      // First block is random salt, no encryption needed.  We use the
      // same random number generator used to generate the key bits.  To
      // ensure that the same key will yield different salt, we will also
      // mix in current time and PID.
      struct timeval tv;
      gettimeofday(&tv, NULL);
      memcpy(ibaa_memory, &tv, sizeof(tv));
      ibaa_memory[99] = getpid();
      for(int i = 0; i < 8192; i++)
         NextRand();
      Word salt[4];
      for(int i = 0; i < 4; i++)
      {
         NextRand();
         salt[i] = ibaa_value;
      }

      Word *iv = (Word*)memcpy(iv_buffer, salt, 16);
      ConvertByteOrder(4, salt);
      fwrite(salt, 16, 1, outfile);

      // Encrypt blocks until end of file
      int size;
      while( (size = fread(buffer, 1, buffer_size, infile)) > 0 )
      {
         // Append extra random block to the input data.  This guarantees
         // that the last block of the file will be a full block.
         for(int i = 0; i < 16; i++)
         {
            NextRand();
            buffer[size + i] = ibaa_value & 0xff;
         }

         // Get the number of blocks to encrypt (rounded up)
         const int block_count = (size + 15) / 16;

         // Encrypt contents
         ConvertByteOrder(block_count * 4, (Word*)buffer);
         for(int i = 0; i < size; i += 16)
         {
            Word *data = (Word*)(buffer + i);
            for(int b = 0; b < 4; b++)
               data[b] ^= iv[b];
            for(int r = 0; r < rounds; r++)
            {
               for(int b = 0; b < 4; b++)
               {
                  data[b] ^= key[r * 4 + b];
                  Word new_block = data[b];
                  for(int e = 0; e < 4; e++)
                  {
                     new_block = (new_block & ~0xff) | sbox[r & 15][new_block & 0xff];
                     new_block = rotr(new_block, 8);
                  }
                  data[b] = new_block;
               }
               Mix(data[0], data[1], data[2], data[3]);
            }
            iv = (Word*)(buffer + i);
         }

         // Save a copy of the last block to be used as initialization
         // vector for the next iteration.  This is needed because fread
         // will overwrite the buffer in the next iteration.
         iv = (Word*)memcpy(iv_buffer, iv, 16);

         // Write encrypted blocks
         ConvertByteOrder(block_count * 4, (Word*)buffer);
         fwrite(buffer, block_count * 16, 1, outfile);
         if( (size & 15) != 0 )
            break;
      }

      // Append size of last block, encoded as a single byte in
      // plaintext.  This must remain in plaintext since the original
      // file size is relatively easy to guess (one of 16 possibilities),
      // so encrypting this length would give the attacker a known plaintext
      // for testing keys.
      int i = size & 15;
      fputc(i, outfile);
   }
   else
   {
      // Decrypt
      for(int r = 0; r < rounds; r++)
      {
         for(int b = 0; b < 4; b++)
            inverse_key[(rounds - 1 - r) * 4 + b] = key[r * 4 + b];
      }

      for(int r = 0; r < 16; r++)
      {
         const int ir = 15 - r;
         for(int i = 0; i < 256; i++)
            inverse_sbox[ir][sbox[r][i]] = (Byte)i;
      }

      // First block is random salt.  No decryption needed, but byte
      // order may need to be adjusted.
      Word salt[4];
      if( fread(salt, 16, 1, infile) == 1 )
      {
         ConvertByteOrder(4, salt);
         Word *iv = salt;

         // Decrypt blocks
         Word offset = 0;
         int size;
         while( (size = fread(buffer + offset, 1, buffer_size - offset, infile)) > 0 )
         {
            size += offset;

            // Decrypt everything before the last full block
            int decrypt_size = (size & ~15) - 16;
            if( decrypt_size <= 0 )
               break;

            ConvertByteOrder(decrypt_size / 4, (Word*)buffer);
            memcpy(iv_buffer, buffer, decrypt_size);
            for(int i = 0; i < decrypt_size; i += 16)
            {
               Decrypt(iv, (Word*)(buffer + i));
               iv = (Word*)(((Byte*)iv_buffer) + i);
            }
            ConvertByteOrder(decrypt_size / 4, (Word*)buffer);

            // Preserve the last block for initialization vector in the next round
            iv = (Word*)memcpy(salt, iv, 16);

            // Write decrypted bytes
            fwrite(buffer, decrypt_size, 1, outfile);

            // Shift last block to the beginning of the buffer.  Next fread
            // will append data to after this first block.
            size -= decrypt_size;
            memmove(buffer, buffer + decrypt_size, size);
            if( size != 16 )
               break;
            offset = 16;
         }

         // Give up on decrypting the last block if file was truncated
         if( size == 17 && buffer[16] < 16 )
         {
            size = (int)buffer[16];
            if( size == 0 )
               size = 16;

            // Decrypt the last block
            ConvertByteOrder(4, (Word*)buffer);
            Decrypt(iv, (Word*)buffer);
            ConvertByteOrder(4, (Word*)buffer);
            fwrite(buffer, size, 1, outfile);
         }
      }
   }

   fclose(infile);
   fclose(outfile);
   return 0;
}