/*
 *   Copyright (C) (2004) (Mario Juric) <mjuric@astro.princeton.edu>
 * 
 *    This program is free software; you can redistribute it and/or 
 *    modify it under the terms of the GNU General Public License 
 *    as published by the Free Software Foundation; either 
 *    version 2 of the License, or (at your option) any later 
 *    version.
 * 
 *    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., 675 Mass Ave, Cambridge, MA 02139, USA.
 *
*/

#include "seekcrypt_internal.h"

#include <stdlib.h>
#include <string.h>
#include <errno.h>

/******************************** ENCINDEX methods *********************************/
int enc_bookmark_stride = 1024;

ENCINDEX *encindex_alloc()
{
	ENCINDEX *idx = malloc(sizeof(*idx));
	if(idx == NULL) return NULL;
	
	memset(idx, 0, sizeof(*idx));
	return idx;
}

void encindex_free(ENCINDEX *this)
{
    if(this == NULL) return;

	free(this->indexfile);
	free(this->index);

	free(this);
}

/* Load the saved state of encryption engine */
void encindex_unserialize(ENCINDEX *this, MCRYPT td, int idx)
{
    mcrypt_enc_set_state(td, &this->index[idx * this->blocksize], this->blocksize);
}

/* Save the current state of encryption engine */
int encindex_serialize(ENCINDEX *this, MCRYPT td, int idx)
{
    int size = this->blocksize;
    if(idx * this->blocksize + size > this->len*this->blocksize)
    {
        set_error("Writing outside index array bounds!\n");
        return -1;
    }    

    mcrypt_enc_get_state(td, &this->index[idx * this->blocksize], &size);

    if(size != this->blocksize)
    {
        set_error("Block sizes differ (this->blocksize=%d and mcrypt_enc_get_state returned %d)!", this->blocksize, size);
        return -1;
    }
}

/* magic identifier of the index file */
static char magic[4] = "ix0\0";

int enc_isave(ENCINDEX *idx, const char *indexfile, const char *passphrase)
{
   	ENCFILE *ef = enc_fopen(indexfile, "wb", passphrase);
   	if(ef == NULL) { return -1; }

    enc_fwrite(magic, 1, 4, ef);
	
    enc_fwrite(&idx->blocksize, 1, sizeof(idx->blocksize), ef);
    enc_fwrite(&idx->bookmark_stride, 1, sizeof(idx->bookmark_stride), ef);
    enc_fwrite(&idx->len, 1, sizeof(idx->len), ef);

    enc_fwrite(idx->index, 1, idx->len * idx->blocksize, ef);

    enc_fclose(ef);

    return 0;
}

ENCINDEX *enc_iload(const char *indexfile, const char *passphrase)
{
   	ENCFILE *ef = enc_fopen(indexfile, "rb", passphrase);
   	if(ef == NULL) { return NULL; }

   	char fmagic[4];
    enc_fread(magic, 1, 4, ef);
    if(!memcmp(magic, fmagic, 4))
    {
    	set_error("Index file has incorrect magic value");
    	goto fail;
    }
	
	ENCINDEX *idx = encindex_alloc();

    enc_fread(&idx->blocksize, 1, sizeof(idx->blocksize), ef);
    enc_fread(&idx->bookmark_stride, 1, sizeof(idx->bookmark_stride), ef);
    enc_fread(&idx->len, 1, sizeof(idx->len), ef);

	idx->index = malloc(idx->len * idx->blocksize);
    enc_fread(idx->index, 1, idx->len * idx->blocksize, ef);
    
    enc_fclose(ef);

    return idx;

fail:
    encindex_free(idx);
    enc_fclose(ef);
    return NULL;
}

int min(int a, int b) { return a < b ? a : b; }

