/* simplecrypt.c - a simple file encryption/decryption utlity
 * see left404.com/2012/07/09/simplecrypt for more information.
 * 
 * Copyright (C) 7/9/2012  Dara Hazeghi
 * This software may be freely modified and redistributed.  It has no warranty.
 *
 * simplecrypt uses the xxtea block cipher for encryption
 * -supports arbitrary-length inputs (automatically padded)
 * -cipher block chaining to encrypt multiple blocks securely
 */

#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <time.h>

/* only available on POSIX */
#include <termios.h>
#include <unistd.h>

/* fixme: better way of removing in-use files upon abnormal exit? */
FILE *in, *out = NULL;
const char *infile, *outfile;

#define BLOCK_LEN   8 /* block length in words - must be > 1 */
#define BLOCK_BYTES BLOCK_LEN * sizeof(int) /* block length in bytes */

#define DELTA 0x9e3779b9
#define MX (((z>>5^y<<2) + (y>>3^z<<4)) ^ ((sum^y) + (key[(p&3)^e] ^ z)))

/* See http://en.wikipedia.org/wiki/XXTEA for details on XXTEA */
void xxtea_encrypt(unsigned int *v, const unsigned int key[4])
{
  unsigned int y, z, sum;
  unsigned int p, rounds, e;
  
  rounds = 6 + 52/BLOCK_LEN;
  sum = 0;
  z = v[BLOCK_LEN-1];
  do {
    sum += DELTA;
    e = (sum >> 2) & 3;
    for (p=0; p<BLOCK_LEN-1; p++) {
      y = v[p+1]; 
      z = v[p] += MX;
    }
    y = v[0];
    z = v[BLOCK_LEN-1] += MX;
  } while (--rounds);  
}

void xxtea_decrypt(unsigned int *v, const unsigned int key[4])
{
  unsigned int y, z, sum;
  unsigned int p, rounds, e;
  
  rounds = 6 + 52/BLOCK_LEN;
  sum = rounds*DELTA;
  y = v[0];
  do {
    e = (sum >> 2) & 3;
    for (p=BLOCK_LEN-1; p>0; p--) {
      z = v[p-1];
      y = v[p] -= MX;
    }
    z = v[BLOCK_LEN-1];
    y = v[0] -= MX;
  } while ((sum -= DELTA) != 0);
}

void error(const char *msg)
{
  printf("error: %s\n", msg);
  if(out)
  {
    fclose(out);
    unlink(outfile);
  }
  exit(1);
}

/* pad buf with n random bytes */
void pad_bytes(unsigned char *buf, int n)
{
  int i;
  srand(time(NULL));
  for(i = 0; i < n; i++)
    buf[i] = rand();
}

/* xor words of dest and src together, store in dest */
void combine_blocks(unsigned int *dest, const unsigned int *src)
{
  int i;
  for(i = 0; i < BLOCK_LEN; i++)
    dest[i] = dest[i] ^ src[i];
}

/* swap two block pointers */
void swap_blocks(unsigned int **a, unsigned int **b)
{
  unsigned int *tmp = *a;
  *a = *b;
  *b = tmp;
}

/* initialize buf with the iv vector */
void init_random_block(unsigned int *buf)
{
  int i;
  srand(time(NULL));
  for(i = 0; i < BLOCK_LEN; i++) buf[i] = rand();
}

/* display a block - fixme */
void print_block(const unsigned int *block)
{
  int i;
  for(i = 0; i < BLOCK_LEN; i++) printf("%X ", block[i]);
  printf("\n");
}

void write_encrypted_block(unsigned int *block, unsigned int *chain, const unsigned int *key, FILE *out)
{
  combine_blocks(block, chain); /* chain blocks */
  xxtea_encrypt(block, key); /* encrypt */
  if(fwrite(block, 1, BLOCK_BYTES, out) != BLOCK_BYTES) error("couldn't write to output"); /* write */
  swap_blocks(&block, &chain); /* cipher text becomes new chained block */
}

