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

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

/* Exim4 dlfunc API header */
#include <local_scan.h>

char * string2hex(unsigned char * input, size_t length) {
    const int growth = 3;
    char * outstring = store_get(growth*length+1);
    memset(outstring, 0, 3*length+1);
    for (int i =0; i<length; i++) {
        sprintf(outstring+i*growth, "%02x ", input[i]);
    }
    return outstring;
}

/*
 * 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 operations 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
    unsigned int cipherlen = messagelen + crypto_secretbox_MACBYTES;
    unsigned char *ciphertext = (unsigned char *) store_get(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(combined_message_len);
    memcpy(combined_message, nonce, crypto_secretbox_NONCEBYTES);
    memcpy(&combined_message[crypto_secretbox_NONCEBYTES], ciphertext, cipherlen);

    // base64-encode the ciphertext
    unsigned int outputsize = sodium_base64_ENCODED_LEN(combined_message_len,
                                                        sodium_base64_VARIANT_ORIGINAL);
    unsigned char *outstring = (unsigned char *) store_get(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(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
    unsigned int cleartextlen = combined_message_len - crypto_secretbox_NONCEBYTES - crypto_secretbox_MACBYTES;
    unsigned char *cleartext = (unsigned char *) store_get(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
    unsigned int pk_buffer_len = crypto_box_PUBLICKEYBYTES; // pkb64_len / 4 * 3 + 1;
    unsigned char *pk = (unsigned char *) store_get(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;
    }

    log_write(0, LOG_MAIN, "[encrypt] PK: %s", string2hex(pk, pk_buffer_len));

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

    log_write(0, LOG_MAIN, "[encrypt] cleartext: %s", message);

    // prepare buffer for ciphertext
    unsigned int cipherlen = messagelen + crypto_box_SEALBYTES;
    unsigned char *ciphertext = store_get(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;
    }

    log_write(0, LOG_MAIN, "[encrypt] ciphertext: %s", string2hex(ciphertext, cipherlen));

    // base64-encode the ciphertext
    unsigned int outputsize = sodium_base64_ENCODED_LEN(cipherlen,
                                                        sodium_base64_VARIANT_ORIGINAL);
    unsigned char *outstring = store_get(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
    unsigned int sk_buffer_len = crypto_box_SECRETKEYBYTES;// skb64_len / 4 * 3;
    unsigned char *sk = (unsigned char *) store_get(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;
    }
    log_write(0, LOG_MAIN, "[decrypt] SK: %s", string2hex(sk, sk_buffer_len));

    // get and convert public key
    unsigned char *pkb64 = argv[1];
    size_t pkb64_len = strlen((const char *) pkb64);
    // reserve space for conversion
    unsigned int pk_buffer_len = crypto_box_PUBLICKEYBYTES; // pkb64_len / 4 * 3;
    unsigned char *pk = (unsigned char *) store_get(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;
    }

    log_write(0, LOG_MAIN, "[decrypt] PK: %s", string2hex(pk, pk_buffer_len));

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

    // base64-decode the ciphertext
    unsigned int ciphertextbuflen = ciphertextb64_len / 4 * 3;
    unsigned char *ciphertext = (unsigned char *) store_get(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;
    }

    log_write(0, LOG_MAIN, "[encrypt] ciphertext: %s", string2hex(ciphertext, ciphertextlen));

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

    // decrypt message
//#define DEBUG
#ifdef DEBUG
    log_write(0, LOG_MAIN, "pid: %d", getpid()); int busywait = 0; while (busywait == 0) {}
#endif
    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;
}