/* blockio.c - Don Yang (uguu.org)

   06/02/12
                                                                           */

#include<assert.h>
#include<stdio.h>
#include<stdlib.h>
#include<string.h>

#define BUFFER_SIZE 0x80
#undef SWAP_BYTE_ORDER

typedef unsigned int Word;
typedef unsigned char Byte;

#define rotr(x,n)   (((x) >> ((int)(n))) | ((x) << (32 - (int)(n))))
#define rotl(x,n)   (((x) << ((int)(n))) | ((x) >> (32 - (int)(n))))
#define bswap(x)    ((rotl(x, 8) & 0x00ff00ff) | (rotr(x, 8) & 0xff00ff00))

/* Random number generator */
static Word RandomNumberStub(void)
{
   return 0x80402010;
}

/* Convert byte order */
static void ConvertByteOrder(int block_count, Word *block)
{
   #ifdef SWAP_BYTE_ORDER
      int i;
      for(i = 0; i < block_count; i++)
         block[i] = bswap(block[i]);
   #endif
}

/* Encrypt a single block */
static void EncryptBlockStub(const Word *iv, Word *block)
{
   int i;

   for(i = 0; i < 4; i++)
   {
      block[i] ^= iv[i];
      /* cipher = no-op */
   }
}

/* Decrypt a single block */
static void DecryptBlockStub(const Word *iv, Word *block)
{
   int i;

   for(i = 0; i < 4; i++)
   {
      /* cipher = no-op */
      block[i] ^= iv[i];
   }
}

/* Encode infile to outfile */
static void Encode(FILE *infile, FILE *outfile)
{
   Word iv_buffer[4], *iv;
   Word salt[4];
   Byte *buffer;
   int size, block_count, i;

   buffer = (Byte*)malloc(BUFFER_SIZE + 16);
   assert(buffer != NULL);

   iv = iv_buffer;
   memset(iv_buffer, 0, sizeof(iv_buffer));

   /* First block is random salt, no encryption needed */
   for(i = 0; i < 4; i++)
      salt[i] = RandomNumberStub();
   memcpy(iv_buffer, salt, sizeof(iv_buffer));
   ConvertByteOrder(4, salt);
   if( fwrite(salt, 16, 1, outfile) != 1 )
      perror("Write error");

   /* Encrypt blocks until end of file */
   while( (size = fread(buffer, 1, BUFFER_SIZE, infile)) > 0 )
   {
      /* Append an 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; i++)
         buffer[size + i] = (Byte)(RandomNumberStub() & 0xff);

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

      /* Encrypt contents */
      ConvertByteOrder(block_count * 4, (Word*)buffer);
      for(i = 0; i < size; i += 16)
      {
         EncryptBlockStub(iv, (Word*)(buffer + i));
         iv = (Word*)(buffer + i);
      }

      /* Save a copy of the last block, to be used as initialization vector
         for the next set of blocks.  This is needed because the byte order
         may be swapped below, and in any case fread will overwrite the
         buffer in the next iteration.                                      */
      assert(iv == (Word*)(buffer + (block_count * 16) - 16));
      iv = (Word*)memcpy(iv_buffer, iv, sizeof(iv_buffer));

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

   /* Append size of last block, encoded as a single byte in plaintext.

      This is not encrypted in any way, since the original file size
      is relatively easy to guess (16-31 bytes smaller than encrypted
      file size).  Encrypting an easy-to-guess data will make it
      easier for attacker to test the encryption key, so we don't
      encrypt this byte.                                              */
   i = size & 15;
   fputc(i, outfile);

   free(buffer);
}

/* Decode infile to outfile */
static void Decode(FILE *infile, FILE *outfile)
{
   Word *iv_buffer, *iv;
   Word salt[4];
   Byte *buffer;
   int size, offset, decrypt_size, i;

   iv_buffer = (Word*)malloc(BUFFER_SIZE);
   buffer = (Byte*)malloc(BUFFER_SIZE);
   assert(buffer != NULL);

   memset(iv_buffer, 0, 16);
   iv = iv_buffer;

   /* First block.  No decryption is needed, but byte order needs to
      be adjusted.                                                   */
   if( fread(salt, 16, 1, infile) != 1 )
      return;
   ConvertByteOrder(4, salt);
   iv = salt;

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

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

      ConvertByteOrder(decrypt_size / 4, (Word*)buffer);
      memcpy(iv_buffer, buffer, decrypt_size);
      for(i = 0; i < decrypt_size; i += 16)
      {
         DecryptBlockStub(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, since it will be overwritten by the memcpy above.      */
      iv = (Word*)memcpy(salt, iv, 16);

      /* Write decrypted bytes */
      if( fwrite(buffer, decrypt_size, 1, outfile) != 1 )
         perror("Write error");

      /* Shift the last block to the beginning of the buffer */
      size -= decrypt_size;
      for(i = 0; i < size; i++)
         buffer[i] = buffer[decrypt_size + i];
      if( size != 16 )
         break;
      offset = 16;
   }

   /* Give up on last block if file was truncated */
   if( size < 17 || buffer[16] >= 16 )
      return;
   size = (int)buffer[16];
   if( size == 0 )
      size = 16;

   /* Decrypt last block */
   ConvertByteOrder(4, (Word*)buffer);
   DecryptBlockStub(iv, (Word*)buffer);
   ConvertByteOrder(4, (Word*)buffer);
   if( fwrite(buffer, size, 1, outfile) != 1 )
      perror("Write error");

   free(buffer);
   free(iv_buffer);
}

/* Program entry point */
int main(int argc, char **argv)
{
   FILE *infile = stdin;
   FILE *outfile = stdout;

   if( argc != 4 )
      return printf("%s <encode|decode> <infile> <outfile>\n", *argv);
   if( argv[2][0] != '-' || argv[2][1] != '\0' )
   {
      if( (infile = fopen(argv[2], "rb")) == NULL )
         return printf("Can not open %s for reading\n", argv[2]);
   }
   if( argv[3][0] != '-' || argv[3][1] != '\0' )
   {
      if( (outfile = fopen(argv[3], "wb+")) == NULL )
      {
         (void)fclose(infile);
         return printf("Can not open %s for writing\n", argv[3]);
      }
   }

   if( argv[1][0] == 'd' )
      Decode(infile, outfile);
   else
      Encode(infile, outfile);

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