/*
 * Copyright (C) 2026 SCHAUAUS GmbH
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included
 * in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */

// compiler values definition
#define PRINT_SERIAL       1
#define DEBUG              0
#define USE_DISPLAY        1
#define HAS_KEYBOARD       0

#define BTN_HOME           9

#define MAX_INPUT_LENGHT 256

// external libraries
#include <Crypto.h>
#include <AES.h>
#include <string.h>

// project files to include
#include "display.h"

//creating an object of AES256 class
AES256_ESP aes256;

void printHelp () {
  #if PRINT_SERIAL
    Serial.println();
    Serial.println("┌-----------------------------------------------------┐");
    Serial.println("| Encrypt / Decrypt seed phrase with AES256 algorithm |");
    Serial.println("└-----------------------------------------------------┘");
    Serial.println();
    Serial.println("WARNING: DO NOT USE THIS ON ANY NETWORK CONNECTED DEVICE AND ENSURE THE DECRYPTION IS WORKING PROPERLY!");
    Serial.println("USE IN SECURE ENVIRONMENTS ONLY!");
    Serial.println("This program is offered as it is and no warranties will be granted in any case!");
    Serial.println("Review the source code thoroughly before using it for real application!");
    Serial.println("Available commands:");
    Serial.println("help                                      - print this help to serial console");
    Serial.println("key                                       - generate random HEX key and print it to the serial console");
    Serial.println("key <your 32byte key>                     - set key (hexadecimal or ascii) for encryption / decryption");
    Serial.println("encrypt <12 or 24 word seed phrase>       - encrypts your 12 ord 24 word plain text seed phrase to HEX string");
    Serial.println("decrypt <encrypted HEX string>            - decrypts your encrypted HEX string back to clear text seed phrase");
    Serial.println("clear                                     - clear all cryptographic information from memory");
    Serial.println("restart                                   - reboot the device for a fresh restart; any working data will be removed from memory");
  #endif

  #if USE_DISPLAY
    display_help();
  #endif
}

void byteToHex (byte b, char* returnVal) {
  static const char hexChars[] = "0123456789ABCDEF";

  //char returnVal[4];
  returnVal[0] = hexChars[b >> 4];
  returnVal[1] = hexChars[b & 0x0F];
  returnVal[2] = ' ';
  returnVal[3] = '\0';
  #if PRINT_SERIAL
    Serial.print(returnVal);
  #endif
}

void setCryptKey (String argKey) {
  #if PRINT_SERIAL
    Serial.println();
    Serial.println("Operation: setting encryption key");
  #endif

  argKey.trim();
  String tempKey = argKey;
  byte cryptKey[32];

  tempKey.replace(" ", ""); // remove any spaces in between characters to check for hexadecimal values

  if(argKey.length() != 32 && tempKey.length() / 2 != 32) { // there was no valid key given; keys need to consist of 32 characters (256bit)
    #if PRINT_SERIAL
      Serial.println("No valid key handed over; generating new key:");
    #endif

    for(int i=0; i < 32; i++) {
      cryptKey[i] = random(256);
    }
  }
  else { // process valid key provided
    #if PRINT_SERIAL
      Serial.println("Valid key detected");
    #endif

    bool keyIsHex = true;
    for(size_t i = 0; i < tempKey.length(); i++)
      if ( !(tempKey[i] >= '0' && tempKey[i] <= '9') &&
           !(tempKey[i] >= 'a' && tempKey[i] <= 'f') &&
           !(tempKey[i] >= 'A' && tempKey[i] <= 'F') )
             keyIsHex = false;

    if (keyIsHex) { // input key is hexadecimal representation
      String byteString = "";
      for(size_t idx = 0; idx < tempKey.length(); idx += 2) {
        byteString = tempKey.substring(idx, idx+2);
        long byteValue = strtol(byteString.c_str(), NULL, 16);
        cryptKey[idx / 2] = (byte)byteValue;
      }
    }
    else { // input key is ASCII-representation
      for(size_t i=0; i<32; i++) {
        cryptKey[i] = argKey[i];
      }
    }
  }
  // set the final key for encryption / decryption
  aes256.setKey(cryptKey, 32);

  #if PRINT_SERIAL
    Serial.println("Key set for encryption / decryption");
    for(size_t i=0; i<32; i++) {
        if(cryptKey[i] < 16) Serial.print("0");
        Serial.print(cryptKey[i], HEX);
        Serial.print(" ");
    }
    Serial.println();
  #endif

  #if USE_DISPLAY
    display_key(cryptKey);
  #endif
}

