#include "session-crypto.h"

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

#include <sodium.h>

#define LOG_MODULE "session-crypto"
#define LOG_ENABLE_DBG 0
#include "log.h"
#include "xmalloc.h"

#define MAGIC "FOOT-ENC1\0"
#define MAGIC_LEN 10

static bool sodium_ready = false;

bool
session_crypto_init(void)
{
    if (sodium_ready)
        return true;
    if (sodium_init() < 0) {
        LOG_ERR("libsodium init failed");
        return false;
    }
    sodium_ready = true;
    return true;
}

static bool
derive_key(const char *password, const unsigned char *salt,
           unsigned char key[crypto_aead_xchacha20poly1305_ietf_KEYBYTES])
{
    return crypto_pwhash(
        key, crypto_aead_xchacha20poly1305_ietf_KEYBYTES,
        password, strlen(password), salt,
        crypto_pwhash_OPSLIMIT_INTERACTIVE,
        crypto_pwhash_MEMLIMIT_INTERACTIVE,
        crypto_pwhash_ALG_ARGON2ID13) == 0;
}

bool
session_crypto_encrypt(const char *password,
                       const unsigned char *plaintext, size_t plaintext_len,
                       unsigned char **out, size_t *out_len)
{
    if (!session_crypto_init())
        return false;
    *out = NULL;
    *out_len = 0;

    const size_t salt_len = crypto_pwhash_SALTBYTES;
    const size_t nonce_len = crypto_aead_xchacha20poly1305_ietf_NPUBBYTES;
    const size_t tag_len = crypto_aead_xchacha20poly1305_ietf_ABYTES;

    const size_t total = MAGIC_LEN + salt_len + nonce_len + plaintext_len + tag_len;
    unsigned char *buf = xmalloc(total);
    memcpy(buf, MAGIC, MAGIC_LEN);

    unsigned char *salt = buf + MAGIC_LEN;
    unsigned char *nonce = salt + salt_len;
    unsigned char *cipher = nonce + nonce_len;

    randombytes_buf(salt, salt_len);
    randombytes_buf(nonce, nonce_len);

    unsigned char key[crypto_aead_xchacha20poly1305_ietf_KEYBYTES];
    if (!derive_key(password, salt, key)) {
        LOG_ERR("crypto_pwhash failed (out of memory?)");
        sodium_memzero(key, sizeof(key));
        free(buf);
        return false;
    }

    /* Bind the file header into the AEAD so any tampering breaks decryption. */
    unsigned long long clen = 0;
    int rc = crypto_aead_xchacha20poly1305_ietf_encrypt(
        cipher, &clen,
        plaintext, plaintext_len,
        buf, MAGIC_LEN + salt_len + nonce_len,  /* AAD = full header */
        NULL,
        nonce,
        key);
    sodium_memzero(key, sizeof(key));

    if (rc != 0) {
        LOG_ERR("xchacha20poly1305 encrypt failed");
        free(buf);
        return false;
    }
    /* Sanity */
    if (clen != plaintext_len + tag_len) {
        free(buf);
        return false;
    }

    *out = buf;
    *out_len = total;
    return true;
}

bool
session_crypto_decrypt(const char *password,
                       const unsigned char *blob, size_t blob_len,
                       unsigned char **out, size_t *out_len)
{
    if (!session_crypto_init())
        return false;
    *out = NULL;
    *out_len = 0;

    const size_t salt_len = crypto_pwhash_SALTBYTES;
    const size_t nonce_len = crypto_aead_xchacha20poly1305_ietf_NPUBBYTES;
    const size_t tag_len = crypto_aead_xchacha20poly1305_ietf_ABYTES;
    const size_t header_len = MAGIC_LEN + salt_len + nonce_len;

    if (blob_len < header_len + tag_len) {
        LOG_ERR("encrypted blob too small (%zu bytes)", blob_len);
        return false;
    }
    if (memcmp(blob, MAGIC, MAGIC_LEN) != 0) {
        LOG_ERR("bad magic in encrypted scrollback");
        return false;
    }

    const unsigned char *salt = blob + MAGIC_LEN;
    const unsigned char *nonce = salt + salt_len;
    const unsigned char *cipher = nonce + nonce_len;
    const size_t cipher_len = blob_len - header_len;

    unsigned char key[crypto_aead_xchacha20poly1305_ietf_KEYBYTES];
    if (!derive_key(password, salt, key)) {
        LOG_ERR("crypto_pwhash failed during decrypt");
        sodium_memzero(key, sizeof(key));
        return false;
    }

    unsigned char *plain = xmalloc(cipher_len);  /* upper bound */
    unsigned long long plen = 0;
    int rc = crypto_aead_xchacha20poly1305_ietf_decrypt(
        plain, &plen,
        NULL,
        cipher, cipher_len,
        blob, header_len,  /* AAD must match what was used in encrypt */
        nonce,
        key);
    sodium_memzero(key, sizeof(key));

    if (rc != 0) {
        LOG_ERR("decryption failed (wrong password or corrupted file)");
        sodium_memzero(plain, cipher_len);
        free(plain);
        return false;
    }

    *out = plain;
    *out_len = (size_t)plen;
    return true;
}