#include "config.h"

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

#include <sys/types.h>
#include <unistd.h>

/*
 * This is a set of workarounds for the different exim local_scan ABI versions and distribution patches.
 *
 * List of local_scan ABI versions per distribution:
 *
 * 2.0  Ubuntu 18-04 (Bionic)
 * 2.0  Debian 10 (Buster)
 * 3.1  Ubuntu 20-04 (Focal)
 * 4.1  Debian 11 (Bullseye)
 */

#if LOCAL_SCAN_VERSION < 3
#define LOCAL_SCAN

#include <exim4/local_scan.h>

#define store_get_untainted(size) (store_get(size))
#define store_get_tainted(size) (store_get(size))
#elif LOCAL_SCAN_VERSION >= 3
#define DLFUNC_IMPL
#include <exim4/local_scan.h>
#define store_get_untainted(size) (store_get(size, FALSE))
#define store_get_tainted(size) (store_get(size, TRUE))
#else
#error "exim local_scan API version unknown"
#endif

/*
 * Encrypt message using crypto_secretbox_easy().
 *
 * Arguments:
 * 1. password
 * 2. cleartext message
 */
int sodium_crypto_secretbox_encrypt_password(uschar **yield, int argc, uschar *argv[]) {
    // ensure libsodium is initialized
    if (sodium_init() == -1) {
        *yield = string_copy(US
        "Unable to initialize libsodium");
        return ERROR;
    }
    // check argument count
    if (argc != 2) {
        *yield = string_sprintf("Wrong number of arguments (got %i, expected 2)", argc);
        return ERROR;
    }
    // get password argument
    unsigned char *password = argv[0];
    size_t passwordlen = strlen((const char *) password);

    // get cleartext message argument
    unsigned char *message = argv[1];
    size_t messagelen = strlen((const char *) message);

    /*
     * Derive a key from the password using a generic hash.
     * This operation needs to be fast (exim holds no state, this might be called multiple times per email).
     * Collisions avoidance or brute force attacks are not a concern here.
     */
    unsigned char keybytes[crypto_secretbox_KEYBYTES];
    sodium_memzero(keybytes, crypto_secretbox_KEYBYTES);
    crypto_generichash(keybytes, crypto_secretbox_KEYBYTES,
                       password, passwordlen, NULL, 0);

    // prepare buffer for ciphertext
    size_t cipherlen = messagelen + crypto_secretbox_MACBYTES;
    unsigned char *ciphertext = (unsigned char *) store_get_untainted(cipherlen);
    sodium_memzero(ciphertext, cipherlen);

    // encrypt message
    unsigned char nonce[crypto_secretbox_NONCEBYTES];
    randombytes_buf(nonce, sizeof nonce);
    if (crypto_secretbox_easy(ciphertext, message, messagelen, nonce, keybytes) != 0) {
        *yield = string_copy((unsigned char *) "Encryption error after crypto_secretbox_easy()");
        return ERROR;
    }

    // combine nonce and ciphertext
    size_t combined_message_len = crypto_secretbox_NONCEBYTES + cipherlen;
    unsigned char *combined_message = store_get_untainted(combined_message_len);
    memcpy(combined_message, nonce, crypto_secretbox_NONCEBYTES);
    memcpy(&combined_message[crypto_secretbox_NONCEBYTES], ciphertext, cipherlen);

    // base64-encode the ciphertext
    size_t outputsize = sodium_base64_ENCODED_LEN(combined_message_len,
                                                  sodium_base64_VARIANT_ORIGINAL);
    unsigned char *outstring = (unsigned char *) store_get_untainted(outputsize);
    //sodium_memzero(outstring, outputsize);
    sodium_bin2base64((char *const) outstring, outputsize,
                      combined_message, combined_message_len,
                      sodium_base64_VARIANT_ORIGINAL);

    // return base64-encoded ciphertext
    *yield = string_copy(outstring);
    return OK;
}

/*
 * Decrypt message encrypted by sodium_crypto_secretbox_encrypt_password(). Arguments:
 * 1. password string
 * 2. encrypted message
 */