ENCINDEX *enc_icreate(ENCFILE *ef, int bookmark_stride)
{
	char *tmpbuf = NULL;
	int nleft, nread;
	ENCINDEX *idx;

	idx = encindex_alloc();
	idx->bookmark_stride = bookmark_stride;

  	nleft = ef->enc_data_length;
	idx->len = nleft / idx->bookmark_stride + (nleft % idx->bookmark_stride != 0); /* this rounds up to higher integer */
	idx->blocksize = ef->blocksize;
	idx->index = malloc(idx->blocksize * idx->len);

	/* initialize hash function module, as a check that the file has been decrypted correctly */
	MHASH tm = mhash_init(ef->checksum_algo);
	if(tm == NULL)
	{
		set_error("mhash initialization failed.\n");
		goto failed;
	}

	/* We don't fseek() because we assume we're at the beginning of the file, and MCRYPT
 	 * has just been initialized with the IV (that is, enc_icreate MUST be called BEFORE
  	 * anything has been read from this file).
	 * That way data[0] entry will be the current contents of MCRYPT
     */
	int atidx = 0;
	encindex_serialize(idx, ef->td, atidx);
	atidx++;

	tmpbuf = malloc(idx->bookmark_stride);
	while(nleft)
	{
		int length = min(idx->bookmark_stride, nleft);

		nread = fread(tmpbuf, 1, length, ef->fp);
		if (nread != length)
    	{
			/* something evil happened, because we've set up the length not to be longer that the file */
	  		set_error("error reading file");
			goto failed;
		}

		mdecrypt_generic(ef->td, tmpbuf, nread);
		mhash(tm, tmpbuf, nread);

		/* store the current decryption engine state */
  		nleft -= length;
  		if(nleft) // last block read (the end of file) has no index entry.
		{
			encindex_serialize(idx, ef->td, atidx);
		}

  		atidx++;
	}

	/* Read and decrypt checksum value */
	void *hash = mhash_end(tm);
	void *fhash = calloc(1, ef->hash_size);
	nread = fread(fhash, 1, ef->hash_size, ef->fp);
	if (nread != ef->hash_size)
	{
		free(fhash);
  		set_error("error reading hash value from file");
		goto failed;
	}
	mdecrypt_generic(ef->td, fhash, nread);

	/* Compare hash value */
	int ckmatch = memcmp(hash, fhash, ef->hash_size) == 0;

	free(fhash);
	free(hash);
	free(tmpbuf);

	if(!ckmatch)
	{
		set_error("Checksums do not match (file corrupted).\n");
		goto failed;
	}

	/* go back to the start of the data, and restore encryption state */
	enc_fseek_index(ef, 0, SEEK_SET, idx);
	return idx;

failed:
    free(tmpbuf);
    encindex_free(idx);
    return NULL;
}

	/* we're done opening the file - now open the index as well */
/*
	ef->idx = enc_iopen(ef, indexfile, index_mode);
	if(ef->idx == NULL)
	{
		set_error("Error opening index!\n");
		goto failed;
	}
*/

int enc_add_index(ENCFILE *ef, const char *indexfile, const char *passphrase, int index_mode)
{
    ENCINDEX *idx;
    int ret = 0;

    /* Try to load the index from file */
    if(index_mode & INDEX_LOAD)
    {
        idx = enc_iload(indexfile, passphrase);
        if(idx != NULL) { ef->idx = idx; return INDEX_LOAD; }
    }

	/* Try to create index */
	if(index_mode & INDEX_CREATE)
	{
	    idx = enc_icreate(ef, enc_bookmark_stride);
	    ret |= INDEX_CREATE;
	}

 	if(idx == NULL) { return 0; }

	/* Try to save if requested */
	if(index_mode & INDEX_SAVE)
	{
	    if(enc_isave(idx, indexfile, passphrase) == 0)
	    {
		    ret |= INDEX_SAVE;
   		}
	}    

    ef->idx = idx;
    return ret;
}


/************************************* ENCFILE ************************************/
int enc_initialized = FALSE;

char *algorithms_directory = NULL;
char *modes_directory = NULL;

char *enc_cipher = DEFAULT_ALGORITHM;
char *enc_mode = DEFAULT_MODE;
char *enc_keymode = DEFAULT_KEYMODE;
int enc_checksum_algo = MHASH_SHA1;

/* if these are set, these are used for IV and salt for the _next_ call of fopen_writable().
 The library takes ownership of both poitners (so they should _not_ be free()d by the user */

char *enc_forced_iv = NULL;
char *enc_forced_salt = NULL;


ENCFILE *encfile_alloc()
{
	ENCFILE *ef = malloc(sizeof(*ef));
	if(ef == NULL) { return NULL; }

	memset(ef, 0, sizeof(*ef));
	return ef;
}

void encfile_free(ENCFILE *this)
{
    if(this == NULL) return;

	encindex_free(this->idx);

 	mcrypt_generic_end(this->td);

	if(this->fp) { fclose(this->fp); }
	
	free(this->salt);
	free(this->key);
	free(this->IV);

	free(this);
}

static int _mcrypt_iv_is_needed( MCRYPT td, char* mode, int noiv)
{
	if (noiv != 0) return 0; /* no IV */

	/* stream mode is treated differently because some stream ciphers use IVs even if the mode itself doesn't. */

	if ( strcmp(mode, "stream")==0 && mcrypt_enc_get_iv_size(td) != 0) { return 1; }
	if ( mcrypt_enc_mode_has_iv(td) != 0 && mcrypt_enc_get_iv_size(td) != 0) { return 1; }

	return 0;
}


