#define _GNU_SOURCE

#include "encryption.h"
#include "shared.h"

#include <sodium.h>

#include <sys/types.h>
#include <sys/stat.h>
#include <stdarg.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <error.h>
#include <fcntl.h>
#include <stdio.h>

#if ___VXGG___ALWAYS_CHECK_LIBSODIUM___ > 0
#if ___VXGG___USE_CLS_CALLBACK___ > 0

void naclfaildefault(void *none) {
    none = none; // Makes gcc happy    
    if(___VXGG___VERBOSE_ERRORS___)
        error(1, ENOTSUP, "<naclfaildefault> Couldn't initialize sodium for some reason. Quitting...");
    exit(EXIT_FAILURE);
}

int checksodiumcb(const vxgg_naclfailcb callback, void *data, unsigned char set) {
    static vxgg_naclfailcb cb = naclfaildefault;
    static void *usr = NULL;
    int ret;

    if(set) {
        cb = callback;
        usr = data;
        return 2;   // libsodium normally returns 1 if the library is already initialized, so this is to signal that the callback has been updated
    }

    if((ret = sodium_init()) < 0) {
        if(cb == NULL)
            error(0, EINVAL, "<checksodiumcb> refusing to run a null callback");
        else
            cb(usr);
    }

    return ret;
}

void vxgg_setsodiumfailcb(vxgg_naclfailcb cb, void *data) {
    checksodiumcb(cb, data, 1);
}

#endif

void checksodium(void) {
    #if ___VXGG___USE_CLS_CALLBACK___ > 0
    checksodiumcb(NULL, NULL, 0);
    #else

    if(sodium_init() < 0) {
        errno = ENOTSUP;
        XALLOC_EXIT("<checksodium> Couldn't initialize sodium for some reason. Quitting...");
    }

    #endif
    
    return;
}

#endif

int maketmp(const char * const dest) {
    if(!dest)
        RETURNWERR(EINVAL, -1);
    return open(dest, (O_TMPFILE | O_WRONLY | O_CLOEXEC | O_SYNC), (S_IRUSR | S_IWUSR));
}

int linkto(const char * const target, int tgfd) {
    char *path = NULL;
    asprintf(&path, "/proc/self/fd/%d", tgfd);
    if(!path)
        ERROR(1, errno, "<linkto> Couldn't get path to move file into system",);
    remove(target); // Make sure an old version isn't sticking around (it's not catastrophic if this fails, but it should be noted or logged somewhere)
    return linkat(AT_FDCWD, path, AT_FDCWD, target, AT_SYMLINK_FOLLOW);
}

int encryptviatmp(const char * const target, const char * const output, const unsigned char key[crypto_secretstream_xchacha20poly1305_KEYBYTES]) {
    #if ___VXGG___ALWAYS_CHECK_LIBSODIUM___ > 0
    checksodium();
    #endif

    int fd = -1, tfd = -1;

    // Open the target file
    if((fd = open(target, O_RDONLY)) < 0)
        return -1;

    // Create a temp file for writing
    char *targetdir = xdirname(output);
    tfd = maketmp(targetdir);
    free(targetdir);
    if(tfd < 0) {
        close(fd);
        return -1;
    }

    FILE *src, *dst;
    if(!(src = fdopen(fd, "rb")))
        ERROR(1, errno, "<encryptviatmp> Couldn't open \"%s\"", , target);
    if(!(dst = fdopen(tfd, "wb")))
        ERROR(1, errno, "<encryptviatmp> Couldn't open \"%s\"", , output);
    if(encrypttofile(src, dst, key) < 0)
        ERROR(1, ENOTRECOVERABLE, "<encryptviatmp> I don't even have a way to cause an error here. How did you do it?",);

    // Link the temp file into the system
    if(linkto(output, tfd) < 0)
        ERROR(1, errno, "<encryptviatmp> Could not link \"%s\" into system after encryption", , output);

    fclose(dst);
    fclose(src);
    // fclose alco closes fd and tfd, as fdopen does not dup the file descriptors

    return 0;
}