int sodium_crypto_secretbox_decrypt_password(uschar **yield, int argc, uschar *argv[]) {
    // ensure libsodium is initialized
    if (sodium_init() == -1) {
        *yield = string_copy(US
        "Unable to initialize libsodium");
        return ERROR;
    }
    // check argument count
    if (argc != 2) {
        *yield =
                string_sprintf
                        ("Wrong number of arguments (got %i, expected 2)", argc);
        return ERROR;
    }
    // get password
    unsigned char *password = argv[0];
    size_t passwordlen = strlen((const char *) password);

    // Derive key from the password.
    unsigned char keybytes[crypto_secretbox_KEYBYTES];
    sodium_memzero(keybytes, crypto_secretbox_KEYBYTES);
    crypto_generichash(keybytes, crypto_secretbox_KEYBYTES,
                       password, passwordlen, NULL, 0);

    // get base64 encoded ciphertext message
    unsigned char *ciphertextb64 = argv[1];
    size_t ciphertextb64_len = strlen((const char *) ciphertextb64);

    // base64-decode the ciphertext
    size_t combined_message_buf_len = ciphertextb64_len / 4 * 3;
    size_t combined_message_len;
    unsigned char *combined_message = (unsigned char *) store_get_untainted(combined_message_buf_len);
    sodium_memzero(combined_message, combined_message_buf_len);
    int b64err = sodium_base642bin(combined_message, combined_message_buf_len,
                                   (const char *) ciphertextb64, ciphertextb64_len,
                                   NULL, &combined_message_len, NULL,
                                   sodium_base64_VARIANT_ORIGINAL);
    if (b64err != 0) {
        *yield = string_copy((unsigned char *) "Error decoding base64 encoded ciphertext");
        return ERROR;
    }

    // extract nonce
    unsigned char nonce[crypto_secretbox_NONCEBYTES];
    memcpy(nonce, combined_message, crypto_secretbox_NONCEBYTES);

    // prepare buffer for cleartext
    size_t cleartextlen = combined_message_len - crypto_secretbox_NONCEBYTES - crypto_secretbox_MACBYTES;
    unsigned char *cleartext = (unsigned char *) store_get_untainted(cleartextlen + 1);
    sodium_memzero(cleartext, cleartextlen + 1);

    // decrypt message
    if (crypto_secretbox_open_easy(cleartext, &combined_message[crypto_secretbox_NONCEBYTES],
                                   combined_message_len - crypto_secretbox_NONCEBYTES, nonce, keybytes) != 0) {
        *yield = string_copy((unsigned char *) "Decryption error after crypto_secretbox_open_easy()");
        return ERROR;
    }

    // return cleartext
    *yield = string_copy(cleartext);
    return OK;
}

// --------------------------------


/*
 * Encrypt message using crypto_box_seal().
 *
 * Arguments:
 * 1. public key
 * 2. cleartext message
 */
int sodium_crypto_box_seal(uschar **yield, int argc, uschar *argv[]) {
    if (sodium_init() == -1) {
        *yield = string_copy(US
        "Unable to initialize libsodium");
        return ERROR;
    }
    if (argc != 2) {
        *yield = string_sprintf("Wrong number of arguments (got %i, expected 2)", argc);
        return ERROR;
    }

    // get and convert public key
    unsigned char *pkb64 = argv[0];
    size_t pkb64_len = strlen((const char *) pkb64);
    // reserve space for conversion
    size_t pk_buffer_len = crypto_box_PUBLICKEYBYTES; // pkb64_len / 4 * 3 + 1;
    unsigned char *pk = (unsigned char *) store_get_untainted(pk_buffer_len);
    sodium_memzero(pk, pk_buffer_len);
    // convert encoded key to raw form
    int b64err = sodium_base642bin(pk, pk_buffer_len,
                                   (const char *) pkb64, pkb64_len,
                                   NULL, NULL, NULL, sodium_base64_VARIANT_ORIGINAL);
    if (b64err == -1) {
        *yield = string_copy((unsigned char *) "Error decoding public key");
        return ERROR;
    }

    // get cleartext message
    unsigned char *message = argv[1];
    size_t messagelen = strlen((const char *) message);

    // prepare buffer for ciphertext
    size_t cipherlen = messagelen + crypto_box_SEALBYTES;
    unsigned char *ciphertext = store_get_untainted(cipherlen);
    sodium_memzero(ciphertext, cipherlen);

    // encrypt message
    if (crypto_box_seal(ciphertext, message, messagelen, pk) == -1) {
        *yield = string_copy((unsigned char *) "Encryption error after crypto_box_seal()");
        return ERROR;
    }

    // base64-encode the ciphertext
    size_t outputsize = sodium_base64_ENCODED_LEN(cipherlen,
                                                  sodium_base64_VARIANT_ORIGINAL);
    unsigned char *outstring = store_get_untainted(outputsize);
    sodium_bin2base64((char *const) outstring, outputsize,
                      ciphertext, cipherlen,
                      sodium_base64_VARIANT_ORIGINAL);

    // return base64-encoded ciphertext
    *yield = string_copy(outstring);
    return OK;
}

