#include <DHT.h> 
#include <Arduino.h>
#include <SPI.h>
#include <esp_system.h> // Include the ESP32 system library for memory functions
#include <esp_heap_caps.h> // Include for detailed memory statistics

#define DHTPIN 4
#define DHTTYPE DHT22
DHT dht(DHTPIN, DHTTYPE);

#define TRIGPIN 12
#define ECHOPIN 14

const float a = 0.9;
const float pi = 3.14;

byte key3[]= "123456789abcdeff";
byte Sbox[16][16]={
   0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5,     // 0x00
    0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0,     // 0x10
    0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
    0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC,     // 0x20
    0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
    0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A,     // 0x30
    0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
    0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0,     // 0x40
    0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
    0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B,     // 0x50
    0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
    0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85,     // 0x60
    0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
    0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5,     // 0x70
    0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
    0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17,     // 0x80
    0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
    0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88,     // 0x90
    0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
    0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C,     // 0xA0
    0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
    0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9,     // 0xB0
    0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
    0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6,     // 0xC0
    0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
    0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E,     // 0xD0
    0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
    0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94,     // 0xE0
    0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
    0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68,     // 0xF0
    0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16

};
byte InvSbox[16][16]={
  
  0x52, 0x09, 0x6A, 0xD5, 0x30, 0x36, 0xA5, 0x38,     // 0x00
    0xBF, 0x40, 0xA3, 0x9E, 0x81, 0xF3, 0xD7, 0xFB,
    0x7C, 0xE3, 0x39, 0x82, 0x9B, 0x2F, 0xFF, 0x87,     // 0x10
    0x34, 0x8E, 0x43, 0x44, 0xC4, 0xDE, 0xE9, 0xCB,
    0x54, 0x7B, 0x94, 0x32, 0xA6, 0xC2, 0x23, 0x3D,     // 0x20
    0xEE, 0x4C, 0x95, 0x0B, 0x42, 0xFA, 0xC3, 0x4E,
    0x08, 0x2E, 0xA1, 0x66, 0x28, 0xD9, 0x24, 0xB2,     // 0x30
    0x76, 0x5B, 0xA2, 0x49, 0x6D, 0x8B, 0xD1, 0x25,
    0x72, 0xF8, 0xF6, 0x64, 0x86, 0x68, 0x98, 0x16,     // 0x40
    0xD4, 0xA4, 0x5C, 0xCC, 0x5D, 0x65, 0xB6, 0x92,
    0x6C, 0x70, 0x48, 0x50, 0xFD, 0xED, 0xB9, 0xDA,     // 0x50
    0x5E, 0x15, 0x46, 0x57, 0xA7, 0x8D, 0x9D, 0x84,
    0x90, 0xD8, 0xAB, 0x00, 0x8C, 0xBC, 0xD3, 0x0A,     // 0x60
    0xF7, 0xE4, 0x58, 0x05, 0xB8, 0xB3, 0x45, 0x06,
    0xD0, 0x2C, 0x1E, 0x8F, 0xCA, 0x3F, 0x0F, 0x02,     // 0x70
    0xC1, 0xAF, 0xBD, 0x03, 0x01, 0x13, 0x8A, 0x6B,
    0x3A, 0x91, 0x11, 0x41, 0x4F, 0x67, 0xDC, 0xEA,     // 0x80
    0x97, 0xF2, 0xCF, 0xCE, 0xF0, 0xB4, 0xE6, 0x73,
    0x96, 0xAC, 0x74, 0x22, 0xE7, 0xAD, 0x35, 0x85,     // 0x90
    0xE2, 0xF9, 0x37, 0xE8, 0x1C, 0x75, 0xDF, 0x6E,
    0x47, 0xF1, 0x1A, 0x71, 0x1D, 0x29, 0xC5, 0x89,     // 0xA0
    0x6F, 0xB7, 0x62, 0x0E, 0xAA, 0x18, 0xBE, 0x1B,
    0xFC, 0x56, 0x3E, 0x4B, 0xC6, 0xD2, 0x79, 0x20,     // 0xB0
    0x9A, 0xDB, 0xC0, 0xFE, 0x78, 0xCD, 0x5A, 0xF4,
    0x1F, 0xDD, 0xA8, 0x33, 0x88, 0x07, 0xC7, 0x31,     // 0xC0
    0xB1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xEC, 0x5F,
    0x60, 0x51, 0x7F, 0xA9, 0x19, 0xB5, 0x4A, 0x0D,     // 0xD0
    0x2D, 0xE5, 0x7A, 0x9F, 0x93, 0xC9, 0x9C, 0xEF,
    0xA0, 0xE0, 0x3B, 0x4D, 0xAE, 0x2A, 0xF5, 0xB0,     // 0xE0
    0xC8, 0xEB, 0xBB, 0x3C, 0x83, 0x53, 0x99, 0x61,
    0x17, 0x2B, 0x04, 0x7E, 0xBA, 0x77, 0xD6, 0x26,     // 0xF0
    0xE1, 0x69, 0x14, 0x63, 0x55, 0x21, 0x0C, 0x7D

};
uint8_t rcon[10] = {0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1B, 0x36};


