#include <Arduino.h>
#include <WiFi.h>
#include <mbedtls/config.h>
#include <mbedtls/entropy.h>
#include <mbedtls/ctr_drbg.h>
#include <mbedtls/rsa.h>

// Define the key size
#define KEY_SIZE_BITS   2048
#define EXPONENT        65537

// Function to encrypt a plaintext string using RSA public key
String encryptString(const char *plaintext, mbedtls_rsa_context *rsa) {
    mbedtls_ctr_drbg_context ctr_drbg;
    mbedtls_entropy_context entropy;
    mbedtls_ctr_drbg_init(&ctr_drbg);
    mbedtls_entropy_init(&entropy);

    const char *pers = "rsa_encrypt";
    mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                          (const unsigned char *)pers, strlen(pers));

    size_t plaintext_len = strlen(plaintext);
    size_t encrypted_len = 0;
    unsigned char encrypted[MBEDTLS_MPI_MAX_SIZE];

    // Encrypt using RSA public key
    int ret = mbedtls_rsa_pkcs1_encrypt(rsa, mbedtls_ctr_drbg_random, &ctr_drbg,
                                         MBEDTLS_RSA_PUBLIC, plaintext_len,
                                         (const unsigned char *)plaintext,
                                         encrypted);
    if (ret != 0) {
        Serial.printf("Failed to encrypt: -0x%x\n", -ret);
        mbedtls_ctr_drbg_free(&ctr_drbg);
        mbedtls_entropy_free(&entropy);
        return "";
    }

    mbedtls_ctr_drbg_free(&ctr_drbg);
    mbedtls_entropy_free(&entropy);

    return String((char *)encrypted);
}

// Function to decrypt an encrypted string using RSA private key
String decryptString(const String &encrypted_str, mbedtls_rsa_context *rsa) {
    mbedtls_ctr_drbg_context ctr_drbg;
    mbedtls_entropy_context entropy;
    mbedtls_ctr_drbg_init(&ctr_drbg);
    mbedtls_entropy_init(&entropy);

    const char *pers = "rsa_decrypt";
    mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                          (const unsigned char *)pers, strlen(pers));

    size_t encrypted_len = encrypted_str.length();
    unsigned char encrypted[MBEDTLS_MPI_MAX_SIZE];
    memcpy(encrypted, encrypted_str.c_str(), encrypted_len);

    // Decrypt using RSA private key
    size_t decrypted_len = 0;
    unsigned char decrypted[MBEDTLS_MPI_MAX_SIZE];
    int ret = mbedtls_rsa_pkcs1_decrypt(rsa, mbedtls_ctr_drbg_random, &ctr_drbg,
                                         MBEDTLS_RSA_PRIVATE, &decrypted_len,
                                         encrypted, decrypted, sizeof(decrypted));
    if (ret != 0) {
        Serial.printf("Failed to decrypt: -0x%x\n", -ret);
        mbedtls_ctr_drbg_free(&ctr_drbg);
        mbedtls_entropy_free(&entropy);
        return "";
    }

    mbedtls_ctr_drbg_free(&ctr_drbg);
    mbedtls_entropy_free(&entropy);

    decrypted[decrypted_len] = '\0';  // Null-terminate the decrypted data
    return String((char *)decrypted);
}

void setup() {
    Serial.begin(115200);
    while (!Serial);  // Wait for Serial to be ready

    mbedtls_rsa_context rsa;
    mbedtls_entropy_context entropy;
    mbedtls_ctr_drbg_context ctr_drbg;
    const char *pers = "rsa_keygen";

    mbedtls_rsa_init(&rsa, MBEDTLS_RSA_PKCS_V15, 0);

    mbedtls_entropy_init(&entropy);
    mbedtls_ctr_drbg_init(&ctr_drbg);

    // Seed the random number generator
    if (mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                               (const unsigned char *)pers, strlen(pers)) != 0) {
        Serial.println("Failed to seed random number generator");
        return;
    }

    // Generate RSA key pair
    int ret = mbedtls_rsa_gen_key(&rsa, mbedtls_ctr_drbg_random, &ctr_drbg,
                                   KEY_SIZE_BITS, EXPONENT);
    if (ret != 0) {
        Serial.printf("Failed to generate RSA key pair: -0x%x\n", -ret);
        mbedtls_rsa_free(&rsa);
        mbedtls_ctr_drbg_free(&ctr_drbg);
        mbedtls_entropy_free(&entropy);
        return;
    }

    // Print the public key components
    Serial.println("RSA key generation successful!");
    Serial.print("Public key (N): ");
    for (size_t i = 0; i < rsa.N.n * sizeof(mbedtls_mpi_uint); i++) {
        Serial.printf("%02x", ((unsigned char *)rsa.N.p)[i]);
    }
    Serial.println();
    Serial.printf("Public key (E): %d\n", rsa.E);

    // Encrypt and decrypt a test string
    const char *plaintext = "Hello, world!";
    Serial.printf("Original: %s\n", plaintext);

    String encrypted_str = encryptString(plaintext, &rsa);
    Serial.printf("Encrypted: %s\n", encrypted_str.c_str());

    String decrypted_str = decryptString(encrypted_str, &rsa);
    Serial.printf("Decrypted: %s\n", decrypted_str.c_str());

    // Free resources
    mbedtls_rsa_free(&rsa);
    mbedtls_ctr_drbg_free(&ctr_drbg);
    mbedtls_entropy_free(&entropy);

    // End of setup, loop not used
}

void loop() {
    // Empty loop, not used in this example
}