An AES C++ Wrapper

aes.hpp

#include <iostream>
#include <stdio.h>
#include <string.h>
#include <openssl/aes.h>
#include <openssl/rand.h>

#define CRYPTO_AES_KEY_LEN_128 128
#define CRYPTO_AES_KEY_LEN_192 192
#define CRYPTO_AES_KEY_LEN_256 256

#define CRYPTO_PRINT_ERROR fprintf(stderr, "[x] %s\n", strerror(errno))

class CryptoAES{
public:
  void GenerateKey(int bits){
    InitKey(bits);
    RAND_bytes((unsigned char *)key, GetKeySizeBytes());
    ClearIVs();
    RAND_bytes(iv, AES_BLOCK_SIZE);
    InitIVs();
  }
  void ClearKey(){
    memset(key, 0, key_size_bytes);
  }
  int GetKeySizeBytes(){
    return key_size_bytes;
  }
   void SetKeySize(int bits){
    key_size = bits;
    key_size_bytes = bits / 8;
  }
  int GetKeySizeBits(){
    return key_size;
  }
  unsigned char *GetIV(){
    return iv;
  }
  void *GetKey(){
    return key;
  }
  void InitKey(int bits){
    SetKeySize(bits);
    key = malloc(GetKeySizeBytes());
    ClearKey();
  }
  bool SetIV(void *src){
    if (memcpy(iv, src, AES_BLOCK_SIZE) == NULL){
      CRYPTO_PRINT_ERROR;
      return false;
    }
    InitIVs();
    return true;
  }
  bool SetKey(void *key, int bits){
    InitKey(bits);
    if (memcpy(key, key, GetKeySizeBytes())== NULL){
        return false;
      }
      return true;
  }
  void *GetPlainText(){
    return plaintext;
  }
  void InitIVs(){
    memcpy(iv_encrypt, iv, AES_BLOCK_SIZE);
    memcpy(iv_decrypt, iv, AES_BLOCK_SIZE);
  }
  void ClearIVs(){
    memset(iv, 0, AES_BLOCK_SIZE);
    memset(iv_encrypt, 0, AES_BLOCK_SIZE);
    memset(iv_decrypt, 0, AES_BLOCK_SIZE);
  }
  void SetPlainTextSize(int size){
    plaintext_size = size;
  }
  int GetPlainTextSize(){
    return plaintext_size;
  }
  void InitPlainText(){
    plaintext = malloc(GetPlainTextSize());
    ClearPlainText();
  }
  void InitCipherText(){
    ciphertext = malloc(GetCipherTextSize());
    ClearCipherText();
  }
  void SetCipherTextEncryptSize(){
    ciphertext_size = ((GetPlainTextSize()/16 + 1) * 16);
  }
  void SetPlainTextDecryptSize(){
    plaintext_size = ciphertext_size;
  }
  int GetCipherTextSize(){
    return ciphertext_size;
  }
  void SetCipherTextSize(int size){
    ciphertext_size = size;
  }
  void SetPlainText(void *src, int src_size){
    SetPlainTextSize(src_size);
    InitPlainText();
    memcpy(plaintext, src, (size_t)src_size);
  }
  void SetCipherText(void *src, int src_size){
    SetCipherTextSize(src_size);
    InitCipherText();
    memcpy(ciphertext, src, GetCipherTextSize());
  }
  void EncryptBuffer(void *src, int src_size){
    SetPlainText(src, src_size);
    SetCipherTextEncryptSize();
    InitCipherText();
    AES_set_encrypt_key((unsigned char *)key, GetKeySizeBits(), &key_encrypt);
    AES_cbc_encrypt((unsigned char *)plaintext,
                    (unsigned char *)ciphertext,
                    GetPlainTextSize(),
                    &key_encrypt,
                    iv_encrypt,
                    AES_ENCRYPT);
  }
  void Encrypt(){
    SetCipherTextEncryptSize();
    InitCipherText();
    AES_set_encrypt_key((unsigned char *)key, GetKeySizeBits(), &key_encrypt);
    AES_cbc_encrypt((unsigned char *)plaintext,
                    (unsigned char *)ciphertext,
                    GetPlainTextSize(),
                    &key_encrypt,
                    iv_encrypt,
                    AES_ENCRYPT);
  }
  void WriteCipherTextFile(char *file_name){
    WriteFile(file_name, ciphertext, GetCipherTextSize());
  }
  void WritePlainTextFile(char *file_name){
    WriteFile(file_name, plaintext, GetPlainTextSize() - AES_BLOCK_SIZE);
  }
  bool ReadKeyFile(char *file_name){
    FILE *fp = fopen(file_name, "rb");
    int file_size = GetFileSize(fp);
    if ((file_size != CRYPTO_AES_KEY_LEN_128 / 8) &&
        (file_size != CRYPTO_AES_KEY_LEN_192 / 8) &&
        (file_size != CRYPTO_AES_KEY_LEN_256 / 8)){
      fprintf(stderr,
              "[x] key size of %d bytes is not supported",
              file_size);
      return false;
    }
    InitKey(file_size * 8);
    fread(key, file_size, 1, fp);
    fclose(fp);
    return true;
  }
  bool ReadIVFile(char *file_name){
    FILE *fp = fopen(file_name, "rb");
    int file_size = GetFileSize(fp);
    if (file_size != AES_BLOCK_SIZE){
      fprintf(stderr, "[x] iv size of %d is not supported", file_size);
      return false;
    }
    ClearIVs();
    fread(iv, AES_BLOCK_SIZE, 1, fp);
    fclose(fp);
    InitIVs();
    return true;
  }
  void ReadCipherTextFile(char *file_name){
    FILE *fp = fopen(file_name, "rb");
    SetCipherTextSize(GetFileSize(fp));
    InitCipherText();
    fread(ciphertext, GetCipherTextSize(), 1, fp);
    fclose(fp);
  }
  void ReadPlainTextFile(char *file_name){
    FILE *fp = fopen(file_name, "rb");
    SetPlainTextSize(GetFileSize(fp));
    InitPlainText();
    fread(plaintext, GetPlainTextSize(), 1, fp);
    fclose(fp);
  }
  void WriteKeyFile(char *file_name){
    WriteFile(file_name, key, GetKeySizeBytes());
  }
  void WriteIVFile(char *file_name){
    WriteFile(file_name, iv, AES_BLOCK_SIZE);
  }
  void DecryptBuffer(void *src, int src_size, int plaintext_size){
    SetCipherText(src, src_size);
    SetPlainTextSize(plaintext_size);
    InitPlainText();
    AES_set_decrypt_key((unsigned char *)key, GetKeySizeBits(), &key_decrypt);
    AES_cbc_encrypt((unsigned char *)ciphertext,
                    (unsigned char *)plaintext,
                    GetCipherTextSize(),
                    &key_decrypt,
                    iv_decrypt,
                    AES_DECRYPT);
  }
  void Decrypt(int size){
    SetPlainTextSize(size + AES_BLOCK_SIZE);
    InitPlainText();
    AES_set_decrypt_key((unsigned char *)key, GetKeySizeBits(), &key_decrypt);
    AES_cbc_encrypt((unsigned char *)ciphertext,
                    (unsigned char *)plaintext,
                    GetCipherTextSize(),
                    &key_decrypt,
                    iv_decrypt,
                    AES_DECRYPT);
  }
  void ClearCipherText(){
    memset(ciphertext, 0, ciphertext_size);
  }
  void ClearPlainText(){
    memset(plaintext, 0, GetPlainTextSize());
  }
  void CleanupKey(){
    ClearKey();
    free(key);
  }
  void CleanupCipherText(){
    ClearCipherText();
    free(ciphertext);
  }
  void CleanupPlainText(){
    ClearPlainText();
    free(plaintext);
  }
  void Cleanup(){
    CleanupKey();
    ClearIVs();
    CleanupCipherText();
    CleanupPlainText();
  }
  void DebugDecryptedPlainText(){
    hex_dump((char *)"plaintext", plaintext, GetPlainTextSize() - AES_BLOCK_SIZE);
  }
  void DebugPlainText(){
    hex_dump((char *)"plaintext", plaintext, GetPlainTextSize());
  }
  void DebugCipherText(){
    hex_dump((char *)"ciphertext", ciphertext, GetCipherTextSize());
  }
  void DebugIV(){
    hex_dump((char *)"iv", iv, AES_BLOCK_SIZE);
  }
  void DebugKey(){
    hex_dump((char *)"key", key, GetKeySizeBytes());
  }
private:
  void hex_dump(char *desc, void *addr, size_t len){
    int i;
    unsigned char buff[17];
    const unsigned char *pc = (const unsigned char*)addr;
    if (desc != NULL){
      printf ("%s:\n", desc);
    }
    for (i = 0; i < (int)len; i++) {
      if ((i % 16) == 0) {
        if (i != 0){
          printf ("  %s\n", buff);
        }
        printf ("  %04x ", i);
      }
      printf (" %02x", pc[i]);
      if ((pc[i] < 0x20) || (pc[i] > 0x7e)){
        buff[i % 16] = '.';
      } else{
        buff[i % 16] = pc[i];
      }
      buff[(i % 16) + 1] = '\0';
    }
    while ((i % 16) != 0){
      printf ("   ");
      i++;
    }
    printf ("  %s\n", buff);
  }
  void WriteFile(char *file_name, void *data, int size){
    FILE *fp = fopen(file_name, "wb");
    fwrite(data, size, 1, fp);
    fclose(fp);
  }
  int GetFileSize(FILE *fp){
    int prev=ftell(fp);
    fseek(fp, 0L, SEEK_END);
    int sz=ftell(fp);
    fseek(fp,prev,SEEK_SET);
    return sz;
  }
  void *ciphertext;
  size_t ciphertext_size;
  void *plaintext;
  size_t plaintext_size;
  void *key;
  AES_KEY key_encrypt;
  AES_KEY key_decrypt;
  int key_size;
  int key_size_bytes;
  unsigned char iv[AES_BLOCK_SIZE];
  unsigned char iv_encrypt[AES_BLOCK_SIZE];
  unsigned char iv_decrypt[AES_BLOCK_SIZE];
};