ENCFILE *enc_fopen(const char *filename, const char *rwmode, const char *passphrase)
{
	ENCFILE *ef;
  	ef = encfile_alloc();
	if(ef == NULL) { return NULL; }

	if(strchr(rwmode, 'r')) { ef->rwmode |= READABLE; }
	if(strchr(rwmode, 'w')) { ef->rwmode |= WRITABLE; }

	/* open file */
	ef->fp = fopen(filename, rwmode);
	if (ef->fp == NULL)
	{
		set_error("fopen() failed on file %s", filename);
		goto fail;
	}

	/* go to specialized initialization */
	     if(ef->rwmode & WRITABLE) { return enc_fopen_writable(ef, passphrase); }
	else if(ef->rwmode & READABLE) { return enc_fopen_readable(ef, passphrase); }

fail:
	encfile_free(ef);
	return NULL;
}

ENCFILE *enc_fopen_readable(ENCFILE *ef, const char *passphrase)
{
	char cipher[50], mode[50], keymode[50], l_salt[100];
	int noiv, len, j;

	/* load encrypted file headers */
	noiv = 0;
	if(read_header(ef->fp, cipher, mode, keymode, &ef->key_size, l_salt, &ef->salt_size, &noiv, &ef->checksum_algo) != 0)
	{
		set_error("No valid file headers found.");
		goto failed;
	}
	ef->salt = memdup(l_salt, ef->salt_size);

	/* initialize decryption module based on things read from the header */
	ef->td = mcrypt_module_open(cipher, /*algorithms_directory*/ NULL, mode, /*modes_directory*/ NULL);
	if(ef->td == MCRYPT_FAILED)
	{
		set_error("Mcrypt failed to open module.");
		goto failed;
	}

	/* if we did not get the keysize from the header... */
	if(ef->key_size == 0) { ef->key_size = mcrypt_enc_get_key_size(ef->td); }

	/* size of the encryption block (i.e. 256-bit encryption == 32) */
	ef->blocksize = mcrypt_enc_get_block_size(ef->td);

	/* this block of code hashes the passphrase to get the encryption key */
	len = strlen(passphrase);
	ef->key = passphrase_to_key(passphrase, &len, keymode, ef->key_size, ef->salt, ef->salt_size);
	if(ef->key == NULL)
	{
		set_error("There was an error in key generation\n");
		goto failed;
	}

	/* load/get IV */
	ef->IV_size = mcrypt_enc_get_iv_size(ef->td);
	if(_mcrypt_iv_is_needed(ef->td, mode, noiv) != 0)
	{
		ef->IV = read_iv(ef->fp, ef->IV_size);
	}
	else
	{
		ef->IV = calloc(1, ef->IV_size);
		// printf("NO IV needed (noiv=%d).\n", noiv);
	}

	/* decryption initialization, not that we gathered all of the information about
	the encrypted file. When this call returns, td is initialized to correct mcrypt handle.
	*/
	j = mcrypt_generic_init(ef->td, ef->key, len, ef->IV);
	if (j < 0)
	{
		set_error("mcrypt_generic_init failed");
		goto failed;
	}

	/* store the file offset where decrypted data begin (i.e. - here) */
	ef->enc_data_offset = ftell(ef->fp);


	/* get hash size for CRC-like checks*/
	/* allocate a location for file CRC */
	ef->hash_size = mhash_get_block_size(ef->checksum_algo);

	/* seek to the end of the file to figure out the encrypted data size */
	fseek(ef->fp, -ef->hash_size, SEEK_END);
	ef->enc_data_length = ftell(ef->fp) - ef->enc_data_offset;
	fseek(ef->fp, ef->enc_data_offset, SEEK_SET);

	return ef;

failed:
    encfile_free(ef);
    return NULL;
}

