/*
 * crypto.c
 * vim: expandtab:ts=4:sts=4:sw=4
 *
 * Copyright (C) 2019 Paul Fariello <paul@fariello.eu>
 *
 * This file is part of Profanity.
 *
 * Profanity is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Profanity is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Profanity.  If not, see <https://www.gnu.org/licenses/>.
 *
 * In addition, as a special exception, the copyright holders give permission to
 * link the code of portions of this program with the OpenSSL library under
 * certain conditions as described in each individual source file, and
 * distribute linked combinations including the two.
 *
 * You must obey the GNU General Public License in all respects for all of the
 * code used other than OpenSSL. If you modify file(s) with this exception, you
 * may extend this exception to your version of the file(s), but you are not
 * obligated to do so. If you do not wish to do so, delete this exception
 * statement from your version. If you delete this exception statement from all
 * source files in the program, then also delete it here.
 *
 * krypt.sh: libgcrypt replaced with OpenSSL/LibreSSL backend.
 */
#include "config.h"

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

#include <openssl/evp.h>
#include <openssl/hmac.h>
#include <openssl/rand.h>

#include <signal/signal_protocol.h>
#include <signal/signal_protocol_types.h>

#include "log.h"
#include "omemo/omemo.h"
#include "omemo/crypto.h"

#define AES256_GCM_TAG_LENGTH  16
#define AES256_GCM_BUFFER_SIZE 1024

int
omemo_crypto_init(void)
{
    /* OpenSSL/LibreSSL self-initializes. Verify RAND is seeded. */
    unsigned char test[1];
    if (RAND_bytes(test, 1) != 1)
        return -1;
    return 0;
}

int
omemo_random_func(uint8_t* data, size_t len, void* user_data)
{
    return RAND_bytes(data, (int)len) == 1 ? 0 : -1;
}

/*
 * HMAC-SHA256
 */

int
omemo_hmac_sha256_init_func(void** hmac_context, const uint8_t* key, size_t key_len, void* user_data)
{
    HMAC_CTX* ctx = HMAC_CTX_new();
    if (!ctx)
        return -1;

    if (HMAC_Init_ex(ctx, key, (int)key_len, EVP_sha256(), NULL) != 1) {
        HMAC_CTX_free(ctx);
        return -1;
    }

    *hmac_context = ctx;
    return 0;
}

int
omemo_hmac_sha256_update_func(void* hmac_context, const uint8_t* data, size_t data_len, void* user_data)
{
    return HMAC_Update((HMAC_CTX*)hmac_context, data, data_len) == 1 ? 0 : -1;
}

int
omemo_hmac_sha256_final_func(void* hmac_context, signal_buffer** output, void* user_data)
{
    unsigned char md[32];
    unsigned int  len = sizeof(md);

    if (HMAC_Final((HMAC_CTX*)hmac_context, md, &len) != 1)
        return -1;

    *output = signal_buffer_create(md, len);
    return *output ? 0 : -1;
}

void
omemo_hmac_sha256_cleanup_func(void* hmac_context, void* user_data)
{
    HMAC_CTX_free((HMAC_CTX*)hmac_context);
}

/*
 * SHA-512 digest
 */

int
omemo_sha512_digest_init_func(void** digest_context, void* user_data)
{
    EVP_MD_CTX* ctx = EVP_MD_CTX_new();
    if (!ctx)
        return -1;

    if (EVP_DigestInit_ex(ctx, EVP_sha512(), NULL) != 1) {
        EVP_MD_CTX_free(ctx);
        return -1;
    }

    *digest_context = ctx;
    return 0;
}

int
omemo_sha512_digest_update_func(void* digest_context, const uint8_t* data, size_t data_len, void* user_data)
{
    return EVP_DigestUpdate((EVP_MD_CTX*)digest_context, data, data_len) == 1 ? 0 : -1;
}

