#include <AESLib.h>

AESLib aesLib;

const char* plaintext = "Looks like key but it's not me.";
int loopcount = 0;

char cleartext[256] = {0};
char ciphertext[512];

byte aes_key[] = { 0x4d, 0x33, 0x6E, 0x74, 0x30, 0x72, 0x5F, 0x32, 0x30, 0x32, 0x32, 0, 0, 0, 0, 0 };
// AES Encryption Key

// General initialization vector (you must use your own IV's in production for full security!!!)
byte aes_iv[N_BLOCK] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };

// Sample strings as generated by node.js server
String server_b64iv = "AAAAAAAAAAAAAAAAAAAAAAAA=="; // same as aes_iv  but in Base-64 form as received from server
String server_b64msg = "j0RFVdlKjYrwx17qzHdt40ZS4hxckx0riP4SNy21X3U="; // CBC/Zeropadding; same as aes_iv  but in Base-64 form as received from server




void print_key_iv() {
  Serial.print("AES IV: ");
  for (unsigned int i = 0; i < sizeof(aes_iv); i++) {
    Serial.print(aes_iv[i], DEC);
    if ((i + 1) < sizeof(aes_iv)) {
      Serial.print(",");
    }
  }

  Serial.println("");
}

String encrypt_impl(char * msg, byte iv[]) {
  int msgLen = strlen(msg);
  char encrypted[2 * msgLen] = {0};
  aesLib.encrypt64((const byte*)msg, msgLen, encrypted, aes_key, sizeof(aes_key), iv);
  return String(encrypted);
}

String decrypt_impl(char * msg, byte iv[]) {
  int msgLen = strlen(msg);
  char decrypted[msgLen] = {0}; // half may be enough
  aesLib.decrypt64(msg, msgLen, (byte*)decrypted, aes_key, sizeof(aes_key), iv);
  return String(decrypted);
}

// Generate IV (once)
void aes_init() {

  Serial.flush();
  
  delay(1000);

  Serial.println("\n=======\n");

  ///
  
  Serial.println("\n1) AES init... paddingMode::ZeroLength");  
  aesLib.set_paddingmode(paddingMode::ZeroLength);  

  byte enc_iv_A[N_BLOCK] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
  Serial.println("Encrypting \"Looks like key but it's not me.\" using null-IV with ZeroLength padding");
  String encrypted1 = encrypt_impl((char*)plaintext, enc_iv_A);
  Serial.print("Encrypted(1): "); Serial.println(encrypted1);
  print_key_iv();

  aesLib.set_paddingmode(paddingMode::ZeroLength);  

  byte dec_iv_B[N_BLOCK] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
  Serial.println("Decrypting \"Looks like key but it's not me.\" using null-IV ZeroLength padding");
  String decrypted = decrypt_impl((char*)encrypted1.c_str(), dec_iv_B); // aes_iv fails here, incorrectly decoded...
  Serial.print("Cleartext: ");
  Serial.println(decrypted);
  Serial.println("In first iteration this should work (using untouched dec_iv_B) ^^^");

  ///
  
  Serial.println("\n2) AES init... paddingMode::CMS");
  aesLib.set_paddingmode(paddingMode::CMS);  

  Serial.println("Encrypting \"Looks like key but it's not me.\" using null-IV with CMS padding");
  byte enc_iv_X[N_BLOCK] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
  String encrypted2 = encrypt_impl((char*)plaintext, enc_iv_X );
  Serial.print("Encrypted (2): "); Serial.println(encrypted2);
  print_key_iv();

  aesLib.set_paddingmode(paddingMode::CMS);

  Serial.println("Decrypting \"Looks like key but it's not me.\" using null-IV CMS padding");
  byte enc_iv_Y[N_BLOCK] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
  decrypted = decrypt_impl((char*)encrypted2.c_str(), enc_iv_Y);
  Serial.print("Cleartext: ");
  Serial.println(decrypted);
  
  ///

  Serial.println("\n3) AES init... from Server, paddingMode::CMS");
  aesLib.set_paddingmode(paddingMode::CMS);  

  // Decode IV from server to aes_iv instead of directly using zeros...
  int ivLen = base64_decode((char*)server_b64iv.c_str(), (char *)aes_iv, server_b64iv.length());
  Serial.println("Server IV should be null-IV: ");
  print_key_iv();
  Serial.print("Decoded Server IV bytes to aes_iv: "); Serial.println(ivLen);

  Serial.print("B64 Ciphertext: "); Serial.println(server_b64msg);  

  String decrypted_string = decrypt_impl((char*)server_b64msg.c_str(), aes_iv);

  Serial.print("Server message decrypted using server IV and CMS, cleartext: ");
  Serial.println(decrypted_string);
}

void setup() 
{
  // put your setup code here, to run once:
  Serial.begin(115200);
  Serial.println("Hello, STM32!");
  aes_init();
}

void loop() {

  loopcount++;

  if (loopcount > 5) return; // prevent week-long logs
  
  Serial.println("");

  //sprintf(cleartext, "START; %i \n", loopcount);
  sprintf(cleartext, "Looks like key but it's not me.");

  aesLib.set_paddingmode(paddingMode::CMS);

  // Encrypt Data
  byte enc_iv[N_BLOCK] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 }; // iv_block gets written to, provide own fresh copy...
  String encrypted = encrypt_impl(cleartext, enc_iv);
  sprintf(ciphertext, "%s", encrypted.c_str());
  Serial.print("Base64 encoded Ciphertext: ");
  Serial.println(encrypted);

  // Decrypt Data
  byte dec_iv[N_BLOCK] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 }; // iv_block gets written to, provide own fresh copy...
  String decrypted = decrypt_impl(ciphertext, dec_iv);
  Serial.print("Base64-decoded Cleartext: ");
  Serial.println(decrypted);

  delay(5000);

  // Reset MCU to restart from beggining
  while(true);

}