void setup() {
Serial.begin(115200);
  dht.begin();
  pinMode(TRIGPIN, OUTPUT);
  pinMode(ECHOPIN, INPUT);

 

  Serial.println("Setup Complete.");
}


void expandKey(uint8_t *initialKey, uint8_t *expandedKeys, int rounds) {
  memcpy(expandedKeys, initialKey, 16);

  for (int round = 1; round <= rounds; round++) {
    uint8_t temp[4];

    // Rotate and substitute the last column of the previous key
    temp[0] = Sbox[expandedKeys[(round - 1) * 16 + 13] / 16][expandedKeys[(round - 1) * 16 + 13] % 16];
    temp[1] = Sbox[expandedKeys[(round - 1) * 16 + 14] / 16][expandedKeys[(round - 1) * 16 + 14] % 16];
    temp[2] = Sbox[expandedKeys[(round - 1) * 16 + 15] / 16][expandedKeys[(round - 1) * 16 + 15] % 16];
    temp[3] = Sbox[expandedKeys[(round - 1) * 16 + 12] / 16][expandedKeys[(round - 1) * 16 + 12] % 16]; // Rotate

    temp[0] ^= rcon[round - 1]; // Add round constant

    // Generate the next round key
    for (int i = 0; i < 16; i++) {
      if (i < 4) {
        expandedKeys[round * 16 + i] = expandedKeys[(round - 1) * 16 + i] ^ temp[i];
      } else {
        expandedKeys[round * 16 + i] = expandedKeys[(round - 1) * 16 + i] ^ expandedKeys[round * 16 + i - 4];
      }
    }
  }
}

byte gfMul(byte a, byte b) {
  byte p = 0;
  byte hiBitSet;

  for (byte i = 0; i < 8; i++) {
    if (b & 1) {
      p ^= a;
    }
    hiBitSet = a & 0x80;
    a <<= 1;
    if (hiBitSet) {
      a ^= 0x1b;  // AES modulo polynomial x^8 + x^4 + x^3 + x + 1
    }
    b >>= 1;
  }

  return p;
}


 


void applyPadding(byte *paddedInput, const byte *input, int inputLen) {
  // Copy the input data into the padded block
  memcpy(paddedInput, input, inputLen);
  // Apply zero padding
  
  for (int i = inputLen; i < 16; i++) {
    paddedInput[i] = 0x00;
  }
}



void removePadding(byte *output, const byte *paddedInput) {
  int unpaddedLength = 16;
  // Find the unpadded length by detecting the first 0x00 from the end
  while (unpaddedLength > 0 && paddedInput[unpaddedLength - 1] == 0x00) {
    unpaddedLength--;
  }
  // Copy the data excluding the padding
  memcpy(output, paddedInput, unpaddedLength);
}

