#ifndef _CRYPT
#define _CRYPT

#include "aes.hpp"

class crypt : public aes {
private:
	EVP_CIPHER_CTX* aesEncryptCtx;

public:
	crypt(std::string filename);
	crypt(const crypt& a);
	void init_all();
	int aesEncrypt(const unsigned char *msg, size_t msgLen, unsigned char **encMsg);
	~crypt();
};

crypt::crypt(std::string filename) : aes(filename) {
    aesEncryptCtx = EVP_CIPHER_CTX_new();
    //aesEncryptCtx = (EVP_CIPHER_CTX*)malloc(sizeof(EVP_CIPHER_CTX));
    if(aesEncryptCtx == NULL) {
        exit(FAILURE);
    }
    init_all();
}

crypt::crypt(const crypt& _a) : aes::aes() {
    this->clone(_a);
    aesEncryptCtx = EVP_CIPHER_CTX_new();
    //aesEncryptCtx = (EVP_CIPHER_CTX*)malloc(sizeof(EVP_CIPHER_CTX));
    if(aesEncryptCtx == NULL) {
        exit(FAILURE);
    }

	int res = EVP_CIPHER_CTX_copy(aesEncryptCtx, _a.aesEncryptCtx);
    if(!res)
    	exit(FAILURE);
}

void crypt::init_all() {
	setExtension(".enc");
	aes::init_all();
	EVP_CIPHER_CTX_init(aesEncryptCtx);
}

int crypt::aesEncrypt(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, aes::aesKey, aes::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;
    }

    return encMsgLen + blockLen;
}

crypt::~crypt() {
	aes::clear_all();
	EVP_CIPHER_CTX_cleanup(aesEncryptCtx);
    EVP_CIPHER_CTX_free(aesEncryptCtx);
    //free(aesEncryptCtx);
}

#endif