int
omemo_sha512_digest_final_func(void* digest_context, signal_buffer** output, void* user_data)
{
    unsigned char md[64];
    unsigned int  len = sizeof(md);

    if (EVP_DigestFinal_ex((EVP_MD_CTX*)digest_context, md, &len) != 1)
        return -1;

    *output = signal_buffer_create(md, len);
    return *output ? 0 : -1;
}

void
omemo_sha512_digest_cleanup_func(void* digest_context, void* user_data)
{
    EVP_MD_CTX_free((EVP_MD_CTX*)digest_context);
}

/*
 * AES-256-CBC-PKCS5 encrypt/decrypt
 * Used by libsignal-protocol-c for session key material.
 */

int
omemo_encrypt_func(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* user_data)
{
    EVP_CIPHER_CTX* ctx        = NULL;
    unsigned char*  ciphertext = NULL;
    int             outl = 0, finl = 0;
    int             ret  = OMEMO_ERR_UNSUPPORTED_CRYPTO;

    switch (key_len) {
    case 32: break;
    default: return OMEMO_ERR_UNSUPPORTED_CRYPTO;
    }

    switch (cipher) {
    case SG_CIPHER_AES_CBC_PKCS5: break;
    default: return OMEMO_ERR_UNSUPPORTED_CRYPTO;
    }

    /* worst case: plaintext + one full block of PKCS7 padding */
    ciphertext = malloc(plaintext_len + 16);
    if (!ciphertext)
        return -1;

    ctx = EVP_CIPHER_CTX_new();
    if (!ctx)
        goto out;

    if (EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv) != 1)
        goto out;

    if (EVP_EncryptUpdate(ctx, ciphertext, &outl, plaintext, (int)plaintext_len) != 1)
        goto out;

    if (EVP_EncryptFinal_ex(ctx, ciphertext + outl, &finl) != 1)
        goto out;

    *output = signal_buffer_create(ciphertext, outl + finl);
    ret = *output ? SG_SUCCESS : -1;

out:
    EVP_CIPHER_CTX_free(ctx);
    free(ciphertext);
    return ret;
}

int
omemo_decrypt_func(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* user_data)
{
    EVP_CIPHER_CTX* ctx       = NULL;
    unsigned char*  plaintext = NULL;
    int             outl = 0, finl = 0;
    int             ret  = OMEMO_ERR_UNSUPPORTED_CRYPTO;

    switch (key_len) {
    case 32: break;
    default: return OMEMO_ERR_UNSUPPORTED_CRYPTO;
    }

    switch (cipher) {
    case SG_CIPHER_AES_CBC_PKCS5: break;
    default: return OMEMO_ERR_UNSUPPORTED_CRYPTO;
    }

    plaintext = malloc(ciphertext_len);
    if (!plaintext)
        return -1;

    ctx = EVP_CIPHER_CTX_new();
    if (!ctx)
        goto out;

    if (EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv) != 1)
        goto out;

    if (EVP_DecryptUpdate(ctx, plaintext, &outl, ciphertext, (int)ciphertext_len) != 1)
        goto out;

    if (EVP_DecryptFinal_ex(ctx, plaintext + outl, &finl) != 1)
        goto out;

    *output = signal_buffer_create(plaintext, outl + finl);
    ret = *output ? SG_SUCCESS : -1;

out:
    EVP_CIPHER_CTX_free(ctx);
    free(plaintext);
    return ret;
}

/*
 * AES-128-GCM encrypt — used for the actual OMEMO message payload.
 * Tag is returned separately from ciphertext.
 */