void encryptData(byte *output, const byte *input, int inputLen) {
  byte paddedInput[16];
  applyPadding(paddedInput, input, inputLen); // Apply padding

  byte state[16];
  memcpy(state, paddedInput, 16);

  byte expandedKeys[176]; // Expanded key array
  expandKey(key3, expandedKeys, 10); // Expand the keys for encryption

  // Initial Round: AddRoundKey
  for (int i = 0; i < 16; i++) {
    state[i] ^= key3[i]; // Use expanded key
  }

  // Main Rounds
  for (int round = 1; round < 10; round++) {
    // SubBytes
    for (int i = 0; i < 16; i++) {
      state[i] = Sbox[state[i] / 16][state[i] % 16];
    }

    // ShiftRows
    byte temp[16];
    temp[0] = state[0];
    temp[1] = state[5];
    temp[2] = state[10];
    temp[3] = state[15];
    temp[4] = state[4];
    temp[5] = state[9];
    temp[6] = state[14];
    temp[7] = state[3];
    temp[8] = state[8];
    temp[9] = state[13];
    temp[10] = state[2];
    temp[11] = state[7];
    temp[12] = state[12];
    temp[13] = state[1];
    temp[14] = state[6];
    temp[15] = state[11];
    memcpy(state, temp, 16);

    // MixColumns
    for (int col = 0; col < 4; col++) {
      temp[col * 4 + 0] = gfMul(state[col * 4 + 0], 2) ^ gfMul(state[col * 4 + 1], 3) ^ state[col * 4 + 2] ^ state[col * 4 + 3];
      temp[col * 4 + 1] = state[col * 4 + 0] ^ gfMul(state[col * 4 + 1], 2) ^ gfMul(state[col * 4 + 2], 3) ^ state[col * 4 + 3];
      temp[col * 4 + 2] = state[col * 4 + 0] ^ state[col * 4 + 1] ^ gfMul(state[col * 4 + 2], 2) ^ gfMul(state[col * 4 + 3], 3);
      temp[col * 4 + 3] = gfMul(state[col * 4 + 0], 3) ^ state[col * 4 + 1] ^ state[col * 4 + 2] ^ gfMul(state[col * 4 + 3], 2);
    }
    memcpy(state, temp, 16);

    // AddRoundKey
    for (int i = 0; i < 16; i++) {
      state[i] ^= expandedKeys[round * 16 + i];
    }
  }

  // Final Round
  for (int i = 0; i < 16; i++) {
    state[i] = Sbox[state[i] / 16][state[i] % 16];
  }

  byte temp[16];
  temp[0] = state[0];
  temp[1] = state[5];
  temp[2] = state[10];
  temp[3] = state[15];
  temp[4] = state[4];
  temp[5] = state[9];
  temp[6] = state[14];
  temp[7] = state[3];
  temp[8] = state[8];
  temp[9] = state[13];
  temp[10] = state[2];
  temp[11] = state[7];
  temp[12] = state[12];
  temp[13] = state[1];
  temp[14] = state[6];
  temp[15] = state[11];
  memcpy(state, temp, 16);
   // AddRoundKey
  for (int i = 0; i < 16; i++) {
    state[i] ^= expandedKeys[10 * 16 + i]; // Final expanded key
  }

  // Output the encrypted state
  memcpy(output, state, 16);

}   
/*void encryptECB(byte *output, const byte *input, int length) {
  int numBlocks = (length + 15) / 16; // Calculate the number of 16-byte blocks
  byte buffer[16];

  for (int i = 0; i < numBlocks; i++) {
    int blockSize = (i == numBlocks - 1 && length % 16 != 0) ? length % 16 : 16;
    
    // Apply padding for the last block (only when necessary)
   
    if (i == numBlocks - 1 && length % 16 != 0) {
      applyPadding(buffer, &input[i * 16], blockSize);
      encryptData(&output[i * 16], buffer); // Encrypt the padded block
    } else {
      encryptData(&output[i * 16], &input[i * 16]); // Encrypt the full 16-byte block
    }    
  }
}
*/
/*void encryptECB(byte *output, const byte *input, int length) {
  byte buffer[16] = {0}; 
    // Copy input data to padded block
  memcpy(buffer, input, length); 

  // Encrypt the padded block
  encryptData(output, buffer); 
    
  }*/

