#include "secret.h"

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

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

#include "util.h"

#define SECRET_PREFIX "enc:v1:"
#define SECRET_KEY_LEN 32
#define SECRET_NONCE_LEN 12
#define SECRET_TAG_LEN 16

static int hex_val(int c)
{
    if (c >= '0' && c <= '9')
        return c - '0';
    if (c >= 'a' && c <= 'f')
        return c - 'a' + 10;
    if (c >= 'A' && c <= 'F')
        return c - 'A' + 10;
    return -1;
}

static int secret_key(unsigned char out[SECRET_KEY_LEN])
{
    const char *s = getenv("BOKFD_SECRET_KEY");
    if (!s || !*s)
        return -1;
    size_t n = strlen(s);
    if (n == 2 * SECRET_KEY_LEN) {
        int hex = 1;
        for (size_t i = 0; i < n && hex; i++)
            if (hex_val((unsigned char)s[i]) < 0)
                hex = 0;
        if (hex) {
            for (size_t i = 0; i < SECRET_KEY_LEN; i++)
                out[i] = (unsigned char)((hex_val(s[2 * i]) << 4) |
                                         hex_val(s[2 * i + 1]));
            return 0;
        }
    }
    if (n % 4 == 1 || n > 172)
        return -1;
    char padded[176];
    memcpy(padded, s, n);
    while (n % 4)
        padded[n++] = '=';
    padded[n] = '\0';
    unsigned char *raw = NULL;
    size_t raw_n = 0;
    if (util_b64_decode(padded, n, &raw, &raw_n) != 0)
        return -1;
    int ok = raw_n == SECRET_KEY_LEN;
    if (ok)
        memcpy(out, raw, SECRET_KEY_LEN);
    free(raw);
    return ok ? 0 : -1;
}

int secret_available(void)
{
    unsigned char key[SECRET_KEY_LEN] = { 0 };
    int ok = secret_key(key) == 0;
    OPENSSL_cleanse(key, sizeof key);
    return ok;
}

int secret_encrypt(const char *plain, char **out)
{
    if (out)
        *out = NULL;
    if (!plain || !out)
        return -1;
    size_t plen = strlen(plain);
    if (plen > (size_t)INT_MAX)
        return -1;
    unsigned char key[SECRET_KEY_LEN], nonce[SECRET_NONCE_LEN];
    if (secret_key(key) != 0)
        return -1;
    if (RAND_bytes(nonce, SECRET_NONCE_LEN) != 1) {
        OPENSSL_cleanse(key, sizeof key);
        return -1;
    }
    size_t clen = plen + SECRET_TAG_LEN;
    unsigned char *ct = xmalloc(clen);
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
    int len = 0, total = 0, ok = 0;
    if (ctx &&
        EVP_EncryptInit_ex(ctx, EVP_aes_256_gcm(), NULL, NULL, NULL) == 1 &&
        EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, SECRET_NONCE_LEN,
                            NULL) == 1 &&
        EVP_EncryptInit_ex(ctx, NULL, NULL, key, nonce) == 1 &&
        EVP_EncryptUpdate(ctx, ct, &len, (const unsigned char *)plain,
                          (int)plen) == 1) {
        total = len;
        ok = EVP_EncryptFinal_ex(ctx, ct + total, &len) == 1;
        total += len;
    }
    if (ok)
        ok = EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, SECRET_TAG_LEN,
                                 ct + total) == 1;
    EVP_CIPHER_CTX_free(ctx);
    OPENSSL_cleanse(key, sizeof key);
    if (!ok || total != (int)plen) {
        free(ct);
        return -1;
    }
    total += SECRET_TAG_LEN;
    char *nonce_b64 = util_b64(nonce, sizeof nonce);
    char *ct_b64 = util_b64(ct, (size_t)total);
    free(ct);
    size_t olen =
        strlen(SECRET_PREFIX) + strlen(nonce_b64) + 1 + strlen(ct_b64) + 1;
    char *s = xmalloc(olen);
    snprintf(s, olen, "%s%s:%s", SECRET_PREFIX, nonce_b64, ct_b64);
    free(nonce_b64);
    free(ct_b64);
    *out = s;
    return 0;
}

int secret_decrypt(const char *stored, char **out)
{
    if (out)
        *out = NULL;
    if (!stored || !out)
        return -1;
    if (strncmp(stored, SECRET_PREFIX, strlen(SECRET_PREFIX)) != 0)
        return -1;
    const char *p = stored + strlen(SECRET_PREFIX);
    const char *sep = strchr(p, ':');
    if (!sep || sep == p || !sep[1])
        return -1;
    unsigned char key[SECRET_KEY_LEN];
    if (secret_key(key) != 0)
        return -1;
    unsigned char *nonce = NULL, *ct = NULL;
    size_t nonce_n = 0, ct_n = 0;
    if (util_b64_decode(p, (size_t)(sep - p), &nonce, &nonce_n) != 0) {
        OPENSSL_cleanse(key, sizeof key);
        return -1;
    }
    if (util_b64_decode(sep + 1, strlen(sep + 1), &ct, &ct_n) != 0) {
        free(nonce);
        OPENSSL_cleanse(key, sizeof key);
        return -1;
    }
    if (nonce_n != SECRET_NONCE_LEN || ct_n < SECRET_TAG_LEN ||
        ct_n - SECRET_TAG_LEN > (size_t)INT_MAX) {
        free(nonce);
        free(ct);
        OPENSSL_cleanse(key, sizeof key);
        return -1;
    }
    size_t plen = ct_n - SECRET_TAG_LEN;
    unsigned char *pt = xmalloc(plen + 1);
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
    int len = 0, total = 0, ok = 0;
    if (ctx &&
        EVP_DecryptInit_ex(ctx, EVP_aes_256_gcm(), NULL, NULL, NULL) == 1 &&
        EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, SECRET_NONCE_LEN,
                            NULL) == 1 &&
        EVP_DecryptInit_ex(ctx, NULL, NULL, key, nonce) == 1 &&
        EVP_DecryptUpdate(ctx, pt, &len, ct, (int)plen) == 1) {
        total = len;
        if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, SECRET_TAG_LEN,
                                ct + plen) == 1)
            ok = EVP_DecryptFinal_ex(ctx, pt + total, &len) == 1;
        total += len;
    }
    EVP_CIPHER_CTX_free(ctx);
    OPENSSL_cleanse(key, sizeof key);
    free(nonce);
    free(ct);
    if (!ok || total != (int)plen) {
        OPENSSL_cleanse(pt, plen + 1);
        free(pt);
        return -1;
    }
    pt[total] = '\0';
    *out = (char *)pt;
    return 0;
}

int secret_key_is_secret(const char *key)
{
    return key && strcmp(key, "smtp_password") == 0;
}