/*
 * crypto.c
 * vim: expandtab:ts=4:sts=4:sw=4
 *
 * Copyright (C) 2019 Paul Fariello <paul@fariello.eu>
 *
 * SPDX-License-Identifier: GPL-3.0-or-later WITH OpenSSL-exception
 */
#include "config.h"

#include <assert.h>
#ifdef HAVE_LIBOMEMO_C
#include <omemo/signal_protocol.h>
#include <omemo/signal_protocol_types.h>
#else
#include <signal/signal_protocol.h>
#include <signal/signal_protocol_types.h>
#endif

#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)
{
    if (!gcry_check_version(GCRYPT_VERSION)) {
        return GPG_ERR_UNKNOWN_VERSION;
    }

    gcry_error_t ret;
    ret = gcry_control(GCRYCTL_SUSPEND_SECMEM_WARN);
    if (ret != 0)
        return ret;

    ret = gcry_control(GCRYCTL_INIT_SECMEM, 16384, 0);
    if (ret != 0)
        return ret;

    ret = gcry_control(GCRYCTL_RESUME_SECMEM_WARN);
    if (ret != 0)
        return ret;

    ret = gcry_control(GCRYCTL_INITIALIZATION_FINISHED, 0);
    if (ret != 0)
        return ret;

    /* Ask for a first random buffer to ensure CSPRNG is initialized.
     * Thus we control the memleak produced by gcrypt initialization and we can
     * suppress it without having false negatives */
    gcry_free(gcry_random_bytes_secure(1, GCRY_VERY_STRONG_RANDOM));

    return 0;
}

int
omemo_random_func(uint8_t* data, size_t len, void* user_data)
{
    gcry_randomize(data, len, GCRY_VERY_STRONG_RANDOM);
    return 0;
}

int
omemo_hmac_sha256_init_func(void** hmac_context, const uint8_t* key, size_t key_len, void* user_data)
{
    gcry_error_t res;
    gcry_mac_hd_t hd;

    res = gcry_mac_open(&hd, GCRY_MAC_HMAC_SHA256, 0, NULL);
    if (res != GPG_ERR_NO_ERROR) {
        log_error("OMEMO: %s", gcry_strerror(res));
        return OMEMO_ERR_GCRYPT;
    }

    *hmac_context = hd;
    res = gcry_mac_setkey(hd, key, key_len);
    if (res != GPG_ERR_NO_ERROR) {
        log_error("OMEMO: %s", gcry_strerror(res));
        return OMEMO_ERR_GCRYPT;
    }

    return 0;
}

int
omemo_hmac_sha256_update_func(void* hmac_context, const uint8_t* data, size_t data_len, void* user_data)
{
    gcry_error_t res;

    res = gcry_mac_write(hmac_context, data, data_len);
    if (res != GPG_ERR_NO_ERROR) {
        log_error("OMEMO: %s", gcry_strerror(res));
        return OMEMO_ERR_GCRYPT;
    }

    return 0;
}

int
omemo_hmac_sha256_final_func(void* hmac_context, signal_buffer** output, void* user_data)
{
    gcry_error_t res;
    size_t mac_len = 32;
    unsigned char out[mac_len];

    res = gcry_mac_read(hmac_context, out, &mac_len);
    if (res != GPG_ERR_NO_ERROR) {
        log_error("OMEMO: %s", gcry_strerror(res));
        return OMEMO_ERR_GCRYPT;
    }

    *output = signal_buffer_create(out, mac_len);
    return 0;
}

void
omemo_hmac_sha256_cleanup_func(void* hmac_context, void* user_data)
{
    gcry_mac_close(hmac_context);
}

int
omemo_sha512_digest_init_func(void** digest_context, void* user_data)
{
    gcry_error_t res;
    gcry_md_hd_t hd;

    res = gcry_md_open(&hd, GCRY_MD_SHA512, 0);
    if (res != GPG_ERR_NO_ERROR) {
        log_error("OMEMO: %s", gcry_strerror(res));
        return OMEMO_ERR_GCRYPT;
    }

    *digest_context = hd;

    return 0;
}

int
omemo_sha512_digest_update_func(void* digest_context, const uint8_t* data, size_t data_len, void* user_data)
{
    gcry_md_write(digest_context, data, data_len);

    return 0;
}