void decryptData(byte *output, const byte *input) {
  byte state[16];
  memcpy(state, input, 16);

  byte expandedKeys[176]; // Expanded key array
  expandKey(key3, expandedKeys, 10); // Expand the keys for decryption

  // Initial Round: AddRoundKey
  for (int i = 0; i < 16; i++) {
    state[i] ^= expandedKeys[10 * 16 + i]; // Use the last expanded key
  }

  // Final Round (without MixColumns)
  // Inverse ShiftRows
  byte temp[16];
  temp[0] = state[0];
  temp[1] = state[13];
  temp[2] = state[10];
  temp[3] = state[7];
  temp[4] = state[4];
  temp[5] = state[1];
  temp[6] = state[14];
  temp[7] = state[11];
  temp[8] = state[8];
  temp[9] = state[5];
  temp[10] = state[2];
  temp[11] = state[15];
  temp[12] = state[12];
  temp[13] = state[9];
  temp[14] = state[6];
  temp[15] = state[3];
  memcpy(state, temp, 16);

  // Inverse SubBytes
  for (int i = 0; i < 16; i++) {
    state[i] = InvSbox[state[i] / 16][state[i] % 16];
  }

  // Main Rounds
  for (int round = 9; round > 0; round--) {
    // AddRoundKey
    for (int i = 0; i < 16; i++) {
      state[i] ^= expandedKeys[round * 16 + i];
    }

    // Inverse MixColumns
    for (int col = 0; col < 4; col++) {
      temp[col * 4 + 0] = gfMul(state[col * 4 + 0], 14) ^ gfMul(state[col * 4 + 1], 11) ^ gfMul(state[col * 4 + 2], 13) ^ gfMul(state[col * 4 + 3], 9);
      temp[col * 4 + 1] = gfMul(state[col * 4 + 0], 9) ^ gfMul(state[col * 4 + 1], 14) ^ gfMul(state[col * 4 + 2], 11) ^ gfMul(state[col * 4 + 3], 13);
      temp[col * 4 + 2] = gfMul(state[col * 4 + 0], 13) ^ gfMul(state[col * 4 + 1], 9) ^ gfMul(state[col * 4 + 2], 14) ^ gfMul(state[col * 4 + 3], 11);
      temp[col * 4 + 3] = gfMul(state[col * 4 + 0], 11) ^ gfMul(state[col * 4 + 1], 13) ^ gfMul(state[col * 4 + 2], 9) ^ gfMul(state[col * 4 + 3], 14);
    }
    memcpy(state, temp, 16);

    // Inverse ShiftRows
    temp[0] = state[0];
    temp[1] = state[13];
    temp[2] = state[10];
    temp[3] = state[7];
    temp[4] = state[4];
    temp[5] = state[1];
    temp[6] = state[14];
    temp[7] = state[11];
    temp[8] = state[8];
    temp[9] = state[5];
    temp[10] = state[2];
    temp[11] = state[15];
    temp[12] = state[12];
    temp[13] = state[9];
    temp[14] = state[6];
    temp[15] = state[3];
    memcpy(state, temp, 16);

    // Inverse SubBytes
    for (int i = 0; i < 16; i++) {
      state[i] = InvSbox[state[i] / 16][state[i] % 16];
    }
  }

  // Initial Round: AddRoundKey
  for (int i = 0; i < 16; i++) {
    state[i] ^= key3[i];
  }

  // Remove padding from the decrypted state
  removePadding(output, state);
}

// New function for SAC
void testSAC(byte *input) {
  byte flippedInput[16];
  byte encryptedOriginal[16], encryptedFlipped[16];
  byte diffCount = 0;

  // Flip each bit one at a time
  for (int i = 0; i < 16; i++) {
    for (int bit = 0; bit < 8; bit++) {
      memcpy(flippedInput, input, 16);
      flippedInput[i] ^= (1 << bit); // Flip one bit

      // Encrypt the original and flipped inputs
      encryptData(encryptedOriginal, input,16);  // Encrypt the original data
      encryptData(encryptedFlipped, flippedInput,16);  // Encrypt the flipped data

      // Count the bit differences between the encrypted original and flipped
      for (int j = 0; j < 16; j++) {
        byte diff = encryptedOriginal[j] ^ encryptedFlipped[j];
        diffCount += __builtin_popcount(diff);
      }
    }
  }

  float avalanchePercentage = (diffCount / 128.0) * 100.0; // 128 bits in total (16 bytes * 8 bits per byte)
  
  // Define the pass threshold (e.g., 50%)
  float passThreshold = 50.0; // 50% bit difference threshold for passing

  // Print result based on the avalanche percentage
  Serial.print("SAC Avalanche Percentage: ");
  Serial.println(avalanchePercentage);

  if (avalanchePercentage >= passThreshold) {
    Serial.println("Test Passed");
  } else {
    Serial.println("Test Failed");
  }
}


