M2-PT-DRP/source/utils/crypto/aes/AesKey.hpp

#pragma once

#include <array>
#include <cstdint>
#include <functional>
#include <memory>
#include <random>
#include <openssl/evp.h>


namespace drp::util::crypto {


/**
 * Represent an AES key.
 * Allow for encrypting and decrypting data.
 * @tparam keySize the size of the key (in bytes)
 * @tparam ivSize the size of the initialisation vector (in bytes)
 */
template<std::size_t keySize, std::size_t ivSize, const EVP_CIPHER*(cipherFunction)()>
class AesKey {
public:
    /**
     * Create a random AES key.
     */
    AesKey() {
        // generate a random key
        for (auto& byte : this->_data)
            byte = randomDistribution(randomGenerator);
    }

    explicit AesKey(const std::array<std::uint8_t, keySize>& data) {
        this->_data = data;
    }

    /**
     * Encrypt data with this key.
     * @param plainData the data to encrypt
     * @return the encrypted data. It will always be longer than the original data.
     */
    [[nodiscard]] std::vector<std::uint8_t> encrypt(const std::vector<std::uint8_t>& plainData) const {
        // create an initialization vector
        std::array<std::uint8_t, ivSize> iv {};
        for (auto& byte : iv)
            byte = randomDistribution(randomGenerator);

        // create the cipher context
        const auto context = std::unique_ptr<EVP_CIPHER_CTX, decltype(&EVP_CIPHER_CTX_free)>(
            EVP_CIPHER_CTX_new(),
            EVP_CIPHER_CTX_free
        );

        if (context == nullptr)
            throw std::runtime_error("Error creating EVP_CIPHER_CTX");

        // initialize the encryptor
        if (EVP_EncryptInit_ex(
            context.get(),
            cipherFunction(),
            nullptr,
            this->_data.data(),
            iv.data()
        ) != 1)
            throw std::runtime_error("Error initializing encryption");

        std::vector<std::uint8_t> encryptedData(ivSize + plainData.size() + EVP_CIPHER_block_size(EVP_aes_256_cbc()));
        std::copy(iv.begin(), iv.end(), encryptedData.begin());
        int length;

        // encrypt the data
        if (EVP_EncryptUpdate(
            context.get(),
            encryptedData.data() + ivSize,
            &length,
            plainData.data(),
            static_cast<int>(plainData.size())
        ) != 1)
            throw std::runtime_error("Error encrypting data");
        int encryptedDataLength = length;

        // finalize the encryption
        if (EVP_EncryptFinal_ex(
            context.get(),
            encryptedData.data() + ivSize + encryptedDataLength,
            &length
        ) != 1)
            throw std::runtime_error("Error finalizing encryption");
        encryptedDataLength += length;
        encryptedData.resize(ivSize + encryptedDataLength);

        return encryptedData;
    }

    /**
     * Decrypt data with this key.
     * @param rawEncryptedData the encrypted data to decrypt.
     * @return the decrypted data.
     */
    [[nodiscard]] std::vector<std::uint8_t> decrypt(const std::vector<std::uint8_t>& rawEncryptedData) const {
        // create a cipher context
        const auto context = std::unique_ptr<EVP_CIPHER_CTX, decltype(&EVP_CIPHER_CTX_free)>(
            EVP_CIPHER_CTX_new(),
            EVP_CIPHER_CTX_free
        );

        std::array<std::uint8_t, ivSize> iv;
        std::copy(rawEncryptedData.begin(), rawEncryptedData.begin() + ivSize, iv.data());
        const std::vector encryptedData(rawEncryptedData.begin() + ivSize, rawEncryptedData.end());

        // initialize the decryptor
        if (EVP_DecryptInit_ex(
            context.get(),
            cipherFunction(),
            nullptr,
            this->_data.data(),
            iv.data()
        ) != 1)
            throw std::runtime_error("Error initializing decryptor");

        std::vector<std::uint8_t> plainData(encryptedData.size());
        int length;

        // decrypt the data
        if (EVP_DecryptUpdate(
            context.get(),
            plainData.data(),
            &length,
            encryptedData.data(),
            static_cast<int>(encryptedData.size())
        ) != 1)
            throw std::runtime_error("Error decrypting data");

        // finalize the decryption
        if (EVP_DecryptFinal_ex(
            context.get(),
            plainData.data(),
            &length
        ) != 1)
            throw std::runtime_error("Error finalizing decryptor");
        plainData.resize(length);

        return plainData;
    }

private:
    static std::mt19937 randomGenerator;
    static std::uniform_int_distribution<std::uint8_t> randomDistribution;

    std::array<std::uint8_t, keySize> _data;
};


template<std::size_t keySize, std::size_t ivSize, const EVP_CIPHER*(cipherFunction)()>
std::mt19937 AesKey<keySize, ivSize, cipherFunction>::randomGenerator = std::mt19937(std::random_device{}());

template<std::size_t keySize, std::size_t ivSize, const EVP_CIPHER*(cipherFunction)()>
std::uniform_int_distribution<std::uint8_t> AesKey<keySize, ivSize, cipherFunction>::randomDistribution = std::uniform_int_distribution(
    std::numeric_limits<std::uint8_t>::min(),
    std::numeric_limits<std::uint8_t>::max()
);


using AesKey256 = AesKey<256/8, 16, EVP_aes_256_cbc>;


}