#include "aes.hpp"

enum process_flag_t {
  AES_DEC,
  AES_ENC
};

/* Key is: 47f1ed9166c996b2f553b147be3fbc20 (128-bit)*/
const uint8_t key[] = {0x47, 0xF1, 0xED, 0x91, 0x66, 0xC9, 0x96, 0xB2, 
                       0xF5, 0x53, 0xB1, 0x47, 0xBE, 0x3F, 0xBC, 0x20};

AES_ctx ctx;

void setup() {
  Serial.begin(115200);
  Serial.setTimeout(50);

  pinMode(D13, OUTPUT);
  digitalWrite(D13, LOW);
  
  AES_init_ctx(&ctx, key);
}

uint8_t cLUT(char c) {
  return c >= 'A' ? c - 'A' + 10 : c - '0';
}

char vLUT(uint8_t n) {
  return n >= 10 ? n + 'A' - 10 : n + '0';;
}

String process(String& str, process_flag_t flag) {
  String res("0x");

  uint8_t buffer[AES_BLOCKLEN];
  memset(buffer, 0, AES_BLOCKLEN);
  str.remove(0, 6);

  for (int i = 0; i < str.length(); i += 2) {
    buffer[i / 2] = (cLUT(str.charAt(i)) << 4) | cLUT(str.charAt(i + 1)); 
  }

  digitalWrite(D13, HIGH);
  switch (flag) {
    case AES_ENC: {
      AES_ECB_encrypt(&ctx, buffer);
      break;
    }
    case AES_DEC: {
      AES_ECB_decrypt(&ctx, buffer);
      break;
    }
  }
  digitalWrite(D13, LOW);
  
  for (int i = 0; i < AES_BLOCKLEN; ++i) {
    res += vLUT((buffer[i] >> 4) & 0xF);
    res += vLUT(buffer[i] & 0xF);
  }


  return res;
}

void loop() {
  String msg = Serial.readString();
  if (msg.length() > 6) {
    if (msg.startsWith("DEC")) {
      Serial.println(process(msg, AES_DEC));
    }

    if (msg.startsWith("ENC")) {
      Serial.println(process(msg, AES_ENC));
    }
  }

  delay(5);
}