void decryptData (String data) {
  #if PRINT_SERIAL
    Serial.println();
    Serial.println("Operation: decrypting data");
  #endif

  data.replace(" ", "");
  byte dataBlock[aes256.blockSize()];
  byte plainText[aes256.blockSize()];
  String byteString = "";
  String fullString = "";

  // check whether the input data is a hexadecimal string
  bool dataIsHex = true;
  for(size_t i = 0; i < data.length(); i++)
    if ( !(data[i] >= '0' && data[i] <= '9') &&
          !(data[i] >= 'a' && data[i] <= 'f') &&
          !(data[i] >= 'A' && data[i] <= 'F') )
            dataIsHex = false;

  if(!dataIsHex) {
    #if PRINT_SERIAL
      Serial.println("Abort decryption! No valid input data!");
    #endif
    return;
  }

  #if PRINT_SERIAL
    Serial.println("decrypted data block: ");
  #endif

  for(int i=0; i < data.length(); i += aes256.blockSize() * 2) {
    for(int j=0; j<aes256.blockSize() * 2; j+=2) {
          byteString = data.substring(i + j, i + j + 2);
          long byteValue = strtol(byteString.c_str(), NULL, 16);
          dataBlock[j / 2] = (byte)byteValue;
    }
    aes256.decryptBlock(plainText, dataBlock);
    for(int i=0; i<sizeof(plainText); i++){
      if(plainText[i] == '\0') break;
      #if PRINT_SERIAL
        Serial.write(plainText[i]);
      #endif
      fullString += (char)plainText[i];
    }
  }
  #if PRINT_SERIAL
    Serial.println();
  #endif
  #if USE_DISPLAY
    display_decrypt_text(fullString, BTN_HOME);
  #endif
}

void encryptData (String data) {
  #if PRINT_SERIAL
    Serial.println();
    Serial.println("Operation: encrypting data");
  #endif

  byte dataBlock[aes256.blockSize()];
  byte cypher[aes256.blockSize()];
  String fullString = "";
  char hexStr[4];

  #if PRINT_SERIAL
    Serial.println("encrypted data block: ");
  #endif

  for(int i=0; i<data.length(); i+=aes256.blockSize()) {
    for(int j=0; j<aes256.blockSize(); j++) {
      dataBlock[j] = data[i + j];
      if((i + j) > data.length()) dataBlock[j] = '\0';
    }
    aes256.encryptBlock(cypher, dataBlock);
    for(int j=0; j<sizeof(cypher); j++) {
      byteToHex(cypher[j], hexStr);
      fullString += hexStr;
    }
  }
  #if PRINT_SERIAL
    Serial.println();
  #endif
  display_crypt_block(fullString, BTN_HOME);
}

void clearCryptograficInformation () {
  #if PRINT_SERIAL
    Serial.println();
    Serial.println("Operation: clear all cryptographic information from memory");
  #endif

  aes256.clear();
  #if PRINT_SERIAL
    Serial.println("Memory cleared");
  #endif

  display_home();
}

void restartDevice () {
    #if PRINT_SERIAL
      Serial.println("\r\nInitiating reboot of this device!\r\n");
    #endif

    ESP.restart();
}

void readCmd () {
  if ( Serial.available() ) {
    String inp = Serial.readString();
    inp.trim();
    String cmd = "";
    String arg = "";

    if(inp.indexOf(' ') > 0) {
      cmd = inp.substring(0, inp.indexOf(' '));
      arg = inp.substring(inp.indexOf(' ') + 1);
      arg.trim();
    }
    else
      cmd = inp;

    if (cmd == "help") printHelp();
    else if(cmd == "key") setCryptKey(arg);
    else if (cmd == "encrypt") encryptData(arg);
    else if (cmd == "decrypt") decryptData(arg);
    else if (cmd == "clear") clearCryptograficInformation();
    else if (cmd == "restart") restartDevice();
    else {
      #if PRINT_SERIAL
        Serial.println("\r\nUnknown command! Try again!");
      #endif
    }
  }
}

void readHome() {
  if(!digitalRead(BTN_HOME)) { // Home-Button pressed (pulled high with 10K resistor on the board)
    #if USE_DISPLAY
      display_home();
    #endif

  #if PRINT_SERIAL
    Serial.println("Waiting for command (send \"help\" for available commands)");
  #endif
  delay(1000);
  }
}

void setup() {
  Serial.begin(115200);
  while(!Serial);
  pinMode(BTN_HOME, INPUT);
  display_setup();
  Serial.println("Waiting for command");
  #if PRINT_SERIAL
    Serial.println("send \"help\" for available commands");
  #endif
}

void loop() {
  readCmd();
  readHome(); // check for Home-button being pressed
}