int decryptto(const char * const encrypted, const char * const target, const unsigned char key[crypto_secretstream_xchacha20poly1305_KEYBYTES]) {
    #if ___VXGG___ALWAYS_CHECK_LIBSODIUM___ > 0
        checksodium();
    #endif

    if(!encrypted)
        RETURNWERR(EINVAL, -1);
    if(!target)
        RETURNWERR(EINVAL, -1);
    if(!key)
        RETURNWERR(EINVAL, -1);

    FILE *src, *dst;
    if(!(src = fopen(encrypted, "rb")))
        ERROR(1, errno, "<decryptto> Could not open \"%s\" for decryption", , encrypted);

    int fdst = maketmp(target);
    if(!fdst)
        ERROR(1, errno, "<decryptto> Could not get temp file for decryption", );
    if(!(dst = fdopen(fdst, "wb")))
        ERROR(1, errno, "<decryptto> Could not open \"%s\" for writing decrypted data", , target);
    
    if(decrypttofile(src, dst, key) < 0)
        ERROR(1, errno, "<decryptto> How did you even cause an error?",);

    // Link temp into system
    if(linkto(target, fdst) < 0)
        ERROR(1, errno, "<decryptto> Could not link \"%s\" into system", , target);

    fclose(dst);
    fclose(src);
    // fclose alco closes fd and tfd, as fdopen does not dup the file descriptors


    return 0;
}

int encrypttofile(FILE *src, FILE *dst, const unsigned char key[crypto_secretstream_xchacha20poly1305_KEYBYTES]) {
    unsigned char  buf[CHUNKSIZE], cbuf[CHUNKSIZE + crypto_secretstream_xchacha20poly1305_ABYTES];
    unsigned char  header[crypto_secretstream_xchacha20poly1305_HEADERBYTES];
    crypto_secretstream_xchacha20poly1305_state state;
    unsigned long long cbuflen;
    unsigned char tag;
    size_t bytesread;
    int eof;

    #if ___VXGG___ALWAYS_CHECK_LIBSODIUM___ > 0
        checksodium();
    #endif

    if(!src)
        RETURNWERR(EINVAL, -1);
    if(!dst)
        RETURNWERR(EINVAL, -1);
    if(!key)
        RETURNWERR(EINVAL, -1);

    // Write the header
    crypto_secretstream_xchacha20poly1305_init_push(&state, header, key);
    if(fwrite(header, 1, sizeof(header), dst) < sizeof(header))
        if(ferror(dst))
            ERROR(1, errno, "<encrypttofile> Could not write header",);

    // Encrypt each chunk
    do {
        if((bytesread = fread(buf, 1, sizeof(buf), src)) < sizeof(buf))
            if(ferror(src))
                ERROR(1, errno, "<encrypttofile> Could not read from source",);
        eof = feof(src);
        tag = eof ? crypto_secretstream_xchacha20poly1305_TAG_FINAL : 0;

        crypto_secretstream_xchacha20poly1305_push(&state, cbuf, &cbuflen, buf, bytesread, NULL, 0, tag);
        if(fwrite(cbuf, 1, (size_t)cbuflen, dst) < (size_t)cbuflen)
            if(ferror(dst))
                ERROR(1, errno, "<encrypttofile> Could not write to target",);
    } while (!eof);
    
    return 0;
}

int decrypttofile(FILE *src, FILE *dst, const unsigned char key[crypto_secretstream_xchacha20poly1305_KEYBYTES]) {
    unsigned char  cbuf[CHUNKSIZE + crypto_secretstream_xchacha20poly1305_ABYTES], buf[CHUNKSIZE];
    unsigned char  header[crypto_secretstream_xchacha20poly1305_HEADERBYTES];
    crypto_secretstream_xchacha20poly1305_state state;
    unsigned long long buflen;
    unsigned char  tag;
    size_t bytesread;
    int eof;
    
    #if ___VXGG___ALWAYS_CHECK_LIBSODIUM___ > 0
        checksodium();
    #endif

    if(!src)
        RETURNWERR(EINVAL, -1);
    if(!dst)
        RETURNWERR(EINVAL, -1);
    if(!key)
        RETURNWERR(EINVAL, -1);

    // Read the header
    if(fread(header, 1, sizeof(header), src) < sizeof(header))
        if(ferror(src))
            ERROR(1, errno, "<decrypttofile> Couldn't read header", );

    // Make sure the header isn't fuckey
    if(crypto_secretstream_xchacha20poly1305_init_pull(&state, header, key) != 0)
        ERROR(1, errno, "<decrypttofile> Incomplete header", );

    // Decrypt each chunk
    do {
        if((bytesread = fread(cbuf, 1, sizeof(cbuf), src)) < sizeof(cbuf))
            if(ferror(src))
                ERROR(1, errno, "<decrypttofile> Ran into problem reading for decryption", );
        eof = feof(src);

        if (crypto_secretstream_xchacha20poly1305_pull(&state, buf, &buflen, &tag, cbuf, bytesread, NULL, 0) != 0)
            ERROR(1, errno, "<decrypttofile> Corrupted chunk", );

        if(tag == crypto_secretstream_xchacha20poly1305_TAG_FINAL && !eof)
                ERROR(1, errno, "<decrypttofile> End of stream before end of file", );
        if(eof && tag != crypto_secretstream_xchacha20poly1305_TAG_FINAL)
                ERROR(1, errno, "<decrypttofile> End of file before end of stream", );
        
        fwrite(buf, 1, (size_t)buflen, dst);
    } while(! eof);

    return 0;
}