int
aes128gcm_encrypt(unsigned char* ciphertext, size_t* ciphertext_len,
                  unsigned char* tag, size_t* tag_len,
                  const unsigned char* const plaintext, size_t plaintext_len,
                  const unsigned char* const iv, const unsigned char* const key)
{
    EVP_CIPHER_CTX* ctx  = NULL;
    int             outl = 0, finl = 0;
    int             ret  = -1;

    ctx = EVP_CIPHER_CTX_new();
    if (!ctx)
        goto out;

    if (EVP_EncryptInit_ex(ctx, EVP_aes_128_gcm(), NULL, NULL, NULL) != 1)
        goto out;

    if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, AES128_GCM_IV_LENGTH, NULL) != 1)
        goto out;

    if (EVP_EncryptInit_ex(ctx, NULL, NULL, key, iv) != 1)
        goto out;

    if (EVP_EncryptUpdate(ctx, ciphertext, &outl, plaintext, (int)plaintext_len) != 1)
        goto out;

    if (EVP_EncryptFinal_ex(ctx, ciphertext + outl, &finl) != 1)
        goto out;

    *ciphertext_len = outl + finl;

    if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, AES128_GCM_TAG_LENGTH, tag) != 1)
        goto out;

    *tag_len = AES128_GCM_TAG_LENGTH;
    ret = 0;

out:
    EVP_CIPHER_CTX_free(ctx);
    return ret;
}

/*
 * AES-128-GCM decrypt.
 * Tag is passed as a separate last parameter (not appended to ciphertext).
 * iv_len is honoured so both 12-byte and 16-byte IVs work.
 */

int
aes128gcm_decrypt(unsigned char* plaintext, size_t* plaintext_len,
                  const unsigned char* const ciphertext, size_t ciphertext_len,
                  const unsigned char* const iv, size_t iv_len,
                  const unsigned char* const key, const unsigned char* const tag)
{
    EVP_CIPHER_CTX* ctx  = NULL;
    int             outl = 0, finl = 0;
    int             ret  = -1;

    ctx = EVP_CIPHER_CTX_new();
    if (!ctx)
        goto out;

    if (EVP_DecryptInit_ex(ctx, EVP_aes_128_gcm(), NULL, NULL, NULL) != 1)
        goto out;

    if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, (int)iv_len, NULL) != 1)
        goto out;

    if (EVP_DecryptInit_ex(ctx, NULL, NULL, key, iv) != 1)
        goto out;

    if (EVP_DecryptUpdate(ctx, plaintext, &outl, ciphertext, (int)ciphertext_len) != 1)
        goto out;

    /* must set tag before calling Final */
    if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG,
            AES128_GCM_TAG_LENGTH, (void*)tag) != 1)
        goto out;

    if (EVP_DecryptFinal_ex(ctx, plaintext + outl, &finl) != 1) {
        log_error("OMEMO: aes128gcm_decrypt tag verification failed");
        goto out;
    }

    *plaintext_len = outl + finl;
    ret = 0;

out:
    EVP_CIPHER_CTX_free(ctx);
    return ret;
}

/*
 * AES-256-GCM file encrypt/decrypt.
 * Single function: gboolean encrypt selects direction.
 * Returns 0 on success, -1 on failure.
 */