/*
 * Decrypt message encrypted by sodium_crypto_box_seal().
 *
 * Arguments:
 * 1. private key
 * 2. public key
 * 3. encrypted message
 */
int sodium_crypto_box_seal_open(uschar **yield, int argc, uschar *argv[]) {
    if (sodium_init() == -1) {
        *yield = string_copy(US
        "Unable to initialize libsodium");
        return ERROR;
    }
    if (argc != 3) {
        *yield = string_sprintf("Wrong number of arguments (got %i, expected 2)", argc);
        return ERROR;
    }

    // get and convert private key
    unsigned char *skb64 = argv[0];
    size_t skb64_len = strlen((const char *) skb64);
    // reserve space for conversion
    size_t sk_buffer_len = crypto_box_SECRETKEYBYTES;// skb64_len / 4 * 3;
    unsigned char *sk = (unsigned char *) store_get_untainted(sk_buffer_len);
    sodium_memzero(sk, sk_buffer_len);
    // convert encoded key to raw form
    int b64err = sodium_base642bin(sk, sk_buffer_len,
                                   (const char *) skb64, skb64_len,
                                   NULL, NULL, NULL, sodium_base64_VARIANT_ORIGINAL);
    if (b64err == -1) {
        *yield = string_copy((unsigned char *) "Error decoding private key");
        return ERROR;
    }

    // get and convert public key
    unsigned char *pkb64 = argv[1];
    size_t pkb64_len = strlen((const char *) pkb64);
    // reserve space for conversion
    size_t pk_buffer_len = crypto_box_PUBLICKEYBYTES; // pkb64_len / 4 * 3;
    unsigned char *pk = (unsigned char *) store_get_untainted(pk_buffer_len);
    sodium_memzero(pk, pk_buffer_len);
    // convert encoded key to raw form
    b64err = sodium_base642bin(pk, pk_buffer_len,
                               (const char *) pkb64, pkb64_len,
                               NULL, NULL, NULL, sodium_base64_VARIANT_ORIGINAL);
    if (b64err == -1) {
        *yield = string_copy((unsigned char *) "Error decoding public key");
        return ERROR;
    }

    // get encrypted message
    unsigned char *ciphertextb64 = argv[2];
    size_t ciphertextb64_len = strlen((const char *) ciphertextb64);

    // base64-decode the ciphertext
    size_t ciphertextbuflen = ciphertextb64_len / 4 * 3;
    unsigned char *ciphertext = (unsigned char *) store_get_untainted(ciphertextbuflen);
    size_t ciphertextlen;
    sodium_memzero(ciphertext, ciphertextbuflen);
    b64err = sodium_base642bin(ciphertext, ciphertextbuflen,
                               (const char *) ciphertextb64, ciphertextb64_len,
                               NULL, &ciphertextlen, NULL,
                               sodium_base64_VARIANT_ORIGINAL);
    if (b64err == -1) {
        *yield = string_copy((unsigned char *) "Error decoding base64 encoded ciphertext");
        return ERROR;
    }

    // prepare buffer for cleartext
    size_t cleartextlen = ciphertextlen - crypto_box_SEALBYTES;
    unsigned char *cleartext = (unsigned char *) store_get_untainted(cleartextlen + 1);
    sodium_memzero(cleartext, cleartextlen + 1);

    // decrypt message
    if (crypto_box_seal_open(cleartext, ciphertext, ciphertextlen, pk, sk) != 0) {
        *yield = string_copy((unsigned char *) "Decryption error after crypto_box_seal_open()");
        return ERROR;
    }

    // return cleartext
    *yield = string_copy(cleartext);
    return OK;
}