int genpassword(char **str, unsigned int words) {
    // Early returns
    if(words < 1) {return 0;}
    #if ___VXGG___ALWAYS_CHECK_LIBSODIUM___ > 0
        checksodium();
    #endif

    // Bootstrap the first word
    char *lstr = NULL, *tmp = NULL;
    if(asprintf(&lstr, "%s", PASSWORD_WORDS[randombytes_uniform(PASSWORD_WORDS_LEN)]) < 0)
        return -1;

    // Concat the rest of the words into the password (without leaking memory)
    int ret;
    for(unsigned int i = 1; i < words; i++) {    
        ret = asprintf(&tmp, "%s %s", lstr, PASSWORD_WORDS[randombytes_uniform(PASSWORD_WORDS_LEN)]);
        sodium_memzero(lstr, strlen(lstr) + 1);
        free(lstr);
        if(ret < 0)
            return -1;

        lstr = tmp;
    }

    *str = lstr;
    return words;
}

// sodium_malloc wrapper. Calls `error()` or `abort()` depnding on the value of `___VXGG___XALLOC_EXIT_ON_ERROR___`. Will make sure libsodium is initialized if `___VXGG___ALWAYS_CHECK_LIBSODIUM___ > 0`
void* xsodium_malloc(size_t size) {
    #if ___VXGG___ALWAYS_CHECK_LIBSODIUM___ > 0
        checksodium();
    #endif

    void *mem = sodium_malloc(size);
    if(mem == NULL)
        XALLOC_EXIT("<xsodium_malloc> could not allocate memory... Quitting", );

    return mem;
}


/*
int main(void) {
    // Example code for creating a temp file, writing to it, then linking it back into the fs
    const char *dir = ".", *testmsg = "we do a little testing\n";
    char *path = NULL;

    int fd = maketmp(dir);
    if(fd < 0)
        error(1, errno, "Couldn't make temp file at %s", dir);

    if(write(fd, testmsg, strlen(testmsg)) < 0)
        error(1, errno, "write broke");
    
    asprintf(&path, "/proc/self/fd/%d", fd);
    linkat(AT_FDCWD, path, AT_FDCWD, "./test", AT_SYMLINK_FOLLOW);
    free(path);

    // Apparently, I don't have the CAP_DAC_READ_SEARCH capibility. Thanks for the solution, linux man pages

    if(close(fd) < 0)
        error(1, errno, "close broke");
    //*///

    /*// Example code for getting a password using genpassword
    checksodium();

    char *password = NULL;
    genpassword(&password, 20);
    printf("%s\n", (password != NULL) ? password : "Couldn't get a password");
    free(password);
    //*///

    /*// Example code for generating a password, derriving a secret key from it, and storing things properly
    
    // Initialization
    checksodium();
    char *pass = NULL, hpass[crypto_pwhash_STRBYTES];

    if(genpassword(&pass, 20) < 0) {
        error(1, 0, "Could not generate password, quitting...");
        abort(); // Makes gcc happy. Not sure why gcc randomly decides that error() isn't a proper exit, but hey whatever
    }
    sodium_mlock(pass, strlen(pass) + 1);
    printf("Password:%s\n", pass);

    // Store the password
    if(crypto_pwhash_str(hpass, pass, strlen(pass) + 1, crypto_pwhash_OPSLIMIT_MODERATE, crypto_pwhash_MEMLIMIT_MODERATE) != 0)
        error(1, errno, "Couldn't generate password, quitting...");
        // Don't know if I want to use MODERATE or SENSITIVE for this. SENSITIVE takes a little bit on my laptop, which honestly 
        // shouldn't be a problem, but it annoys me. MODERATE is quick and snappy, or at least quick enough that the slowdown is
        // barely noticable. I might do MODERATE for testing and SENSITIVE for release

    sodium_munlock(pass, strlen(pass) + 1);
    free(pass);

    printf("Hashed password: %s\n", hpass);

    // Check if the password from the user is correct
    char *uin = NULL; int size = -1;
    if((size = readwholebuffer(&uin, 1, STDIN_FILENO)) < 0)
        error(1, errno, "Could not read from stdin");
    sodium_mlock(uin, size);

    printf("Valid password? %s\n", (crypto_pwhash_str_verify(hpass, uin, size) == 0) ? "True" : "False");


    sodium_munlock(uin, strlen(uin) + 1);
    free(uin);

    return 0;
}
*/