// 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 1 when compiling on big-endian machines
const Word swap_byte_order = 1,

// Number of encryption rounds
           rounds = 64,

// Preallocated buffers.  16k blocks (as opposed to anything larger)
// appears to offer the best performance on mingw.
           buffer_size = 0x4000;
Byte buffer[buffer_size + 16],
     sbox[16][256], inverse_sbox[16][256];
Word iv_buffer[buffer_size / 4], *iv,
     key[rounds * 4], inverse_key[rounds * 4], salt[4],
     ibaa_memory[256], ibaa_accumulator, ibaa_index, ibaa_value,
     i, j, k, x, y, z;

int size, s2, s3;

tm *lt, t;

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

   const char *text[4] =
   {
   // Sbox parameters.  See sbox_params.c
      "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>",
      "[yyyy-mm-dd] <key> [input] [output]",
      "read error",
      "write error"
   };

// Rotate right
Word rotr(Word p, int q)
{
   return (p >> q) | (p << (32 - q));
}

// Rotate left
Word rotl(Word p, int q)
{
   return (p << q) | (p >> (32 - q));
}

// Convert byte order for a block of data
void ConvertByteOrder(Word block_count, Word *data)
{
   if( swap_byte_order )
      for(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.
void NextRand()
{
   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
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
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
void InitSbox(int group)
{

   for(x = 0; x < 16; x++)
   {
      size = text[0][(group * 16 + x) * 3] ^ 81;
      s2 = text[0][(group * 16 + x) * 3 + 1] ^ 25;
      s3 = text[0][(group * 16 + x) * 3 + 2] ^ 16;
      for(i = 0; i < 256; i++)
      {
         j = ((i + size) & 0x7f) | (i & 0x80);
         j = (((j >> 1) + s2) & 0x7f) | ((j << 7) & 0x80);
         sbox[x][i] = (((j >> 1) + s3) & 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(y = 0; y < 2; y++)
         for(z = 0; z < 8; z++)
            for(i = 0; i < 256; i++)
            {
               k = sbox[x][i] ^ sbox[x][i ^ (1 << z)];

               // 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( !(k & (k - 1)) )
               {
                  j = i ^ (1 << ((z + 1) % 8));
                  s3 = sbox[x][i];
                  sbox[x][i] = sbox[x][j];
                  sbox[x][j] = s3;
               }
            }
   }
}

// Decrypt a single block
void Decrypt(Word *data)
{
   for(x = 0; x < rounds; x++)
   {
      Unmix(data[0], data[1], data[2], data[3]);
      for(y = 0; y < 4; y++)
      {
         k = data[y];
         for(z = 0; z < 4; z++)
         {
            k = rotl(k, 8);
            k = (k & ~0xff) | inverse_sbox[x & 15][k & 0xff];
         }
         data[y] = k ^ inverse_key[x * 4 + y];
      }
   }
   for(x = 0; x < 4; x++)
      data[x] ^= iv[x];
}

int Read(int size, void *output)
{
   return fread(output, 1, size, infile);
}

void Write(int size)
{
   fwrite(buffer, size, 1, outfile);
}

void Print(int msg)
{
   puts(text[msg]);
}

void Close()
{
   fclose(infile);
}

// Program entry point
int main(int argc, char **argv)
{
   if( argc < 2 )
      Print(1);
   else
   {
      int arg_offset;
      memset(&t, 0, sizeof(t));
      if( sscanf(argv[1], "%d-%d-%d", &t.tm_year, &t.tm_mon, &t.tm_mday) - 3 )
      {
         // Set decryption mode
         arg_offset = 1;

         // Set decryption timestamp from current time
         time_t 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".
         lt = localtime(&s);
         t.tm_year = lt->tm_year;
         t.tm_mon = lt->tm_mon;
         t.tm_mday = lt->tm_mday;
      }
      else
      {
         // 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>.
      }

      // 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")) )
      {
         Print(2);
         return 1;
      }
      memset(ibaa_memory, 10, 1024);
      Read(1024, ibaa_memory);
      Close();
      ConvertByteOrder(256, ibaa_memory);

      ibaa_accumulator = 1;

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

      // Generate the remaining keys by passing them through the
      // encryption function.
      InitSbox(8);
      for(i = 1; i < rounds; i++)
      {
         for(j = 0; j < 4; j++)
         {
            x = key[(i - 1) * 4 + j];
            for(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]) )
      {
         if( !(infile = fopen(argv[arg_offset], "rb")) )
         {
            Print(2);
            return 1;
         }
      }
      else
      {
         infile = stdin;
      }
      arg_offset++;
      if( argc > arg_offset &&
          (argv[arg_offset][0] - '-' && argv[arg_offset][1]) )
      {
         if( !(outfile = fopen(argv[arg_offset], "wb+")) )
         {
            Close();
            Print(3);
            return 1;
         }
      }
      else
      {
         outfile = stdout;
      }

      // Initialize sbox to current moon phase
      z = 2551443;
      InitSbox((((mktime(&t) - 434928) % z) * 8) / z);

      if( arg_offset - 4 )
      {
         // Decrypt
         for(i = 0; i < rounds; i++)
            for(j = 0; j < 4; j++)
               inverse_key[(rounds - 1 - i) * 4 + j] = key[i * 4 + j];

         for(j = 0; j < 16; j++)
            for(i = 0; i < 256; i++)
               inverse_sbox[15 - j][sbox[j][i]] = (Byte)i;

         // First block is random salt.  No decryption needed, but byte
         // order may need to be adjusted.
         if( Read(16, salt) == 16 )
         {
            ConvertByteOrder(4, salt);
            iv = salt;

            // Decrypt blocks
            size = 16;
            for(j = 0;
                size == 16 && (size = Read(buffer_size - j, buffer + j)) > 0; j = 16)
            {
               size += j;

               // Decrypt everything before the last full block
               s2 = (size & ~15) - 16;
               if( s2 > 0 )
               {

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

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

                  // Write decrypted bytes
                  Write(s2);

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

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

               // Decrypt the last block
               ConvertByteOrder(4, (Word*)buffer);
               Decrypt((Word*)buffer);
               ConvertByteOrder(4, (Word*)buffer);
               Write(j);
            }
         }
      }
      else
      {
         // 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.
         timeval tv;
         gettimeofday(&tv, NULL);
         memcpy(ibaa_memory, &tv, sizeof(tv));
         ibaa_memory[99] = getpid();
         for(i = 0; i < 8192; i++)
            NextRand();
         for(i = 0; i < 4; salt[i++] = ibaa_value)
            NextRand();

         iv = (Word*)memcpy(iv_buffer, salt, 16);
         ConvertByteOrder(4, salt);
         memcpy(buffer, salt, 16);
         Write(16);

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

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

            // Encrypt contents
            ConvertByteOrder(s2 * 4, (Word*)buffer);
            for(i = 0; i < (Word)size; i += 16)
            {
               Word *data = (Word*)(buffer + i);
               for(j = 0; j < 4; j++)
                  data[j] ^= iv[j];
               for(j = 0; j < rounds; j++)
               {
                  for(k = 0; k < 4; k++)
                  {
                     x = data[k] ^ key[j * 4 + k];
                     for(z = 0; z < 4; z++)
                     {
                        x = (x & ~0xff) | sbox[j & 15][x & 0xff];
                        x = rotr(x, 8);
                     }
                     data[k] = x;
                  }
                  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(s2 * 4, (Word*)buffer);
            Write(s2 * 16);
         }

         // 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.
         buffer[0] = size & 15;
         Write(1);
      }

      Close();
      infile = outfile;
      Close();
   }
   return 0;
}