int
omemo_sha512_digest_final_func(void* digest_context, signal_buffer** output, void* user_data)
{
    gcry_error_t res;
    unsigned char out[64];

    res = gcry_md_extract(digest_context, GCRY_MD_SHA512, out, 64);
    if (res != GPG_ERR_NO_ERROR) {
        log_error("OMEMO: %s", gcry_strerror(res));
        return OMEMO_ERR_GCRYPT;
    }

    *output = signal_buffer_create(out, 64);
    return 0;
}

void
omemo_sha512_digest_cleanup_func(void* digest_context, void* user_data)
{
    gcry_md_close(digest_context);
}

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)
{
    gcry_cipher_hd_t hd;
    unsigned char* padded_plaintext;
    unsigned char* ciphertext;
    size_t ciphertext_len;
    int mode;
    int algo;
    uint8_t padding = 0;

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

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

    gcry_cipher_open(&hd, algo, mode, GCRY_CIPHER_SECURE);

    gcry_cipher_setkey(hd, key, key_len);

    switch (cipher) {
    case SG_CIPHER_AES_CBC_PKCS5:
        gcry_cipher_setiv(hd, iv, iv_len);
        padding = 16 - (plaintext_len % 16);
        break;
    default:
        assert(FALSE);
    }

    padded_plaintext = malloc(plaintext_len + padding);
    memcpy(padded_plaintext, plaintext, plaintext_len);
    memset(padded_plaintext + plaintext_len, padding, padding);

    ciphertext_len = plaintext_len + padding;
    ciphertext = malloc(ciphertext_len);
    gcry_cipher_encrypt(hd, ciphertext, ciphertext_len, padded_plaintext, plaintext_len + padding);

    *output = signal_buffer_create(ciphertext, ciphertext_len);
    free(padded_plaintext);
    free(ciphertext);

    gcry_cipher_close(hd);

    return SG_SUCCESS;
}

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)
{
    int ret = SG_SUCCESS;
    gcry_cipher_hd_t hd;
    unsigned char* plaintext = NULL;
    size_t plaintext_len;
    int mode;
    int algo;
    uint8_t padding = 0;

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

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

    gcry_cipher_open(&hd, algo, mode, GCRY_CIPHER_SECURE);

    gcry_cipher_setkey(hd, key, key_len);

    switch (cipher) {
    case SG_CIPHER_AES_CBC_PKCS5:
        gcry_cipher_setiv(hd, iv, iv_len);
        break;
    default:
        assert(FALSE);
    }

    plaintext_len = ciphertext_len;
    // PKCS#5/PKCS#7 padding requires at least one byte of plaintext;
    // reject malformed empty ciphertext before any plaintext[len - 1]
    // access could underflow / read out-of-bounds.
    if (plaintext_len == 0) {
        ret = SG_ERR_INVAL;
        goto out;
    }
    plaintext = malloc(plaintext_len);
    gcry_cipher_decrypt(hd, plaintext, plaintext_len, ciphertext, ciphertext_len);

    switch (cipher) {
    case SG_CIPHER_AES_CBC_PKCS5:
        padding = plaintext[plaintext_len - 1];
        break;
    default:
        assert(FALSE);
    }

    // padding byte must address bytes that exist in plaintext.
    if (padding == 0 || (size_t)padding > plaintext_len) {
        ret = SG_ERR_UNKNOWN;
        goto out;
    }

    for (int i = 0; i < padding; i++) {
        if (plaintext[plaintext_len - 1 - i] != padding) {
            ret = SG_ERR_UNKNOWN;
            goto out;
        }
    }

    *output = signal_buffer_create(plaintext, plaintext_len - padding);

out:
    free(plaintext);

    gcry_cipher_close(hd);

    return ret;
}

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)
{
    gcry_error_t res;
    gcry_cipher_hd_t hd;

    res = gcry_cipher_open(&hd, GCRY_CIPHER_AES128, GCRY_CIPHER_MODE_GCM, GCRY_CIPHER_SECURE);
    if (res != GPG_ERR_NO_ERROR) {
        goto out;
    }
    res = gcry_cipher_setkey(hd, key, AES128_GCM_KEY_LENGTH);
    if (res != GPG_ERR_NO_ERROR) {
        goto out;
    }
    res = gcry_cipher_setiv(hd, iv, AES128_GCM_IV_LENGTH);
    if (res != GPG_ERR_NO_ERROR) {
        goto out;
    }

    res = gcry_cipher_encrypt(hd, ciphertext, *ciphertext_len, plaintext, plaintext_len);
    if (res != GPG_ERR_NO_ERROR) {
        goto out;
    }

    res = gcry_cipher_gettag(hd, tag, *tag_len);
    if (res != GPG_ERR_NO_ERROR) {
        goto out;
    }

out:
    gcry_cipher_close(hd);
    return res;
}

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)
{
    gcry_error_t res;
    gcry_cipher_hd_t hd;

    res = gcry_cipher_open(&hd, GCRY_CIPHER_AES128, GCRY_CIPHER_MODE_GCM, GCRY_CIPHER_SECURE);
    if (res != GPG_ERR_NO_ERROR) {
        goto out;
    }

    res = gcry_cipher_setkey(hd, key, AES128_GCM_KEY_LENGTH);
    if (res != GPG_ERR_NO_ERROR) {
        goto out;
    }

    res = gcry_cipher_setiv(hd, iv, iv_len);
    if (res != GPG_ERR_NO_ERROR) {
        goto out;
    }

    res = gcry_cipher_decrypt(hd, plaintext, *plaintext_len, ciphertext, ciphertext_len);
    if (res != GPG_ERR_NO_ERROR) {
        goto out;
    }

    res = gcry_cipher_checktag(hd, tag, AES128_GCM_TAG_LENGTH);
    if (res != GPG_ERR_NO_ERROR) {
        goto out;
    }

out:
    gcry_cipher_close(hd);
    return res;
}