/* encrypt contents of file in using key, and send to out */
void encrypt_file(FILE *in, FILE *out, const unsigned int *key)
{
  unsigned int len, total;
  unsigned int b[BLOCK_LEN], c[BLOCK_LEN], *buf = b, *chain = c;
  
  fseek(in, 0, SEEK_END); /* find file size */
  total = ftell(in);
  rewind(in);
  
  init_random_block(chain); /* create and write initialization vector */
  if(fwrite(chain, 1, BLOCK_BYTES, out) != BLOCK_BYTES) error("couldn't write to output");
  
  
  buf[0] = total; /* length stored in first word of first block */
  if((len = fread(buf+1, 1, BLOCK_BYTES - sizeof(int), in) + sizeof(int)) != BLOCK_BYTES)
    goto end_block;
    
  write_encrypted_block(buf, chain, key, out);
  
  while((len = fread(buf, 1, BLOCK_BYTES, in)) == BLOCK_BYTES)
  {
    write_encrypted_block(buf, chain, key, out);
  }

  if(len != 0) /* if file is not a multiple of BLOCK_LEN bytes */
  {
end_block:
    pad_bytes((unsigned char *)buf + len, BLOCK_BYTES - len);
    write_encrypted_block(buf, chain, key, out);
  }
}

void read_encrypted_block(unsigned int *block, unsigned int *chain, unsigned int *old, const unsigned int *key, FILE *in)
{
  if(fread(block, 1, BLOCK_BYTES, in) != BLOCK_BYTES) error("wrong password or malformed input");
  memcpy(old, block, BLOCK_BYTES); /* save cipher text */
  xxtea_decrypt(block, key); /* decrypt */
  combine_blocks(block, chain); /* unchain */
  swap_blocks(&old, &chain); /* old ciphertext -> new chain block */
}

/* decrypt contents of file in using key, and send to out */
void decrypt_file(FILE *in, FILE *out, const unsigned int *key)
{
  unsigned int total;
  unsigned int b[BLOCK_LEN], c[BLOCK_LEN], o[BLOCK_LEN], *buf = b, *chain = c, *old = o;
  
  if(fread(chain, 1, BLOCK_BYTES, in) != BLOCK_BYTES) error("invalid input file");
  
  read_encrypted_block(buf, chain, old, key, in);
  total = buf[0];
  
  if(total < BLOCK_BYTES - sizeof(int))
  {
    if(fwrite(buf + 1, 1, total, out) != total) error("couldn't write to output");
    return;
  }
  
  if(fwrite(buf + 1, 1, BLOCK_BYTES - sizeof(int), out) != BLOCK_BYTES - sizeof(int)) error("couldn't write to output");
  total -= (BLOCK_BYTES - sizeof(int));
  
  while(total >= BLOCK_BYTES) /* continue until only a non-full block remains */
  {
    read_encrypted_block(buf, chain, old, key, in);
    if(fwrite(buf, 1, BLOCK_BYTES, out) != BLOCK_BYTES) error("couldn't write to output");
    total -= BLOCK_BYTES;
  }
  
  if(total != 0) /* file wasn't multiple of BLOCK_LEN bytes */
  {
    read_encrypted_block(buf, chain, old, key, in);
    if(fwrite(buf, 1, total, out) != total) error("couldn't write to output");
  }
}

int main(int argc, const char *argv[])
{
  int enc = 0, len;
  unsigned int key[5] = { 0 }; /* key must be zeroed out */
  struct termios passterm, normterm;
  char *k = (char *)&key;

  if(argc != 4)
  {
    printf("usage: %s -e/-d infile outfile\n", argv[0]);
    exit(1);
  }
  
  if(strcmp(argv[1], "-e") == 0)
    enc = 1;
  
  else if(strcmp(argv[1], "-d") == 0)
    enc = 0;
    
  else
  {
    printf("invalid option %s\n", argv[1]);
    exit(1);
  }
  
  infile = argv[2];
  outfile = argv[3];
  
  tcgetattr(STDOUT_FILENO, &normterm);
	passterm = normterm;
	passterm.c_lflag &= ~(ECHO | ICANON);
  tcsetattr(STDOUT_FILENO, TCSANOW, &passterm); /* hide typed password */
  
  printf("enter the key (max: 16 characters): ");
  fgets(k, 18, stdin);
  
  tcsetattr(STDOUT_FILENO, TCSANOW, &normterm);
  putchar('\n');
  
  len = strlen(k);
  if(len == 17 && k[len-1] != '\n') error("key too long");
  
  if(!(in = fopen(infile, "rb"))) error("couldn't open input file");
  if((out = fopen(outfile, "rb"))) error("output file already exists");
  if(!(out = fopen(outfile, "wb"))) error("couldn't open output file");
  
  if(enc) encrypt_file(in, out, key);
  else decrypt_file(in, out, key);
  
  fclose(out);
  fclose(in);
  
  return 0;
}