/*
 * Copyright (c) 2023-2024 Ian Marco Moffett and the Osmora Team.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 * 3. Neither the name of Hyra nor the names of its
 *    contributors may be used to endorse or promote products derived from
 *    this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include <sys/mman.h>
#include <openssl/evp.h>
#include <openssl/aes.h>
#include <openssl/rand.h>
#include <crypto/aes.h>
#include <stdio.h>

/*
 * Encrypt 'data' with AES-256-GCM.
 *
 * @data: Data to be encrypted.
 * @len: Length of data.
 * @key: Private key to use.
 * @res: Will contains the new ciphertext and IV.
 *
 * TODO: Transparently authenticate data with AAD.
 */
int
aes256_encrypt(const unsigned char *data, size_t len,
    const unsigned char *key, struct aes_message *res)
{
    EVP_CIPHER_CTX *ctx = NULL;
    size_t ciphertext_len;
    int error, tmp, *lenres;

    if (res == NULL) {
        printf("aes256_encrypt: 'res' is NULL!\n");
        return -1;
    }

    lenres = &res->ciphertext_len;
    ciphertext_len = 0;
    res->ciphertext = malloc(len + AES_GCM_TAG_SIZE);
    res->ciphertext_len = ciphertext_len;

    /* Generate random bytes for IV */
    if (RAND_bytes(res->iv, AES_IV_SIZE) <= 0) {
        printf("aes256_encrypt: Failed to generate AES-256 IV\n");
        free(res->ciphertext);
        return -1;
    }

    ctx = EVP_CIPHER_CTX_new();
    if (ctx == NULL) {
        printf("aes256_encrypt: Failed to create EVP PKEY context\n");
        free(res->ciphertext);
        return -1;
    }

    /* Initialize encryption operation */
    error = EVP_EncryptInit_ex(ctx, EVP_aes_256_gcm(), NULL, NULL, NULL);
    if (error <= 0) {
        printf("aes256_encrypt: Failed to initialize AES context\n");
        EVP_CIPHER_CTX_free(ctx);
        free(res->ciphertext);
        return -1;
    }

    /* Setup the IV length */
    error = EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, AES_IV_SIZE, NULL);
    if (error <= 0) {
        printf("aes256_encrypt: Failed to initialize context IV length\n");
        EVP_CIPHER_CTX_free(ctx);
        free(res->ciphertext);
        return -1;
    }

    /* Initialize key and IV */
    error = EVP_EncryptInit_ex(ctx, NULL, NULL, key, res->iv);
    if (error <= 0) {
        printf("aes256_encrypt: Failed to initialize key and IV\n");
        EVP_CIPHER_CTX_free(ctx);
        free(res->ciphertext);
        return -1;
    }

    error = EVP_EncryptUpdate(ctx, res->ciphertext, &tmp, data, len);
    if (error <= 0) {
        printf("aes256_encrypt: Failed to add plaintext\n");
        EVP_CIPHER_CTX_free(ctx);
        free(res->ciphertext);
        return -1;
    }

    res->ciphertext_len = tmp;
    error = EVP_EncryptFinal_ex(ctx, res->ciphertext + tmp, &tmp);
    if (error <= 0) {
        printf("aes256_encrypt: Failed to finish encryption\n");
        EVP_CIPHER_CTX_free(ctx);
        free(res->ciphertext);
        return -1;
    }

    res->ciphertext_len += tmp;
    error = EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, AES_GCM_TAG_SIZE,
        res->tag);

    if (error <= 0) {
        printf("aes256_encrypt: Failed get GCM tag\n");
        EVP_CIPHER_CTX_free(ctx);
        free(res->ciphertext);
        return -1;
    }

    EVP_CIPHER_CTX_free(ctx);
    return 0;
}


/*
 * Decrypt 'data' with AES-256-GCM.
 *
 * @amp: Contains ciphertext and IV.
 * @key: Private key to use.
 * @res: Will be set to a buffer containing plaintext.
 */
int
aes256_decrypt(struct aes_message *amp, const unsigned char *key,
    unsigned char **res)
{
    EVP_CIPHER_CTX *ctx = NULL;
    int len, error;
    unsigned char *plaintext;

    if (res == NULL) {
        printf("aes256_decrypt: 'res' is NULL\n");
        return -1;
    }

    plaintext = malloc(amp->ciphertext_len);
    if (plaintext == NULL) {
        printf("aes256_decrypt: Failed to allocate plaintext memory\n");
        return -1;
    }

    /* Lock plaintext in memory */
    if (mlock(plaintext, amp->ciphertext_len) != 0) {
        printf("aes256_decrypt: Failed to lock plaintext memory\n");
        return -1;
    }

    ctx = EVP_CIPHER_CTX_new();
    if (ctx == NULL) {
        printf("aes256_decrypt: Failed to create cipher context\n");
        free(plaintext);
        return -1;
    }

    error = EVP_DecryptInit_ex(ctx, EVP_aes_256_gcm(), NULL, NULL, NULL);
    if (error <= 0) {
        printf("aes256_decrypt: Failed to initialize decryption\n");
        EVP_CIPHER_CTX_free(ctx);
        free(plaintext);
        return -1;
    }

    /* Set IV length */
    error = EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, AES_IV_SIZE, NULL);
    if (error <= 0) {
        printf("aes256_decrypt: Failed to set IV length\n");
        EVP_CIPHER_CTX_free(ctx);
        free(plaintext);
        return -1;
    }

    /* Init key and IV */
    error = EVP_DecryptInit_ex(ctx, NULL, NULL, key, amp->iv);
    if (error <= 0) {
        printf("aes256_decrypt: Failed to set key and IV\n");
        EVP_CIPHER_CTX_free(ctx);
        free(plaintext);
        return -1;
    }

    error = EVP_DecryptUpdate(ctx, plaintext, &len, amp->ciphertext,
        amp->ciphertext_len);

    if (error <= 0) {
        printf("aes256_decrypt: Failed to set plaintext\n");
        EVP_CIPHER_CTX_free(ctx);
        free(plaintext);
        return -1;
    }

    /* Set GCM tag */
    error = EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, AES_GCM_TAG_SIZE,
        amp->tag);

    if (error <= 0) {
        printf("aes256_decrypt: Failed to set tag\n");
        EVP_CIPHER_CTX_free(ctx);
        free(plaintext);
        return -1;
    }

    error = EVP_DecryptFinal_ex(ctx, plaintext + len, &len);
    if (error <= 0) {
        printf("aes256_decrypt: Could not verify plaintext\n");
        EVP_CIPHER_CTX_free(ctx);
        free(plaintext);
        return -1;
    }

    *res = plaintext;
    return 0;
}