ENCFILE *enc_fopen_writable(ENCFILE *ef, const char *passphrase)
{
    int keysize;
    int salt_size;
    byte *IV = NULL;

	if(!enc_initialized)
	{
	    set_error("Encryption module unititialized. Call enc_initialize() first.");
	    return NULL;
	}

	/* Open encryption modules */
	ef->td = mcrypt_module_open(enc_cipher, NULL /*algorithms_directory*/, enc_mode, NULL /*modes_directory*/);
	if (ef->td == MCRYPT_FAILED)
	{
    	set_error("Mcrypt failed to open module.\n");
    	goto fail;
	}

    ef->blocksize = mcrypt_enc_get_block_size(ef->td);

	/* determine keysize */
	keysize = mcrypt_enc_get_key_size(ef->td);

    /* Generate the salt */
    if(!enc_forced_salt)
    {
	    ef->salt = malloc(SALT_SIZE);	/* 20 bytes salt */
	    mcrypt_randomize(ef->salt, SALT_SIZE, FALSE); /* Fill the salt with random data */
	} else {
	    ef->salt = enc_forced_salt;
	    enc_forced_salt = NULL;
	}

	/* Generate IV */
	int ivsize = mcrypt_enc_get_iv_size(ef->td);
	if(enc_forced_iv)
 	{
		IV = enc_forced_iv;
		enc_forced_iv = NULL;
 	} else {
		IV = malloc(ivsize);
		mcrypt_randomize(IV, ivsize, FALSE);
	}

	/* Determine how to hash the key */
	int ki = identify_algorithm(enc_keymode);
	if (ki == -1)
	{
		fprintf(stderr, "Error in keymode '%s'.\n", enc_keymode);
		goto fail;
	}

	/* Determine salt size for keygen algorithm */
	if (mhash_keygen_uses_salt(ki) == 1)
	{
		salt_size = mhash_get_keygen_salt_size(ki);
		if (salt_size == 0) { salt_size = SALT_SIZE; }
	} else {
		salt_size = 0;
	}

	/* Write file header */
	if (write_header(ef->fp, enc_cipher, enc_mode, enc_keymode, &keysize, ef->salt, salt_size) != 0)
	{
		set_error("Error writing file\n");
		goto fail;
	}

	/* Write IV */
	if (_mcrypt_iv_is_needed(ef->td, enc_mode, FALSE) != 0)
 	{
		if (write_iv(ef->fp, IV, ivsize) != 0)
  		{
			set_error("Error writing file\n");
			goto fail;
		}
	}

	fflush(ef->fp);


	/* Hash passphrase to key */
	int len = strlen(passphrase);
	ef->key = passphrase_to_key(passphrase, &len, enc_keymode, keysize, ef->salt, salt_size);
	if (ef->key == NULL)
 	{
		set_error("There was an error in key generation\n");
		goto fail;
	}

	/* initialize encryption and hash engines */
	int tc = mcrypt_generic_init(ef->td, ef->key, len, IV);
	if (tc < 0)
 	{
		mcrypt_perror(tc);
		goto fail;
	}

	ef->tm = mhash_init(enc_checksum_algo);
	ef->hash_size = mhash_get_block_size(enc_checksum_algo);
	if (ef->tm==NULL)
 	{
         set_error("mhash initialization failed.\n");
         goto fail;
	}

	if(IV) { Bzero(IV, ivsize); free(IV); }

	return ef;

fail:
	encfile_free(ef);
	free(IV);
	
	return NULL;
}

void enc_info(ENCFILE *ef)
{
	// fprintf(stderr, "Encrypted data length: %d\n", ef->enc_data_length);
	// fprintf(stderr, "Hash size:%d\n", ef->hash_size);
	// hexdump("IV ", ef->IV, ef->IV_size);
	// printf("Key fixed: %s key_size=%d len=%d\n", passphrase, ef->key_size, len);
	// hexdump("Key", ef->key, ef->key_size);
	// printf("Mcrypt initialized: key_size=%d blocksize=%d\n", ef->key_size, ef->blocksize);
	//printf("Header loaded: %s %s %s \n", cipher, mode, keymode);
	//hexdump("Salt", l_salt, ef->salt_size);
	// fprintf(stderr, "Data offset/length: %d %d\n", ef->enc_data_offset, ef->enc_data_length);


	/*fprintf(stderr, "Hashing passphrase [%s]...\n", passphrase);
	fprintf(stderr, "Salt size for algorithm: %d\n", salt_size);
	hexdump("IV generated", IV, ivsize);
	hexdump("Salt generated", ef->salt, SALT_SIZE);
	fprintf(stderr, "Opened mcrypt module (%s, %s, %s)\n", enc_cipher, enc_mode, enc_keymode);
	fprintf(stderr, "Passphrase len: %d\n", len);
	hexdump("Passphrase hash", ef->key, len);
	*/
}    

void enc_fflush(ENCFILE *ef)
{
    fflush(ef->fp);
}

