/*
 * Copyright (C) 2015 Southern Storm Software, Pty Ltd.
 *
 * 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.
 */
/*
This example explains basic AES128 implementation.
In Arduino serial monitor text appears to be non readable characters, but if you use any other serial terminals you can see the hex values
Example contributor: Aswin
*/

#include <Crypto.h>
#include <AES.h>
#include <string.h>
#include <EEPROM.h>
#include <Ed25519.h>
#include <Curve25519.h>
#include <RNG.h>

const int addr = 0;  //EEPROM Adresse
boolean stats = false;
unsigned long int time;

//byte AESkey[16] = { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F };
AES128 aes128;

uint8_t privateKey[32];
uint8_t publicKey[32];
uint8_t sharedSecret[32];
uint8_t signature[64];

void generateKeyPair(uint8_t* privateKey, uint8_t* publicKey) {
  RNG.begin(analogRead(A0));
  Curve25519::dh1(privateKey, publicKey);
}
void computeSharedSecret(const uint8_t* privateKey, const uint8_t* otherPublicKey, uint8_t* sharedSecret) {
  memcpy(sharedSecret, privateKey, 32);
  Curve25519::dh2(sharedSecret, (uint8_t*)otherPublicKey);
}
void pkcs7pad(byte* buffer, size_t dataLen, size_t blockSize) {
  byte padValue = blockSize - (dataLen % blockSize);
  for (size_t i = dataLen; i < dataLen + padValue; ++i) {
    buffer[i] = padValue;
  }
}
size_t pkcs7unpad(byte* buffer, size_t bufferLen) {
  byte padValue = buffer[bufferLen - 1];
  if (padValue > 0 && padValue <= 16) {
    for (size_t i = bufferLen - padValue; i < bufferLen; ++i) {
      if (buffer[i] != padValue) return bufferLen;
    }
    return bufferLen - padValue;
  }
  return bufferLen;
}
uint8_t hexCharToNibble(char c) {
  if (c >= '0' && c <= '9') return c - '0';
  if (c >= 'A' && c <= 'F') return c - 'A' + 10;
  if (c >= 'a' && c <= 'f') return c - 'a' + 10;
  return 0;
}
void printHex(const byte* data, size_t len) {
  for (size_t i = 0; i < len; i++) {
    if (data[i] < 16) Serial.print("0");
    Serial.print(data[i], HEX);
    //Serial.print(" ");
  }
  Serial.println();
}
bool readHexKeyFromSerial(uint8_t* key, size_t len) {
  char buffer[65] = { 0 };
  size_t count = 0;
  while (count < 64) {
    if (Serial.available()) {
      char c = Serial.read();
      if ((c >= '0' && c <= '9') || (c >= 'A' && c <= 'F') || (c >= 'a' && c <= 'f')) {
        buffer[count++] = c;
        Serial.print(c);  // Echo für Feedback
      }
    }
  }
  Serial.println();
  // Umwandlung in Bytes
  for (size_t i = 0; i < len; ++i) {
    key[i] = (hexCharToNibble(buffer[2 * i]) << 4) | hexCharToNibble(buffer[2 * i + 1]);
  }
  return true;
}
String readSerialLine() {
  String input = "";
  while (true) {
    if (Serial.available()) {
      char c = Serial.read();
      if (c == '\n' || c == '\r') {  // Zeilenende
        if (input.length() > 0) break;
      } else {
        input += c;
      }
    }
  }
  return input;
}
void serialPrintModusWahl() {
  Serial.println("\nWähle den Modus aus:");  //---\n---\n---\n---\n---\n---\n---\n---\n---\n---
  Serial.println("1: Text verschlüsseln");
  Serial.println("2: Text entschlüsseln");
  //Serial.println("3: Debug-Modus");
  Serial.print("t: Ausführliche Statistiken und serielle Ausgaben");
  if (stats) Serial.println(" (aktuell aktiviert)");
  else Serial.println(" (aktuell deaktiviert)");
}
bool serialYes(char eingabe) {
  if (eingabe == 'Y' || eingabe == 'y' || eingabe == 'J' || eingabe == 'j' || eingabe == '1' || eingabe == '\n') return true;
  return false;
}
bool serialNo(char eingabe) {
  if (eingabe == 'N' || eingabe == 'n' || eingabe == '0') return true;
  return false;
}
void schluesselaustausch() {
  //todo
}
void generateAESkey() {
  //todo
}
void exchangeAESkey() {
  //todo
}
/*void signMessage(char* message, uint8_t* signature) {
  Ed25519::sign(signature, privateKey, publicKey_my, (uint8_t)message, strlen(message));
}
bool verifyMessage(char message, uint8_t* signature, uint8_t* senderPublicKey) {
  return Ed25519::verify(signature, senderPublicKey, (uint8_t*)message, strlen(message));
}*/
// Wandelt ein einzelnes Hex-Zeichen ('0'-'9', 'A'-'F', 'a'-'f') in den Zahlenwert um
byte hexCharToByte(char c) {
  if (c >= '0' && c <= '9') return c - '0';
  if (c >= 'A' && c <= 'F') return c - 'A' + 10;
  if (c >= 'a' && c <= 'f') return c - 'a' + 10;
  return 0;
}
// Wandelt einen Hex-String (z.B. "AF32...") in ein Byte-Array um
void hexStringToBytes(const String& hex, byte* out, size_t outLen) {
  for (size_t i = 0; i < outLen; i++) {
    out[i] = (hexCharToByte(hex[2 * i]) << 4) | hexCharToByte(hex[2 * i + 1]);
  }
}
/*void PrintSignature() {
  for (int i = 0; i < 64; i++) {
    if (signature[i] < 0x10) Serial.print('0');
    Serial.print(signature[i], HEX);
  }
}*/
void encrypt() {
  Serial.println("Gib den zu verschlüsselnden Text ein (maximal 254 Zeichen) und drücke ENTER …");
  String text = readSerialLine();
  if (stats) {
    time = millis();
  }
  if (stats) {
    Serial.print("Textlänge:\t");
    Serial.println(text.length());
  }
  if (text.length() > 254) {
    Serial.println("Text zu lang, maximal 254 Zeichen erlaubt.");
    return;
  }
  if (stats) {
    Serial.print("Eingabe:\t");
    Serial.println(text);
    Serial.print("Schlüssel (shared secret):\t");
    printHex(sharedSecret, 32);
  }

  size_t textLen = text.length();
  size_t blockSize = 16;
  size_t numBlocks = (textLen / blockSize) + ((textLen % blockSize) ? 1 : 0);
  size_t paddedLen = numBlocks * blockSize;

  // Klartext vorbereiten
  byte* plaintext = new byte[paddedLen];
  memset(plaintext, 0, paddedLen);
  text.getBytes(plaintext, paddedLen);
  pkcs7pad(plaintext, textLen, blockSize);

  if (stats) {
    Serial.print("Plaintext (HEX): ");
    printHex(plaintext, paddedLen);
  }

  aes128.setKey(sharedSecret, 32);

  // Verschlüsseln
  byte* cypher = new byte[paddedLen];
  for (size_t i = 0; i < paddedLen; i += blockSize) {
    aes128.encryptBlock(cypher + i, plaintext + i);
  }
  Serial.print("Verschlüsselt (HEX):\t");
  printHex(cypher, paddedLen);
  if (stats) {
    Serial.print("Verschlüsselt (Text):\t");
    for (int j = 0; j < paddedLen; j++) {
      Serial.write(cypher[j]);
    }
  }
  //Signatur hier implementieren
  /*signMessage(cypher, signature);
  PrintSignature();*/
  // Speicher aufräumen
  delete[] plaintext;
  delete[] cypher;
  Serial.println();
  if (stats) runtime(text.length());
  Serial.println();
}
void decrypt() {
  Serial.println("Gib den zu entschlüsselnden Text ein (maximal 254 Zeichen) und drücke ENTER …");
  String text = readSerialLine();
  text.trim();

  /*Serial.print("Textlänge:\t");
  Serial.println(text.length());
  if (text.length() > 254) {
    Serial.println("Text zu lang, maximal 254 Zeichen erlaubt.");
    return;
  }*/
  if (text.length() % 32 != 0 && text.length() % 16 != 0) {
    Serial.println("Warnung: Die Hex-Zeichenanzahl sollte ein Vielfaches von 32 (für 16 Byte Blöcke) sein!");
  }

  if (stats) {
    Serial.print("Eingabe:\t");
    Serial.println(text);
  }

  size_t byteLen = text.length() / 2;  //2 Hex-Zeichen pro Byte
  size_t blockSize = 16;
  size_t numBlocks = (byteLen / blockSize) + ((byteLen % blockSize) ? 1 : 0);
  size_t paddedLen = numBlocks * blockSize;

  // Byte-Array für den Ciphertext
  byte* cypher = new byte[paddedLen];
  memset(cypher, 0, paddedLen);

  // Hex-String in Byte-Array umwandeln
  hexStringToBytes(text, cypher, byteLen);

  if (stats) {
    Serial.print("encrypted (HEX): ");
    printHex(cypher, paddedLen);
  }

  aes128.setKey(sharedSecret, 32);


  // Entschlüsseln
  byte* decryptedtext = new byte[paddedLen];
  for (size_t i = 0; i < paddedLen; i += blockSize) {
    aes128.decryptBlock(decryptedtext + i, cypher + i);
  }
  Serial.print("Decrypted (HEX): ");
  printHex(decryptedtext, paddedLen);

  // Padding entfernen
  size_t unpaddedLen = pkcs7unpad(decryptedtext, paddedLen);
  Serial.print("Decrypted (Text): ");
  for (size_t i = 0; i < unpaddedLen; i++) {
    Serial.write(decryptedtext[i]);
  }
  Serial.println();
  if (stats) runtime(unpaddedLen);
  Serial.println();

  // Speicher aufräumen
  delete[] cypher;
  delete[] decryptedtext;
}
void runtime(unsigned int length) {
  Serial.print("Laufzeit:\t");
  Serial.print(millis() - time);
  Serial.println(" ms");
  Serial.print("Laufzeit pro Zeichen:\t");
  Serial.print((millis() - time) / length);
  Serial.println(" ms");
}
void setup() {
  Serial.begin(9600);
  generateKeyPair(privateKey, publicKey);
  Serial.println("Schlüssel wurden generiert\n");
  Serial.println("Eigenes Public Key (an Partner senden):");
  printHex(publicKey, 32);
  Serial.println("Bitte 64 Hex-Zeichen (ohne Leerzeichen) für den anderen Public Key eingeben und Enter drücken:");
  bool done = false;
  while (!done) {
    if (Serial.available()) {
      uint8_t otherPublicKey[32];
      if (readHexKeyFromSerial(otherPublicKey, 32)) {
        Serial.println("Other Public Key empfangen:");
        printHex(otherPublicKey, 32);
        Serial.println("eingelesen");
        computeSharedSecret(privateKey, otherPublicKey, sharedSecret);

        Serial.println("Gemeinsames Geheimnis:");
        printHex(sharedSecret, 32);

        done = true;
        //Serial.println("Fertig! Neustart für neuen Austausch.");
      }
    }
  }
  EEPROM.get(addr, stats);
  serialPrintModusWahl();
}
void loop() {
  // Warte auf Nutzereingabe über die serielle Konsole
  if (Serial.available()) {
    Serial.println();
    char serial = Serial.read();
    if (serial == '1') {
      Serial.println("1: Verschlüsseln");
      encrypt();
    } else if (serial == '2') {
      Serial.println("2: Entschlüsseln");
      decrypt();
    } /*else if (serial == '3') {
      Serial.println("3: Debugging");
      //todo
  }*/
    else if (serial == 't') {
      if (stats) {
        stats = false;
        EEPROM.put(addr, stats);
        Serial.println("Statistiken deaktiviert");
      } else {
        stats = true;
        EEPROM.put(addr, stats);
        Serial.println("Statistiken aktiviert");
      }
    } else {
      Serial.println("Unbekannter Modus, try again");
    }
    while (Serial.available()) {
      Serial.read();
    }
    delay(10);
    serialPrintModusWahl();
    delay(10);
  }
}