#include <stdint.h>
#include <string.h>

// ======================================================
// Mini AES-like 16-bit block cipher
// ======================================================

static const uint8_t SAES_SBOX[16] = {
  0x9, 0x4, 0xA, 0xB,
  0xD, 0x1, 0x8, 0x5,
  0x6, 0x2, 0x0, 0x3,
  0xC, 0xE, 0xF, 0x7
};

static const uint8_t SAES_RSBOX[16] = {
  0xA, 0x5, 0x9, 0xB,
  0x1, 0x7, 0x8, 0xF,
  0x6, 0x0, 0x2, 0x3,
  0xC, 0x4, 0xD, 0xE
};

static const uint8_t SAES_RCON[2] = { 0x1, 0x2 };

static uint8_t gf16mul(uint8_t a, uint8_t b) {
  uint8_t p = 0;

  for (uint8_t i = 0; i < 4; i++) {
    if (b & 1) p ^= a;

    uint8_t carry = a & 0x8;
    a = (a << 1) & 0xF;

    if (carry) a ^= 0x3;

    b >>= 1;
  }

  return p;
}

static uint8_t getNib(uint16_t block, uint8_t pos) {
  return (block >> (12 - pos * 4)) & 0xF;
}

static uint16_t setNib(uint16_t block, uint8_t pos, uint8_t val) {
  uint8_t shift = 12 - pos * 4;

  block &= ~((uint16_t)0xF << shift);
  block |= ((uint16_t)(val & 0xF) << shift);

  return block;
}

// ======================================================
// MiniAES16 class
// ======================================================

class MiniAES16 {
public:

  void setKey(uint16_t key) {
    uint8_t w0 = (key >> 12) & 0xF;
    uint8_t w1 = (key >>  8) & 0xF;
    uint8_t w2 = (key >>  4) & 0xF;
    uint8_t w3 =  key        & 0xF;

    uint8_t w4 = w0 ^ SAES_SBOX[w3] ^ SAES_RCON[0];
    uint8_t w5 = w1 ^ SAES_SBOX[w4];
    uint8_t w6 = w2 ^ SAES_SBOX[w5];
    uint8_t w7 = w3 ^ SAES_SBOX[w6];

    roundKey[0] = pack(w0, w1);
    roundKey[1] = pack(w2, w3);
    roundKey[2] = pack(w4, w5);
    roundKey[3] = pack(w6, w7);
  }

  uint16_t encryptBlock(uint16_t b) {
    b ^= roundKey[0];
    b = subNib(b);
    b = swapRow(b);
    b = mixCols(b);

    b ^= roundKey[1];
    b = subNib(b);
    b = swapRow(b);
    b = mixCols(b);

    b ^= roundKey[2];
    b = subNib(b);
    b = swapRow(b);

    b ^= roundKey[3];

    return b;
  }

  uint16_t decryptBlock(uint16_t b) {
    b ^= roundKey[3];
    b = swapRow(b);
    b = invSubNib(b);

    b ^= roundKey[2];
    b = mixCols(b);
    b = swapRow(b);
    b = invSubNib(b);

    b ^= roundKey[1];
    b = mixCols(b);
    b = swapRow(b);
    b = invSubNib(b);

    b ^= roundKey[0];

    return b;
  }

  // ====================================================
  // CTR-style encryption/decryption
  //
  // Aynı fonksiyon hem encrypt hem decrypt yapar.
  // Nonce payload'un tamamını etkiler.
  // ====================================================

  void cryptCTR(const uint8_t* in, uint8_t* out, uint8_t len, uint16_t nonce) {
    uint16_t counter = 0;
    uint8_t i = 0;

    while (i < len) {
      uint16_t inputBlock = nonce ^ counter;
      uint16_t streamBlock = encryptBlock(inputBlock);

      uint8_t ks_hi = (uint8_t)(streamBlock >> 8);
      uint8_t ks_lo = (uint8_t)(streamBlock & 0xFF);

      out[i] = in[i] ^ ks_hi;

      if (i + 1 < len) {
        out[i + 1] = in[i + 1] ^ ks_lo;
      }

      i += 2;
      counter++;
    }
  }

private:

  uint16_t roundKey[4];

  uint16_t pack(uint8_t hi, uint8_t lo) {
    return ((uint16_t)hi << 12) |
           ((uint16_t)lo <<  8) |
           ((uint16_t)hi <<  4) |
           ((uint16_t)lo);
  }

  uint16_t subNib(uint16_t b) {
    for (uint8_t i = 0; i < 4; i++) {
      b = setNib(b, i, SAES_SBOX[getNib(b, i)]);
    }

    return b;
  }

  uint16_t invSubNib(uint16_t b) {
    for (uint8_t i = 0; i < 4; i++) {
      b = setNib(b, i, SAES_RSBOX[getNib(b, i)]);
    }

    return b;
  }

  uint16_t swapRow(uint16_t b) {
    uint8_t n1 = getNib(b, 1);
    uint8_t n3 = getNib(b, 3);

    b = setNib(b, 1, n3);
    b = setNib(b, 3, n1);

    return b;
  }