int enc_fwrite(void *data, int size, int count, ENCFILE *ef)
{
    int nwrote = 0;
    int length = size*count;

	void *ciphertext = malloc(length);
	memcpy(ciphertext, data, length);

	/* calculate hash of plaintext data block */
	mhash(ef->tm, ciphertext, length);

	/* stream mode only, for now*/
	if (mcrypt_enc_is_block_mode(ef->td) == 1)
	{
	    set_error("Block ciphers not supported");
	    goto fail;
	}

	/* Encrypt & write to file */
	mcrypt_generic(ef->td, ciphertext, length);
	nwrote = fwrite(ciphertext, 1, length, ef->fp);
	if (nwrote != length)
	{
		set_error("fwrite() error");
		goto fail;
 	}

fail:
    Bzero(ciphertext, length);
    free(ciphertext);

    return nwrote/size;
}

int enc_fread(void *out, int size, int count, ENCFILE *ef)
{
	/* prepare output */
	int nread = 0;
	int length = size*count;

	/* check if the length requested is longer than the file */
	int at = enc_ftell(ef);
	if(at + length > ef->enc_data_length)
	{
		/* truncate it */
		length = ef->enc_data_length - at;
	}

	/* read data, decrypt */
	nread = fread(out, 1, length, ef->fp);

	/* something evil happened, because we've set up the length not to be longer that the file */
	if (nread != length)
	{
		set_error("fread() error");
		return 0;
	}

	mdecrypt_generic(ef->td, out, nread);
	return nread / size;
}

int enc_fclose(ENCFILE *ef)
{
	int ret;

	if(ef->rwmode & WRITABLE)
 	{
		ret = enc_fclose_writable(ef);
		if(ret) { return ret; }
	}

	encfile_free(ef);
	return 0;
}

int enc_fclose_writable(ENCFILE *ef)
{
	int ret = 0;

	/* Finish hash calculation */
	void *hash = mhash_end(ef->tm);

	/* Encrypt & write hash at the end of the file */
	void *ciphertext = malloc(ef->hash_size);
	memcpy(ciphertext, hash, ef->hash_size);
	mcrypt_generic(ef->td, ciphertext, ef->hash_size);

	if (fwrite(ciphertext, 1, ef->hash_size, ef->fp) != ef->hash_size)
 	{
		set_error("fwrite");
		ret = 1;
		goto fail;
	}

	enc_fflush(ef);

fail:
    Bzero(ciphertext, ef->hash_size);
    free(ciphertext);
    
	return ret;
}

int enc_ftell(ENCFILE *ef)
{
	return ftell(ef->fp) - ef->enc_data_offset;
}

int enc_feof(ENCFILE *ef)
{
    return enc_ftell(ef) == ef->enc_data_length;
}

int enc_fseek(ENCFILE *ef, long at, int dir)
{
	/* no index - no seek */
    if(ef->idx == NULL)
    {
        set_error("use enc_add_index() before attempting to enc_fseek()");
		return EOF;
    }

	return enc_fseek_index(ef, at, dir, ef->idx);
}

int enc_fseek_index(ENCFILE *ef, long at, int dir, ENCINDEX *idx)
{
    /* convert dir & offset to absolute offset */
    switch(dir)
    {
        case SEEK_END: at += ef->enc_data_length; break;
        case SEEK_CUR: at += enc_ftell(ef); break;
        case SEEK_SET: break;
        default:
		    errno = EINVAL;
            set_error("unknown fseek() direction");
            return EOF;
    }

	/* seeking beyond the end of the file? */
	if(at < 0 || at > ef->enc_data_length)
	{
	    errno = EINVAL;
		return EOF;
	}
    
	/* decrypt from the nearest bookmarked block, to at */
	unsigned char *tmpbuf;

	/* find the nearest bookmark */
	int bookmark = at / idx->bookmark_stride;

	/* set decryptor state */
	encindex_unserialize(idx, ef->td, bookmark);

	/* fseek to correct location */
	int phyoffset = ef->enc_data_offset + bookmark * idx->bookmark_stride;
	if(fseek(ef->fp, phyoffset, SEEK_SET) == EOF)
	{
	    set_error("error doing fseek()");
	    return EOF;
	}    

	/* fread to requested location */
	int length = at % idx->bookmark_stride;
	if(length)
	{
		tmpbuf = malloc(length);
		int nread = enc_fread(tmpbuf, 1, length, ef);
		free(tmpbuf);

		if(nread != length)
		{
		    set_error("error enc_fread()ing from bookmark location to requested offset");
		    return EOF;
		}
	}

	return 0;
}

int enc_init()
{
    init_random();
    enc_initialized = TRUE;
    return 0;
}

void enc_deinit()
{
	deinit_random();
    enc_initialized = FALSE;
}