aes/aes.hpp

#ifndef _AES
#define _AES

#include <iostream>
#include <stdio.h>
#include <string>
#include <sstream>
#include <string.h>
#include <algorithm>
#include <typeinfo>
#include <regex>
#include <fstream>

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

#include <json/json.h>

#include "lib/base64/base64.hpp"

#define FAILURE -1
#define SUCCESS 0

#define AES_KEYLEN 256
#define AES_ROUNDS 6

class aes {
protected:
    unsigned char* aesKey;
    unsigned char* aesIV;
    aes& clone(const aes& _e);
    void init_all();
    void clear_all();
    void setExtension(std::string _ex);
private:
    std::string filename;
    std::string generateNewFilename();
    std::string getFilePath();
    std::string extension = ".new";

public:
    aes() {};
    aes(std::string _filename);
    static std::string writeFile(unsigned char *file, size_t fileLength, const char* Filename);
    std::string writeFile(unsigned char *file, size_t fileLength);
    static int readFile(unsigned char** file, const char* filename);
    int readFile(unsigned char **file);

    std::string exportKey(const char* filename);
    std::string exportKey();
    void importKey(const char* filename);

    std::string PrintAesKey();
    void setAesIV(const unsigned char* _aesIV);
    void setAesKey(const unsigned char* _aesKey);
    unsigned char* getAesKey();
    unsigned char* getAesIV();

    virtual ~aes() = default;
};

aes::aes(std::string _filename) {
    filename = _filename;
    aesKey = (unsigned char*)malloc(AES_KEYLEN/8);
    aesIV = (unsigned char*)malloc(AES_KEYLEN/8);
}

aes& aes::clone(const aes& _e) {
    filename = _e.filename;

    setAesKey(_e.aesKey);
    setAesIV(_e.aesIV);

    return *this;
}

void aes::init_all() {

    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 << "Use of USE_PBKDF parameter" << 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);
}
int aes::readFile(unsigned char** file) {
    return readFile(file, filename.c_str());
}

// peut être déporté le buffer lut avec methode pour travaillé dessus
int aes::readFile(unsigned char** file, const char* filename) {
    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);                         // @FIXME : failed for bigest file // maybe with new
    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;
}

std::string aes::getFilePath() {
    std::regex e("\\b(.*[(\\.txt)]*)\\.[^(txt)].*");
    filename = std::regex_replace(filename, e,"$1");
    return filename;
}

std::string aes::generateNewFilename() {
    return aes::getFilePath()+extension;
}

std::string aes::writeFile(unsigned char *file, size_t fileLength) {
    std::string newFile = generateNewFilename();
    return writeFile(file, fileLength, newFile.c_str());
}

std::string aes::writeFile(unsigned char *file, size_t fileLength, const char* Filename) {
    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);
    return Filename;
}

void aes::setExtension(std::string _ex) {
    extension = _ex;
}

void aes::setAesKey(const unsigned char* _aesKey) {
    if(aesKey == NULL)
        aesKey = (unsigned char*)malloc(AES_KEYLEN/8);
    memcpy(aesKey , _aesKey, AES_KEYLEN/8);
}

void aes::setAesIV(const unsigned char* _aesIV) {
    if(aesIV == NULL)
        aesIV = (unsigned char*)malloc(AES_KEYLEN/8);
    memcpy(aesIV , _aesIV, AES_KEYLEN/8);
}


unsigned char* aes::getAesKey() {
    unsigned char* res = (unsigned char*)malloc(AES_KEYLEN/8);
    memcpy(res, aesKey, AES_KEYLEN/8);
    return res;
}

unsigned char* aes::getAesIV() {
    unsigned char* res = (unsigned char*)malloc(AES_KEYLEN/8);
    memcpy(res, aesIV, AES_KEYLEN/8);
    return res;
}

std::string aes::PrintAesKey() {
    std::stringstream ss;
    std::string res;

    res = "aesKey : ";
    res += base64_encode(aesKey, AES_KEYLEN/8);
    return res;
}

void aes::importKey(const char* filename) {
    Json::Value root;

    std::ifstream keyfile(filename, std::ifstream::binary);
    keyfile >> root;

    setAesKey(reinterpret_cast<const unsigned char*>(base64_decode(root["key"].asString()).c_str()));
    setAesIV(reinterpret_cast<const unsigned char*>(base64_decode(root["IV"].asString()).c_str()));
}

std::string aes::exportKey() {
    return exportKey((getFilePath()+".key").c_str());
}

std::string aes::exportKey(const char* filename) {
    // export the tuples aeskey and aesIV to json file
    Json::Value root;

    root["key"] = base64_encode(aesKey, AES_KEYLEN/8);
    root["IV"] = base64_encode(aesIV, AES_KEYLEN/8);

    std::ofstream out(filename);
    out << root;
    out.close();
    return std::string(filename);
}

void aes::clear_all() {
    free(aesIV);
    free(aesKey);
}

#endif