#include <iostream>
#include <stdio.h>
#include <string>
#include <string.h>

#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/aes.h>
#include <openssl/err.h>
#include <openssl/rand.h>

#define FAILURE -1
#define SUCCESS 0

#define AES_KEYLEN 256
#define AES_ROUNDS 6


int aesDecrypt(EVP_CIPHER_CTX* aesDecryptCtx, unsigned char* aesKey, unsigned char* aesIV, unsigned char *encMsg, size_t encMsgLen, unsigned char **decMsg) {
    size_t decLen   = 0;
    size_t blockLen = 0;

    *decMsg = (unsigned char*)malloc(encMsgLen);
    if(*decMsg == NULL) return FAILURE;

    if(!EVP_DecryptInit_ex(aesDecryptCtx, EVP_aes_256_cbc(), NULL, aesKey, aesIV)) {
        return FAILURE;
    }

    if(!EVP_DecryptUpdate(aesDecryptCtx, (unsigned char*)*decMsg, (int*)&blockLen, encMsg, (int)encMsgLen)) {
        return FAILURE;
    }
    decLen += blockLen;

    if(!EVP_DecryptFinal_ex(aesDecryptCtx, (unsigned char*)*decMsg + decLen, (int*)&blockLen)) {
        return FAILURE;
    }
    decLen += blockLen;

    EVP_CIPHER_CTX_cleanup(aesDecryptCtx);

    return (int)decLen;
}

int aesEncrypt(EVP_CIPHER_CTX* aesEncryptCtx, unsigned char* aesKey, unsigned char* aesIV, const unsigned char *msg, size_t msgLen, unsigned char **encMsg) {
    size_t blockLen  = 0;
    size_t encMsgLen = 0;

    *encMsg = (unsigned char*)malloc(msgLen + AES_BLOCK_SIZE);
    if(encMsg == NULL) return FAILURE;

    if(!EVP_EncryptInit_ex(aesEncryptCtx, EVP_aes_256_cbc(), NULL, aesKey, aesIV)) {
        return FAILURE;
    }

    if(!EVP_EncryptUpdate(aesEncryptCtx, *encMsg, (int*)&blockLen, (unsigned char*)msg, msgLen)) {
        return FAILURE;
    }
    encMsgLen += blockLen;

    if(!EVP_EncryptFinal_ex(aesEncryptCtx, *encMsg + encMsgLen, (int*)&blockLen)) {
        return FAILURE;
    }

    EVP_CIPHER_CTX_cleanup(aesEncryptCtx);

    return encMsgLen + blockLen;
}

void clear_all(EVP_CIPHER_CTX* aesEncryptCtx, EVP_CIPHER_CTX* aesDecryptCtx, unsigned char* aesKey, unsigned char* aesIV) {
	EVP_CIPHER_CTX_cleanup(aesEncryptCtx);
	EVP_CIPHER_CTX_cleanup(aesDecryptCtx);
 
	free(aesEncryptCtx);
    free(aesDecryptCtx);

    free(aesIV);
    free(aesKey);
}

void init_all(EVP_CIPHER_CTX* aesEncryptCtx, EVP_CIPHER_CTX* aesDecryptCtx, unsigned char* aesKey, unsigned char* aesIV) {
	EVP_CIPHER_CTX_init(aesEncryptCtx);
    EVP_CIPHER_CTX_init(aesDecryptCtx);

    unsigned char *aesPass = (unsigned char*)malloc(AES_KEYLEN/8);
    unsigned char *aesSalt = (unsigned char*)malloc(8);

    if(aesKey == NULL || aesIV == NULL || aesPass == NULL || aesSalt == NULL) {
    	exit(FAILURE);
    }

    #define USE_PBKDF
    #ifdef USE_PBKDF
    	std::cerr << "utilisation de USE_PBKDF" << std::endl;
        // Get some random data to use as the AES pass and salt
        if(RAND_bytes(aesPass, AES_KEYLEN/8) == 0) {
            exit(FAILURE);
        }

        if(RAND_bytes(aesSalt, 8) == 0) {
            exit(FAILURE);
        }

        if(EVP_BytesToKey(EVP_aes_256_cbc(), EVP_sha256(), aesSalt, aesPass, AES_KEYLEN/8, AES_ROUNDS, aesKey, aesIV) == 0) {
            exit(FAILURE);
        }
    #else
        if(RAND_bytes(aesKey, AES_KEYLEN/8) == 0) {
            exit(FAILURE);
        }

        if(RAND_bytes(aesIV, AES_KEYLEN/8) == 0) {
            exit(FAILURE);
        }
    #endif

    free(aesPass);
    free(aesSalt);
}

