#include "mqtt_aes.h"

#include <openssl/bio.h>
#include <openssl/evp.h>
#include <openssl/sha.h>

#include <stdlib.h>
#include <string.h>

static int normalize_key(const char *secret_key, unsigned char *key_out, int *key_len_out)
{
    size_t len = strlen(secret_key);
    if (len == 16 || len == 24 || len == 32) {
        memcpy(key_out, secret_key, len);
        *key_len_out = (int)len;
        return 0;
    }

    if (SHA256((const unsigned char *)secret_key, len, key_out) == NULL) {
        return -1;
    }
    *key_len_out = 32;
    return 0;
}

static const EVP_CIPHER *cipher_for_key_len(int key_len)
{
    switch (key_len) {
    case 16:
        return EVP_aes_128_ecb();
    case 24:
        return EVP_aes_192_ecb();
    case 32:
        return EVP_aes_256_ecb();
    default:
        return NULL;
    }
}

static unsigned char *base64_decode(const char *input, size_t *out_len)
{
    BIO *bio = NULL;
    BIO *b64 = NULL;
    size_t len = strlen(input);
    unsigned char *buffer = (unsigned char *)malloc(len * 3 / 4 + 4);

    if (!buffer) {
        return NULL;
    }

    b64 = BIO_new(BIO_f_base64());
    bio = BIO_new_mem_buf(input, (int)len);
    bio = BIO_push(b64, bio);
    BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);

    int decoded = BIO_read(bio, buffer, (int)(len * 3 / 4 + 4));
    BIO_free_all(bio);

    if (decoded <= 0) {
        free(buffer);
        return NULL;
    }

    *out_len = (size_t)decoded;
    return buffer;
}

int mqtt_aes_decrypt_body(const char *b64_cipher, const char *secret_key, char **plain_json_out)
{
    unsigned char key_bytes[32];
    int key_len = 0;
    size_t cipher_len = 0;
    unsigned char *cipher_bytes = NULL;
    EVP_CIPHER_CTX *ctx = NULL;
    unsigned char *plain = NULL;
    int out_len1 = 0;
    int out_len2 = 0;
    int total = 0;

    if (!b64_cipher || !secret_key || !plain_json_out) {
        return -1;
    }
    *plain_json_out = NULL;

    if (normalize_key(secret_key, key_bytes, &key_len) != 0) {
        return -1;
    }

    const EVP_CIPHER *cipher = cipher_for_key_len(key_len);
    if (!cipher) {
        return -1;
    }

    cipher_bytes = base64_decode(b64_cipher, &cipher_len);
    if (!cipher_bytes || cipher_len == 0) {
        free(cipher_bytes);
        return -1;
    }

    ctx = EVP_CIPHER_CTX_new();
    if (!ctx) {
        free(cipher_bytes);
        return -1;
    }

    plain = (unsigned char *)malloc(cipher_len + 16);
    if (!plain) {
        EVP_CIPHER_CTX_free(ctx);
        free(cipher_bytes);
        return -1;
    }

    if (EVP_DecryptInit_ex(ctx, cipher, NULL, key_bytes, NULL) != 1) {
        goto fail;
    }

    if (EVP_DecryptUpdate(ctx, plain, &out_len1, cipher_bytes, (int)cipher_len) != 1) {
        goto fail;
    }

    if (EVP_DecryptFinal_ex(ctx, plain + out_len1, &out_len2) != 1) {
        goto fail;
    }

    total = out_len1 + out_len2;
    plain[total] = '\0';
    *plain_json_out = (char *)plain;
    plain = NULL;

    EVP_CIPHER_CTX_free(ctx);
    free(cipher_bytes);
    return 0;

fail:
    free(plain);
    EVP_CIPHER_CTX_free(ctx);
    free(cipher_bytes);
    return -1;
}
