/*
 * The Cryptonit security software suite is developped by IDEALX
 * Cryptonit Team (http://IDEALX.org/ and http://cryptonit.org).
 *
 * Copyright 2003-2006 IDEALX
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License version
 * 2 as published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
 * 02110-1301, USA.
 */

#include <unistd.h>
#include <stdio.h>

#include "types.h"

#include "aes.h"
#include "aes.c"

#define CIPHER_BLOCK_SIZE AES_BLOCK_SIZE

#include "sha2.h"
#include "sha2.c"

#ifdef DEBUG
void dump(char *name, unsigned char *val, int size)
{
  int i;
  fprintf(stderr, "%s = ", name);
  for(i=0;i<size;i++) {
    if((i%32 == 0) && (i)) {
      fprintf(stderr, "\n");
    }
    fprintf(stderr, "%02x", val[i]);
  }
  fprintf(stderr, "\n");
}
#define DUMP(x,l) \
  dump(#x, x, l); 
#else
#define DUMP(x,l)
#endif


#ifdef DEBUG
static int get_size(FILE *f)
{
  int size;
  fseek(f, 0, SEEK_END);
  size = ftell(f);
  rewind(f);
  return size;
}
#endif


#define BLOCK_SIZE 1024

enum {
  WRONG_ARG_COUNT        = 1,
  ERROR_OPENING_SELF     = 2,
  ERROR_READING_SIZE     = 3,
  ERROR_READING_SALT     = 4,
  ERROR_CREATING_OUTPUT  = 5,
  ERROR_GETTING_PASSWORD = 6,
} error;

#include "pkcs5.c"

int main(int argc, char **argv)
{
  unsigned int size = 0, index = 0, next = 0, file_size = 0;
  int rv = 0;
  hmac_ctx hmac;
  aes_ctx aes;

  char *password = NULL;
  uint64 salt = 0;
  uint8 buffer[BLOCK_SIZE];
  uint8 key[SHA256_DIGEST_LENGTH];
  uint8 hash[SHA256_DIGEST_LENGTH];
  FILE *fin = NULL, *fout = NULL;

  /* Do we have the right amount of arguments ? */
  if(argc != 2) {
    printf("usage:\n\t%s oufile\n\n", argv[0]);
    rv = WRONG_ARG_COUNT;
    goto end;
  }

  /* Open ourselves for reading */
  if((fin = fopen(argv[0], "r+b")) == NULL) {
    rv = ERROR_OPENING_SELF;
    goto end;
  }

  /* Read original file size stored in last four bytes */
#ifdef DEBUG
  file_size = get_size(fin);
  printf("Reading size at %d\n", file_size - (sizeof(unsigned int) + sizeof(uint64)));
#endif
  fseek(fin, -(sizeof(unsigned int) + sizeof(uint64)), SEEK_END);
  if(fread(&size, sizeof(unsigned int), 1, fin) != 1) {
    rv = ERROR_READING_SIZE;
    goto end;
  }
  

  /* Read original salt stored in last four bytes */
#ifdef DEBUG
  printf("Reading salt at %d\n", file_size - sizeof(uint64));
#endif
  //fseek(fin, -sizeof(uint64), SEEK_END);

  if(fread(&salt, sizeof(uint64), 1, fin) != 1) {
    rv = ERROR_READING_SALT;
    goto end;
  }

#ifdef DEBUG
  printf("Crypted file size = %d\n", size);
  printf("Original salt = %u\n", salt);
#endif

  /* Derive key from password */
  if((password = getpass("Input password: ")) == NULL) {
    rv = ERROR_GETTING_PASSWORD; 
    goto end; 
  } 

  /* Derive password into key */
  derive_key(password, strlen(password), 
	     (char *)&salt, 8, key);

  DUMP(key, SHA256_DIGEST_LENGTH);

  /* Rewinding to start of encrypted file : 
   * = 4 bytes : size of original file
   * + SHA256_DIGEST_LENGTH bytes : hash
   * + size bytes : file size
   * + salt size
   * + padding to CIPHER_BLOCK_SIZE
   */

  index = 4 + SHA256_DIGEST_LENGTH + size + sizeof(uint64);
  if((size % CIPHER_BLOCK_SIZE)) {
    index += CIPHER_BLOCK_SIZE - (size % CIPHER_BLOCK_SIZE);
  }
  fseek(fin, -index, SEEK_END);

  /* Initializing HMAC using the drived key */
  hmac_init( key, SHA256_DIGEST_LENGTH, &hmac );

  for(index = 0; index < size; index += next) {
    int extra = 0;
    next =  size - index;
    if(next > BLOCK_SIZE) {
      next = BLOCK_SIZE;
    } else {
      /* Was there any padding ? */
      if(next % CIPHER_BLOCK_SIZE) {
	extra = CIPHER_BLOCK_SIZE - (next % CIPHER_BLOCK_SIZE);
      }
    }
    /* Read encrypted block */
    if(fread(buffer, (next + extra), 1, fin) != 1) {
      rv = 5;
      goto end;
    }
    /* Update the hmac with the encrypted block */
    hmac_update(buffer, (next + extra), &hmac);
  }

  /* Finish HMAC */
  hmac_final(hash, &hmac);

#ifdef DEBUG
  printf("Reading HMAC at %u.\n", index);
#endif

  /* Read original HMAC */
  if(fread(buffer, SHA256_DIGEST_LENGTH, 1, fin) != 1) {
    rv = 7;
    goto end;
  }

#ifdef DEBUG
  dump("HMAC", buffer, SHA256_DIGEST_LENGTH);
#endif

  /* Compare HMAC */
  if((memcmp(buffer, hash, SHA256_DIGEST_LENGTH)) != 0) {
    rv = 8;
    printf("HMAC is invalid!\n");
    goto end;
  } else {
      rv = 0;
      printf("HMAC is valid.\n");
  }

  /* HMAC test successfull, file can be decyphered */

  /* Opening output file for writing */
  if((fout = fopen(argv[1], "w+b")) == NULL) {
    rv = ERROR_CREATING_OUTPUT;
    goto end;
  }
  
  aes_set_key(&aes, key, 256, 0);

  /* Rewinding to start of encrypted file */
  index = 4 + SHA256_DIGEST_LENGTH + size + sizeof(uint64);
  if((size % CIPHER_BLOCK_SIZE)) {
    index += CIPHER_BLOCK_SIZE - (size % CIPHER_BLOCK_SIZE);
  }
  fseek(fin, -index, SEEK_END);
  
  for(index = 0; index < size; index += next) {
    int extra = 0;
    next =  size - index;
    if(next > BLOCK_SIZE) {
      next = BLOCK_SIZE;
    } else {
      /* Was there any padding ? */
      if(next % CIPHER_BLOCK_SIZE) {
	extra = CIPHER_BLOCK_SIZE - (next % CIPHER_BLOCK_SIZE);
      }
    }
    
    /* Read encrypted block */
    if(fread(buffer, (next + extra), 1, fin) != 1) {
      rv = 5;
      goto end;
    }

    /* Decrypt block */
    aes_cbc_decrypt(&aes, key, buffer, (next + extra));
    
    /* Write block */
    if(fwrite(buffer, next, 1, fout) != 1) {
      rv = 6;
      goto end;
    }
  }

 end:
  if(fin)
    fclose(fin);
  if(fout)
    fclose(fout);
  return rv;
}
