/*
 * Created by victoria on 2021-05-13.
 */

#include "aes_openssl.h"
#include <memory>

extern "C" {
#include <openssl/evp.h>
#include <openssl/opensslv.h>
}

using namespace Signal::Crypto;

class EVPCipherCtxWrapper {
public:
    EVPCipherCtxWrapper() {
#if OPENSSL_VERSION_NUMBER >= 0x1010000fL
        ctx = EVP_CIPHER_CTX_new();
#else
        ctx = new EVP_CIPHER_CTX;
        EVP_CIPHER_CTX_init(ctx);
#endif
    }

    ~EVPCipherCtxWrapper() {
        if (good()) {
#if OPENSSL_VERSION_NUMBER >= 0x1010000fL
            EVP_CIPHER_CTX_free(ctx);
#else
            EVP_CIPHER_CTX_cleanup(ctx);
            delete ctx;
#endif
        }
    }

    EVPCipherCtxWrapper(const EVPCipherCtxWrapper &) = delete;
    EVPCipherCtxWrapper(EVPCipherCtxWrapper &&) = delete;
    EVPCipherCtxWrapper &operator=(const EVPCipherCtxWrapper &) = delete;

    [[nodiscard]] bool good() const { return ctx != nullptr; }

    [[nodiscard]] EVP_CIPHER_CTX *operator*() const { return ctx; }

private:
    EVP_CIPHER_CTX *ctx{nullptr};
};

static const EVP_CIPHER *aes_cipher(int cipher, size_t key_len) {
    if (cipher == SG_CIPHER_AES_CBC_PKCS5) {
        if (key_len == 16) {
            return EVP_aes_128_cbc();
        } else if (key_len == 24) {
            return EVP_aes_192_cbc();
        } else if (key_len == 32) {
            return EVP_aes_256_cbc();
        }
    } else if (cipher == SG_CIPHER_AES_CTR_NOPADDING) {
        if (key_len == 16) {
            return EVP_aes_128_ctr();
        } else if (key_len == 24) {
            return EVP_aes_192_ctr();
        } else if (key_len == 32) {
            return EVP_aes_256_ctr();
        }
    }
    return nullptr;
}

int Aes::encrypt(signal_buffer **output, int cipher, const uint8_t *key, size_t key_len, const uint8_t *iv,
                 size_t iv_len, const uint8_t *plaintext, size_t plaintext_len, void *) {
    const EVP_CIPHER *evp_cipher = aes_cipher(cipher, key_len);
    if (!evp_cipher) {
        fprintf(stderr, "invalid AES mode or key size: %zu\n", key_len);
        return SG_ERR_UNKNOWN;
    }

    if (iv_len != 16) {
        fprintf(stderr, "invalid AES IV size: %zu\n", iv_len);
        return SG_ERR_UNKNOWN;
    }

    if (plaintext_len > INT_MAX - EVP_CIPHER_block_size(evp_cipher)) {
        fprintf(stderr, "invalid plaintext length: %zu\n", plaintext_len);
        return SG_ERR_UNKNOWN;
    }

    EVPCipherCtxWrapper ctx{};
    if (!ctx.good()) {
        fprintf(stderr, "could not create context\n");
        return SG_ERR_UNKNOWN;
    }

    auto result = EVP_EncryptInit_ex(*ctx, evp_cipher, nullptr, key, iv);
    if (!result) {
        fprintf(stderr, "cannot initialize cipher\n");
        return SG_ERR_UNKNOWN;
    }

    if (cipher == SG_CIPHER_AES_CTR_NOPADDING) {
        result = EVP_CIPHER_CTX_set_padding(*ctx, 0);
        if (!result) {
            fprintf(stderr, "cannot set padding\n");
            return SG_ERR_UNKNOWN;
        }
    }

    auto out_buf = std::make_unique<uint8_t>(plaintext_len + EVP_CIPHER_block_size(evp_cipher));

    int out_len = 0;
    result = EVP_EncryptUpdate(*ctx, out_buf.get(), &out_len, plaintext,
                               plaintext_len);
    if (!result) {
        fprintf(stderr, "cannot encrypt plaintext\n");
        return SG_ERR_UNKNOWN;
    }

    int final_len = 0;
    result = EVP_EncryptFinal_ex(*ctx, out_buf.get() + out_len, &final_len);
    if (!result) {
        fprintf(stderr, "cannot finish encrypting plaintext\n");
        return SG_ERR_UNKNOWN;
    }

    *output = signal_buffer_create(out_buf.get(), out_len + final_len);

    return result;
}

int Aes::decrypt(signal_buffer **output, int cipher, const uint8_t *key, size_t key_len, const uint8_t *iv,
                 size_t iv_len, const uint8_t *ciphertext, size_t ciphertext_len, void *) {
    const EVP_CIPHER *evp_cipher = aes_cipher(cipher, key_len);
    if (!evp_cipher) {
        fprintf(stderr, "invalid AES mode or key size: %zu\n", key_len);
        return SG_ERR_INVAL;
    }

    if (iv_len != 16) {
        fprintf(stderr, "invalid AES IV size: %zu\n", iv_len);
        return SG_ERR_INVAL;
    }

    if (ciphertext_len > INT_MAX - EVP_CIPHER_block_size(evp_cipher)) {
        fprintf(stderr, "invalid ciphertext length: %zu\n", ciphertext_len);
        return SG_ERR_UNKNOWN;
    }

    EVPCipherCtxWrapper ctx{};
    if (!ctx.good()) {
        fprintf(stderr, "could not create context\n");
        return SG_ERR_UNKNOWN;
    }

    auto result = EVP_DecryptInit_ex(*ctx, evp_cipher, nullptr, key, iv);
    if (!result) {
        fprintf(stderr, "cannot initialize cipher\n");
        return SG_ERR_UNKNOWN;
    }

    if (cipher == SG_CIPHER_AES_CTR_NOPADDING) {
        result = EVP_CIPHER_CTX_set_padding(*ctx, 0);
        if (!result) {
            fprintf(stderr, "cannot set padding\n");
            return SG_ERR_UNKNOWN;
        }
    }

    auto out_buf = std::make_unique<uint8_t>(ciphertext_len + EVP_CIPHER_block_size(evp_cipher));

    int out_len = 0;
    result = EVP_DecryptUpdate(*ctx, out_buf.get(), &out_len, ciphertext, ciphertext_len);
    if (!result) {
        fprintf(stderr, "cannot decrypt ciphertext\n");
        return SG_ERR_UNKNOWN;
    }

    int final_len = 0;
    result = EVP_DecryptFinal_ex(*ctx, out_buf.get() + out_len, &final_len);
    if (!result) {
        fprintf(stderr, "cannot finish decrypting ciphertext\n");
        return SG_ERR_UNKNOWN;
    }

    *output = signal_buffer_create(out_buf.get(), out_len + final_len);

    return result;
}