gcry_error_t
aes256gcm_crypt_file(FILE* in, FILE* out, off_t file_size,
                     unsigned char key[], unsigned char nonce[], gboolean encrypt)
{

    if (!gcry_control(GCRYCTL_INITIALIZATION_FINISHED_P)) {
        fputs("libgcrypt has not been initialized\n", stderr);
        abort();
    }

    if (!encrypt) {
        file_size -= AES256_GCM_TAG_LENGTH;
    }

    gcry_error_t res;
    gcry_cipher_hd_t hd;

    res = gcry_cipher_open(&hd, GCRY_CIPHER_AES256, GCRY_CIPHER_MODE_GCM,
                           GCRY_CIPHER_SECURE);
    if (res != GPG_ERR_NO_ERROR) {
        goto out;
    }

    res = gcry_cipher_setkey(hd, key, OMEMO_AESGCM_KEY_LENGTH);
    if (res != GPG_ERR_NO_ERROR) {
        goto out;
    }

    res = gcry_cipher_setiv(hd, nonce, OMEMO_AESGCM_NONCE_LENGTH);
    if (res != GPG_ERR_NO_ERROR) {
        goto out;
    }

    unsigned char buffer[AES256_GCM_BUFFER_SIZE];

    size_t bytes = 0;
    off_t bytes_read = 0, bytes_available = 0, read_size = 0;
    while (bytes_read < file_size) {
        bytes_available = file_size - bytes_read;
        if (!bytes_available || ferror(in) != 0) {
            break;
        }

        if (bytes_available < AES256_GCM_BUFFER_SIZE) {
            read_size = bytes_available;
            gcry_cipher_final(hd); // Signal last round of bytes.
        } else {
            read_size = AES256_GCM_BUFFER_SIZE;
        }

        bytes = fread(buffer, 1, (size_t)read_size, in);
        bytes_read += (off_t)bytes;

        if (encrypt) {
            res = gcry_cipher_encrypt(hd, buffer, bytes, NULL, 0);
        } else {
            res = gcry_cipher_decrypt(hd, buffer, bytes, NULL, 0);
        }

        if (res != GPG_ERR_NO_ERROR) {
            goto out;
        }

        fwrite(buffer, 1, bytes, out);
    }

    unsigned char tag[AES256_GCM_TAG_LENGTH];

    if (encrypt) {
        // Append authentication tag at the end of the file.
        res = gcry_cipher_gettag(hd, tag, AES256_GCM_TAG_LENGTH);
        if (res != GPG_ERR_NO_ERROR) {
            goto out;
        }

        fwrite(tag, 1, AES256_GCM_TAG_LENGTH, out);

    } else {
        // Read and verify authentication tag stored at the end of the file.
        bytes = fread(tag, 1, AES256_GCM_TAG_LENGTH, in);
        res = gcry_cipher_checktag(hd, tag, bytes);
    }

out:
    gcry_cipher_close(hd);
    return res;
}

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 = gcry_malloc_secure((nonce_size + key_size) * 2 + 1);

    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;
}