  uint16_t mixCols(uint16_t b) {
    for (uint8_t c = 0; c < 2; c++) {
      uint8_t a0 = getNib(b, c * 2);
      uint8_t a1 = getNib(b, c * 2 + 1);

      uint8_t r0 = gf16mul(2, a0) ^ gf16mul(3, a1);
      uint8_t r1 = gf16mul(3, a0) ^ gf16mul(2, a1);

      b = setNib(b, c * 2,     r0);
      b = setNib(b, c * 2 + 1, r1);
    }

    return b;
  }
};

// ======================================================
// Protocol test area
// ======================================================

MiniAES16 aes;

#define DATA_LEN     15
#define NONCE_LEN    2
#define PACKET_LEN   (NONCE_LEN + DATA_LEN)

static const uint8_t DATA[DATA_LEN] = {
  0x48, 0x65, 0x6C, 0x6C, 0x6F,
  0x20, 0x57, 0x6F, 0x72, 0x6C,
  0x64, 0x21, 0x20, 0x41, 0x42
};

// Test için basit nonce üretimi.
// Gerçek sistemde nonce/counter tekrar etmemeli.
uint16_t generateNonce() {
  uint16_t a = analogRead(A0);
  uint16_t t = millis() & 0xFFFF;

  return (a << 6) ^ t ^ (t >> 3);
}

void printHex(const uint8_t* buf, uint8_t len) {
  for (uint8_t i = 0; i < len; i++) {
    if (buf[i] < 0x10) Serial.print('0');
    Serial.print(buf[i], HEX);
    Serial.print(' ');
  }

  Serial.println();
}

void printAscii(const uint8_t* buf, uint8_t len) {
  for (uint8_t i = 0; i < len; i++) {
    char c = (char)buf[i];

    if (c >= 32 && c <= 126) {
      Serial.print(c);
    } else {
      Serial.print('.');
    }
  }

  Serial.println();
}

// ======================================================
// TX side
// Paket oluşturur:
//
// txPacket[0] = nonce high byte
// txPacket[1] = nonce low byte
// txPacket[2...] = encrypted payload
// ======================================================

void buildTxPacket(const uint8_t* plainPayload,
                   uint8_t payloadLen,
                   uint16_t nonce,
                   uint8_t* txPacket) {
  uint8_t encryptedPayload[DATA_LEN];

  aes.cryptCTR(plainPayload, encryptedPayload, payloadLen, nonce);

  txPacket[0] = (uint8_t)(nonce >> 8);
  txPacket[1] = (uint8_t)(nonce & 0xFF);

  memcpy(&txPacket[2], encryptedPayload, payloadLen);
}

// ======================================================
// RX side
// Gelen paketten nonce'u alır.
// Şifreli payload'u çözer.
// ======================================================

void parseRxPacket(const uint8_t* rxPacket,
                   uint8_t payloadLen,
                   uint8_t* decryptedPayload) {
  uint16_t rxNonce = ((uint16_t)rxPacket[0] << 8) | rxPacket[1];

  const uint8_t* encryptedPayload = &rxPacket[2];

  aes.cryptCTR(encryptedPayload, decryptedPayload, payloadLen, rxNonce);
}

// ======================================================
// Setup test
// ======================================================

void setup() {
  Serial.begin(9600);
  while (!Serial);

  uint16_t key = 0xA7B3;
  aes.setKey(key);

  Serial.println();
  Serial.println("MiniAES16 CTR Packet Test");
  Serial.println("================================");

  Serial.print("Key: 0x");
  Serial.println(key, HEX);

  Serial.print("Original DATA HEX:   ");
  printHex(DATA, DATA_LEN);

  Serial.print("Original DATA ASCII: ");
  printAscii(DATA, DATA_LEN);

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

  uint16_t passed = 0;

  for (uint16_t trial = 0; trial < 100; trial++) {
    uint16_t nonce = generateNonce();

    uint8_t txPacket[PACKET_LEN];
    uint8_t decrypted[DATA_LEN];

    memset(txPacket, 0x00, PACKET_LEN);
    memset(decrypted, 0x00, DATA_LEN);

    // TX: nonce + encrypted payload packet oluştur
    buildTxPacket(DATA, DATA_LEN, nonce, txPacket);

    // RX: paketi çöz
    parseRxPacket(txPacket, DATA_LEN, decrypted);

    bool ok = (memcmp(DATA, decrypted, DATA_LEN) == 0);

    if (ok) passed++;

    Serial.print("Trial ");
    Serial.print(trial + 1);
    Serial.print(" | Nonce: 0x");

    if (nonce < 0x1000) Serial.print('0');
    if (nonce < 0x0100) Serial.print('0');
    if (nonce < 0x0010) Serial.print('0');

    Serial.println(nonce, HEX);

    Serial.print("  Plain HEX:       ");
    printHex(DATA, DATA_LEN);

    Serial.print("  Encrypted HEX:   ");
    printHex(&txPacket[2], DATA_LEN);

    Serial.print("  TX Packet HEX:   ");
    printHex(txPacket, PACKET_LEN);

    Serial.print("  Dec HEX:         ");
    printHex(decrypted, DATA_LEN);

    Serial.print("  Dec ASCII:       ");
    printAscii(decrypted, DATA_LEN);

    Serial.println(ok ? "  -> OK" : "  -> ERROR");
    Serial.println();

    delay(10);
  }

  Serial.println("================================");
  Serial.print("Result: ");
  Serial.print(passed);
  Serial.println("/100 passed");
}

void loop() {
}