void writeFile(char *filename, unsigned char *file, size_t fileLength) {
    FILE *fd = fopen(filename, "wb");
    if(fd == NULL) {
        fprintf(stderr, "Failed to open file: %s\n", strerror(errno));
        exit(1);
    }

    size_t bytesWritten = fwrite(file, 1, fileLength, fd);

    if(bytesWritten != fileLength) {
        fprintf(stderr, "Failed to write file\n");
        exit(1);
    }

    fclose(fd);
}

int readFile(char *filename, unsigned char **file) {
    FILE *fd = fopen(filename, "rb");
    if(fd == NULL) {
        fprintf(stderr, "Failed to open file: %s\n", strerror(errno));
        exit(1);
    }

    // Determine size of the file
    fseek(fd, 0, SEEK_END);
    size_t fileLength = ftell(fd);
    fseek(fd, 0, SEEK_SET);

    // Allocate space for the file
    *file = (unsigned char*)malloc(fileLength);
    if(*file == NULL) {
        fprintf(stderr, "Failed to allocate memory\n");
        exit(1);
    }

    // Read the file into the buffer
    size_t bytesRead = fread(*file, 1, fileLength, fd);

    if(bytesRead != fileLength) {
        fprintf(stderr, "Error reading file\n");
        exit(1);
    }

    fclose(fd);

    return fileLength;
}

int main(int argc, char* argv[]) {
	if(argc != 2) {
        fprintf(stderr, "No file argument supplied.\n");
        return 1;
    }

	EVP_CIPHER_CTX* aesEncryptCtx = (EVP_CIPHER_CTX*)malloc(sizeof(EVP_CIPHER_CTX));
	EVP_CIPHER_CTX* aesDecryptCtx = (EVP_CIPHER_CTX*)malloc(sizeof(EVP_CIPHER_CTX));
	if(aesDecryptCtx == NULL || aesEncryptCtx == NULL) {
		return FAILURE;
	}

	unsigned char* aesKey = (unsigned char*)malloc(AES_KEYLEN/8);
	unsigned char* aesIV = (unsigned char*)malloc(AES_KEYLEN/8);

	init_all(aesEncryptCtx, aesDecryptCtx, aesKey, aesIV);

	//***************************************************
	char* filename = argv[1];

    // Read the file to encrypt
    unsigned char *file;
    size_t fileLength = readFile(filename, &file);
    printf("%d bytes to be encrypted\n", (int)fileLength);

    // Encrypt the file
    unsigned char *encryptedFile;
    int encryptedFileLength;
    if((encryptedFileLength = aesEncrypt(aesEncryptCtx, aesKey, aesIV, (const unsigned char*)file, fileLength, &encryptedFile)) == -1) {
        fprintf(stderr, "Encryption failed\n");
        return 1;
    }
    printf("%d bytes encrypted\n", encryptedFileLength);

    // Append .enc to the filename
    char *encryptedFilename = (char*)malloc(strlen(filename) + 5);
    if(encryptedFilename == NULL) {
        fprintf(stderr, "Failed to allocate memory\n");
        return 1;
    }
    sprintf(encryptedFilename, "%s.enc", filename);

    // Write the encrypted file to its own file
    writeFile(encryptedFilename, encryptedFile, encryptedFileLength);
    
    std::cerr << "aesKey : ";
	for (int i = 0; i < 32; i++) {
		std::cerr << std::hex << aesKey[i];
	}
    std::cerr << std::endl;

    printf("Encrypted message written to \"%s\"\n", encryptedFilename);

    free(file);
	//***************************************************

    fileLength = readFile(encryptedFilename, &file);

    // Decrypt the encrypted file
    unsigned char *decryptedFile;
    int decryptedFileLength;
    if((decryptedFileLength = aesDecrypt(aesDecryptCtx, aesKey, aesIV, file, fileLength, &decryptedFile)) == -1) {
        fprintf(stderr, "Decryption failed\n");
        return 1;
    }
    printf("%d bytes decrypted\n", (int)decryptedFileLength);

    // Append .dec to the filename
    char *decryptedFilename = (char*)malloc(strlen(filename) + 5);
    if(decryptedFilename == NULL) {
        fprintf(stderr, "Failed to allocate memory\n");
        return 1;
    }
    sprintf(decryptedFilename, "%s.dec", filename);
    
    // Write the decrypted file to its own file
    writeFile(decryptedFilename, decryptedFile, decryptedFileLength);
    printf("Decrypted file written to \"%s\"\n", decryptedFilename);

    free(decryptedFile);
    free(file);


    //***************************************************
    clear_all(aesEncryptCtx, aesDecryptCtx, aesKey, aesIV);
	return 0;
}