float readDistance() {
  digitalWrite(TRIGPIN, LOW);
  delayMicroseconds(2);
  digitalWrite(TRIGPIN, HIGH);
  delayMicroseconds(10);
  digitalWrite(TRIGPIN, LOW);

  float duration = pulseIn(ECHOPIN, HIGH);
  return duration * 0.034 / 2; // Convert to cm
}

void loop() {
  float temperature = dht.readTemperature();
  float humidity = dht.readHumidity();
  float keyexpansiontime = 0 ;
  if (isnan(temperature) || isnan(humidity)) {
    Serial.println("Failed to read from DHT sensor!");
    return;
  }

  float distance = readDistance();
  
   
  String dataTypes[] = {"Temperature", "Humidity", "Distance"};
  String values[] = {
    String(temperature, 2),
    String(humidity, 2),
    String(distance, 2)
  };
  byte expandedKeys[176]; // Expanded key array
  for (int i = 0; i < 3; i++) {
    Serial.println("\nProcessing: " + dataTypes[i]);

    byte input[16] = {0};
    byte encryptedData[16] = {0};
    byte decryptedData[16] = {0};

    String data = values[i];
    for (int j = 0; j < data.length() && j < 16; j++) {
      input[j] = data[j];
    }
    unsigned long startkeyexpansion = micros();
    expandKey(key3, expandedKeys, 10); // Expand the keys for encryption
    unsigned long endkeyexpansion = micros();
    
    keyexpansiontime = endkeyexpansion - startkeyexpansion ; 


    // Using encryptECB instead of encryptData
    size_t memoryBeforeEncryption = ESP.getFreeHeap();
    unsigned long startTime2 = micros();
    encryptData(encryptedData, input, data.length());  // Using encryptECB
         unsigned long endencryption = micros();    
   unsigned long encryptionTime = endencryption - startTime2;
    size_t memoryAfterEncryption = ESP.getFreeHeap();
    size_t memoryConsumedEncryption = memoryBeforeEncryption - memoryAfterEncryption;

    // Using decryptECB instead of decryptData
    size_t memoryBeforeDecryption = ESP.getFreeHeap();
    unsigned long startTime1 = micros();
    decryptData(decryptedData, encryptedData);  // Using decryptECB
    unsigned long enddecryption = micros();
    unsigned long decryptionTime =enddecryption - startTime1;
    size_t memoryAfterDecryption = ESP.getFreeHeap();
    size_t memoryConsumedDecryption = memoryBeforeDecryption - memoryAfterDecryption;

   float encryptionThroughput = (16.0 / encryptionTime) * 1e6;
    float decryptionThroughput = (16.0 / decryptionTime) * 1e6;

    Serial.print("Original Data: ");
    Serial.println(data);

    Serial.print("Encrypted Data: ");
    for (int j = 0; j < 16; j++) {
      Serial.printf("%02x", encryptedData[j]);
    }
    Serial.println();

    Serial.print("Decrypted Data: ");
    Serial.println((char *)decryptedData);
    Serial.print("Key expansion Time (us): ");
    Serial.println(keyexpansiontime); 
    
    Serial.print("Encryption Time (us): ");
    Serial.println(encryptionTime);
    Serial.print("Decryption Time (us): ");
    Serial.println(decryptionTime);
    Serial.print("Encryption Throughput (bytes/sec): ");
    Serial.println(encryptionThroughput);
    Serial.print("Decryption Throughput (bytes/sec): ");
    Serial.println(decryptionThroughput);

    //Serial.print("Memory Consumed during Encryption (bytes): ");
    //Serial.println(memoryConsumedEncryption);
    //Serial.print("Memory Consumed during Decryption (bytes): ");
    //Serial.println(memoryConsumedDecryption);

    // Test SAC on encrypted data
    testSAC(input);
  }
  delay(5000);
}