int
aes256gcm_crypt_file(FILE* in, FILE* out, off_t file_size,
                     unsigned char key[], unsigned char nonce[], gboolean encrypt)
{
    EVP_CIPHER_CTX* ctx  = NULL;
    unsigned char   buf[AES256_GCM_BUFFER_SIZE];
    unsigned char   outbuf[AES256_GCM_BUFFER_SIZE + 16];
    unsigned char   tag[AES256_GCM_TAG_LENGTH];
    off_t           bytes_remaining;
    int             bytes, outl, finl;
    int             ret = -1;

    ctx = EVP_CIPHER_CTX_new();
    if (!ctx)
        goto out;

    if (encrypt) {
        if (EVP_EncryptInit_ex(ctx, EVP_aes_256_gcm(), NULL, NULL, NULL) != 1)
            goto out;
        if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN,
                OMEMO_AESGCM_NONCE_LENGTH, NULL) != 1)
            goto out;
        if (EVP_EncryptInit_ex(ctx, NULL, NULL, key, nonce) != 1)
            goto out;

        bytes_remaining = file_size;
        while (bytes_remaining > 0) {
            size_t read_size = bytes_remaining < (off_t)sizeof(buf)
                                   ? (size_t)bytes_remaining : sizeof(buf);
            bytes = fread(buf, 1, read_size, in);
            if (bytes <= 0)
                break;
            if (EVP_EncryptUpdate(ctx, outbuf, &outl, buf, bytes) != 1)
                goto out;
            if (outl > 0 && fwrite(outbuf, 1, outl, out) != (size_t)outl)
                goto out;
            bytes_remaining -= bytes;
        }
        if (ferror(in))
            goto out;

        if (EVP_EncryptFinal_ex(ctx, outbuf, &finl) != 1)
            goto out;
        if (finl > 0 && fwrite(outbuf, 1, finl, out) != (size_t)finl)
            goto out;

        /* append authentication tag */
        if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG,
                AES256_GCM_TAG_LENGTH, tag) != 1)
            goto out;
        if (fwrite(tag, 1, AES256_GCM_TAG_LENGTH, out) != AES256_GCM_TAG_LENGTH)
            goto out;

    } else {
        /* decrypt: tag is the last AES256_GCM_TAG_LENGTH bytes of the file */
        if (file_size < AES256_GCM_TAG_LENGTH)
            goto out;
        off_t data_size = file_size - AES256_GCM_TAG_LENGTH;

        if (EVP_DecryptInit_ex(ctx, EVP_aes_256_gcm(), NULL, NULL, NULL) != 1)
            goto out;
        if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN,
                OMEMO_AESGCM_NONCE_LENGTH, NULL) != 1)
            goto out;
        if (EVP_DecryptInit_ex(ctx, NULL, NULL, key, nonce) != 1)
            goto out;

        bytes_remaining = data_size;
        while (bytes_remaining > 0) {
            size_t read_size = bytes_remaining < (off_t)sizeof(buf)
                                   ? (size_t)bytes_remaining : sizeof(buf);
            bytes = fread(buf, 1, read_size, in);
            if (bytes <= 0)
                break;
            if (EVP_DecryptUpdate(ctx, outbuf, &outl, buf, bytes) != 1)
                goto out;
            if (outl > 0 && fwrite(outbuf, 1, outl, out) != (size_t)outl)
                goto out;
            bytes_remaining -= bytes;
        }
        if (ferror(in))
            goto out;

        /* read and verify tag */
        if (fread(tag, 1, AES256_GCM_TAG_LENGTH, in) != AES256_GCM_TAG_LENGTH)
            goto out;
        if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG,
                AES256_GCM_TAG_LENGTH, tag) != 1)
            goto out;
        if (EVP_DecryptFinal_ex(ctx, outbuf, &finl) != 1) {
            log_error("OMEMO: aes256gcm_crypt_file tag verification failed");
            goto out;
        }
        if (finl > 0 && fwrite(outbuf, 1, finl, out) != (size_t)finl)
            goto out;
    }

    ret = 0;

out:
    EVP_CIPHER_CTX_free(ctx);
    return ret;
}

/*
 * Build the aesgcm:// URL fragment: hex(nonce) + hex(key)
 * Original used gcry_malloc_secure; plain malloc suffices here —
 * this string is a URL fragment, not long-lived key material.
 */

char*
aes256gcm_create_secure_fragment(unsigned char* key, unsigned char* nonce)
{
    int   key_size   = OMEMO_AESGCM_KEY_LENGTH;
    int   nonce_size = OMEMO_AESGCM_NONCE_LENGTH;
    char* fragment   = malloc((nonce_size + key_size) * 2 + 1);

    if (!fragment)
        return NULL;

    for (int i = 0; i < nonce_size; i++)
        sprintf(&(fragment[i * 2]), "%02x", nonce[i]);
    for (int i = 0; i < key_size; i++)
        sprintf(&(fragment[(i + nonce_size) * 2]), "%02x", key[i]);

    return fragment;
}