#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 defined ___VXGG___ALWAYS_CHECK_LIBSODIUM___ && ___VXGG___ALWAYS_CHECK_LIBSODIUM___ > 0
void naclfaildefault(void *none) {
    none = none; // Makes gcc happy    
    error(1, ENOTSUP, "Couldn't initialize sodium for some reason. Quitting...");
}

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

    if(callback != NULL) {
        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
    }

    int ret = sodium_init();
    if(ret < 0)
        cb(usr);

    return ret;
}

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

void checksodium(void) {
    #if defined ___VXGG___ALWAYS_CHECK_LIBSODIUM___ && ___VXGG___ALWAYS_CHECK_LIBSODIUM___ > 0
        checksodiumcb(NULL, NULL);
    #else
        if(sodium_init() < 0)
            error(1, ENOTSUP, "Couldn't initialize sodium for some reason. Quitting...");
    #endif
    
    return;
}

// To encrypt:
// 1- Create a temp file with the correct name in the root folder of the partition being encrypted      -- 
    // 1.1- Detect the partition and find the root folder                                                   -- DONE || NOT NECESSARY
    // 1.2- Create the temp file                                                                            -- DONE
// 2- Encrypt the file's contents to the temp file                                                      -- 
    // 2.1- Open the file                                                                                   -- 
    // 2.2- Stream the file's contents into some encryption algo                                            -- 
        // 2.2.1- Pick which encryption algo to use                                                              -- 
        // 2.2.2- Generate a key                                                                                 -- 
            // 2.2.2.1- Create a password to derrive a key from                                                     -- DONE 
    // 2.3- Pipe the output of the encryption into the temp file                                            -- 
// 3- Once the file has been encrypted, hard link it back to the original location, with the right name -- 
// 4- Delete the original file                                                                          -- 
// 5- Delete the temp file                                                                              --


int maketmp(const char * const dest) {
    return open(dest, (O_TMPFILE | O_WRONLY | O_CLOEXEC | O_SYNC), (S_IRUSR | S_IWUSR));
}

int encrypttotmp(const char * const toencrypt) {
    #if defined ___VXGG___ALWAYS_CHECK_LIBSODIUM___ && ___VXGG___ALWAYS_CHECK_LIBSODIUM___ > 0
        checksodium();
    #endif

    struct stat esb;
    int efd = -1;

    // Make sure the file is real and an actual file that can be encrypted
    if(stat(toencrypt, &esb) < 0)
        return -1;
    if(!S_ISREG(esb.st_mode))
        return -2;

    // Open the file as read-only
    if((efd = open(toencrypt, O_RDONLY)) < 0)
        return -3;

    // Need to get a secret key from a password and then set up cryptostream from libsodium

    return 0;
}

int genpassword(char **str, unsigned int words) {
    // Early returns
    if(words < 1)
        return 0;
    #if defined ___VXGG___ALWAYS_CHECK_LIBSODIUM___ && ___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 defined ___VXGG___ALWAYS_CHECK_LIBSODIUM___ && ___VXGG___ALWAYS_CHECK_LIBSODIUM___ > 0
        checksodium();
    #endif

    void *mem = sodium_malloc(size);
    if(mem == NULL) {
        #if defined ___VXGG___XALLOC_EXIT_ON_ERROR___ && ___VXGG___XALLOC_EXIT_ON_ERROR___ > 0
            error(1, errno, "xsodium_malloc: could not allocate memory... Quitting");
        #endif

        abort();
    }

    return mem;
}

#define TESTING
#ifdef TESTING

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);

    //*///

    //*// Example code for generating a key from a password and encrypting a test file

    const char *dir = ".", *fname = "toBeEncrypted.test.txt", *pass = "this is a password";
    char *path = NULL, *message = NULL, *efname = NULL;
    int fd = -1;

    // Message is just going to be a bunch of words from the dictionary
    checksodium();
    if(genpassword(&message, 1000) < 0)
        error(1, errno, "Could not generate message to be encrypted");
    if(!message)
        abort();

    // Get the temp file and write the message into it
    if((fd = maketmp(dir)) < 0)
        error(1, errno, "Could not make temp file to write to");

    // Write to the file and link it into the system
    remove(fname); // Make sure the new file can be linked into
    write(fd, message, strlen(message));
    asprintf(&path, "/proc/self/fd/%d", fd);
    if(!path)
        abort();
    linkat(AT_FDCWD, path, AT_FDCWD, fname, AT_SYMLINK_FOLLOW);

    free(path);
    free(message);
    close(fd);


    // Time for encryption
    unsigned char header[crypto_secretstream_xchacha20poly1305_HEADERBYTES];
    unsigned char key[crypto_secretstream_xchacha20poly1305_KEYBYTES];
    crypto_secretstream_xchacha20poly1305_state state;
    unsigned char salt[crypto_pwhash_SALTBYTES];
    
    // Generate a salt for derriving the key
    randombytes_buf(salt, sizeof(salt));

    // Derrive the key from the password
    if(crypto_pwhash(key, sizeof(key), pass, strlen(pass), salt, crypto_pwhash_OPSLIMIT_MODERATE, crypto_pwhash_MEMLIMIT_MODERATE, crypto_pwhash_ALG_DEFAULT) != 0)
        error(1, ENOMEM, "Not enough memory to generate a key");

    // Initialize the encryption stream's state
    crypto_secretstream_xchacha20poly1305_init_push(&state, header, key);
    

    // Deal with getting the file
    if((fd = open(fname, O_RDONLY)) < 0)
        error(1, errno, "Could not open test file to encrypt");
    int tfd = -1;
    if((tfd = maketmp(dir)) < 0) // Only because linking temp files is being annoying
        error(1, errno, "Could not open temp file for encryption");
    
    struct stat sb;
    if(fstat(fd, &sb) < 0)
        error(1, errno, "stat() error");

    // Read chunks of the file, encrypt them, then write them into the tmp file
    ssize_t bytesread = -1;
    unsigned char *buf = xcalloc(sb.st_blksize + 1, sizeof(*buf));
    unsigned char *cbuf = xcalloc((sb.st_blksize + 1) + crypto_secretstream_xchacha20poly1305_ABYTES, sizeof(*buf));
    while((bytesread = read(fd, buf, sb.st_blksize)) >= 0) {
        crypto_secretstream_xchacha20poly1305_push(&state, cbuf, NULL, buf, bytesread, NULL, 0, (bytesread > 0) ? 0 : crypto_secretstream_xchacha20poly1305_TAG_FINAL);
        if(writewholebuffer(tfd, cbuf, bytesread) < 0)
            error(1, errno, "write() error");

        if(bytesread == 0)
            break;
    }
    if(bytesread < 0)
        error(1, errno, "read() error");

    close(fd);

    asprintf(&efname, "%s.enc", fname);
    asprintf(&path, "/proc/self/fd/%d", tfd);
    if(!path || !efname)
        abort();
    remove(efname); // Make sure an old version isn't sticking around
    linkat(AT_FDCWD, path, AT_FDCWD, efname, AT_SYMLINK_FOLLOW);

    close(tfd);
    free(path);
    free(efname);

    //*///

    return 0;
}

#endif