#include <AESLib.h>

byte aesKey[16] = { 0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66};
AESLib aesLib;

void generateRandomIV(byte* iv, int length) {
  for (int i = 0; i < length; i++) {
    iv[i] = random(0, 256);
  }
}

void printByteArray(byte* arr, int length) {
  for (int i = 0; i < length; i++) {
    if (arr[i] < 0x10) {
      Serial.print("0");
    }
    Serial.print(arr[i], HEX);
    Serial.print(" ");
  }
  Serial.println();
}

void uint64ToBytes(uint64_t value, byte* bytes) {
  bytes[0] = (byte)(value >> 56);
  bytes[1] = (byte)(value >> 48);
  bytes[2] = (byte)(value >> 40);
  bytes[3] = (byte)(value >> 32);
  bytes[4] = (byte)(value >> 24);
  bytes[5] = (byte)(value >> 16);
  bytes[6] = (byte)(value >> 8);
  bytes[7] = (byte)value;
}

void encrypt(uint64_t message, byte* cipherText, int& outputLength) {
  byte messageBytes[8] = {0};
  uint64ToBytes(message, messageBytes);

  Serial.print("Message : ");
  printByteArray(messageBytes, 8);

  int messageLength = sizeof(messageBytes);
  int blockSize = 16;
  int paddedLength = ((messageLength / blockSize) + 1) * blockSize;
  byte paddedMessage[paddedLength];

  // PKCS#7 padded
  memcpy(paddedMessage, messageBytes, messageLength);
  byte paddingValue = paddedLength - messageLength;
  for (int i = messageLength; i < paddedLength; i++) {
    paddedMessage[i] = paddingValue;
  }

  Serial.print("ENC Padded Message : ");
  printByteArray(paddedMessage, paddedLength);

  byte encIV[16] = {0};
  byte encryptedMessage[16] = {0};

  generateRandomIV(encIV, 16);

  // IV changed after use by encrypt, copy it here first
  memcpy(cipherText, encIV, 16);
  
  Serial.print("ENC IV : ");
  printByteArray(encIV, 16);

  //aesLib.set_paddingmode(paddingMode::CMS);
  outputLength = aesLib.encrypt(paddedMessage, paddedLength, encryptedMessage, aesKey, 128, encIV);

  memcpy(cipherText + 16, encryptedMessage, 16);

  Serial.print("ENC Message : ");
  printByteArray(encryptedMessage, 16);

  Serial.print("ENC Final : ");
  printByteArray(cipherText, 32);

  // Set output length
  outputLength += paddedLength;
}

void decrypt(byte* message, byte* output) {
  byte decIV[16] = {0};
  byte padded[16] = {0};
  byte decryptedMessage[16] = {0};

  memcpy(decIV, message, 16);
  memcpy(padded, message + 16, 16);
  
  // byte input[], uint16_t input_length, byte *plain, const byte key[], int bits, byte my_iv[]
  aesLib.decrypt(padded, sizeof(padded), decryptedMessage, aesKey, 128, decIV);
  memcpy(output, decryptedMessage, 8);

  Serial.print("DEC IV : ");
  printByteArray(decIV, 16);
  Serial.print("DEC Padded : ");
  printByteArray(decryptedMessage, 16);
  Serial.print("DEC Message : ");
  printByteArray(output, 8);
}

void setup() {
  // put your setup code here, to run once:
  Serial.begin(9600);
  randomSeed(analogRead(0));

  uint64_t data = (uint64_t)0xFFFFFFFFFFFFFFFF;
  int outputLength;
  byte cipherText[32] = {0};
  byte plainText[8] = {0};
  encrypt(data, cipherText, outputLength);
  Serial.println("=============");
  decrypt(cipherText, plainText);
}

void loop() {
  // put your main code here, to run repeatedly:

}