// 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 = false;

// 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));
}

// Swap byte order
static Word bswap(Word x)
{
   return (rotl(x, 8) & 0x00ff00ff) | (rotr(x, 8) & 0xff00ff00);
}

// 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] = bswap(data[i]);
   }
}

// 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;
               }
            }
         }
      }
   }
}

// Initialize inverse_sbox from sbox
static void InitInverseSbox()
{
   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;
   }
}

// Initialize key from file
static bool InitKey(const char *key_file)
{
   // 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(key_file, "rb")) == NULL )
      return false;
   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]);
   }
   return true;
}

// Initialize decryption keys.  These are the same as encryption keys,
// but applied in reverse order.
static void InitInverseKey()
{
   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];
   }
}

// Open input and output file handles
static bool InitFiles(int arg_offset, int argc, char **argv)
{
   if( argc <= arg_offset ||
       (argv[arg_offset][0] == '-' && argv[arg_offset][1] == '\0') )
   {
      infile = stdin;
   }
   else
   {
      if( (infile = fopen(argv[arg_offset], "rb")) == NULL )
      {
         printf("Can not open %s for reading\n", argv[arg_offset]);
         return false;
      }
   }
   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);
         printf("Can not open %s for writing\n", argv[arg_offset + 1]);
         return false;
      }
   }
   return true;
}

// Encrypt a single block
static void Encrypt(const Word *iv, Word *data)
{
   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]);
   }
}

// 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];
}

// Encrypt a file with cipher block chaining
static void EncryptFile()
{
   // 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)
      {
         Encrypt(iv, (Word*)(buffer + i));
         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);
}

// Decrypt a file
static void DecryptFile()
{
   // 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 )
      return;  // Give up early on any read errors
   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 )
      return;
   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);
}

// Return moon phase for a particular timestamp, 0 = new, 4 = full, etc.
static int GetMoonPhase(int t)
{
   const int synodic_period = (int)(29.530589 * 86400);
   t -= 434928;
   return ((t % synodic_period) * 8) / synodic_period;
}

// 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;
   }

   if( !InitKey(argv[arg_offset]) )
      return printf("Can not initialize key from %s\n", argv[arg_offset]);
   InitFiles(arg_offset + 1, argc, argv);

   // Initialize sbox to current moon phase
   s = mktime(&t);
   InitSbox(GetMoonPhase(s));

   if( arg_offset == 2 )
   {
      // Encrypt
      EncryptFile();
   }
   else
   {
      // Decrypt
      InitInverseKey();
      InitInverseSbox();
      DecryptFile();
   }

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