/*
This is part of WHY2
Copyright (C) 2022 Václav Šmejkal

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

#include <why2/chat/crypto.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/random.h>

#include <why2/memory.h>
#include <why2/misc.h>

#include <gmp.h>

#include <openssl/sha.h>

//DO NOT TOUCH THESE PLS :3
char *rsa_modulus = NULL; //THE RSA MODULUS
char *rsa_d = NULL; //THE RSA d

//LOCAL
void generate_prime(mpz_t x)
{
    //RANDOM
    gmp_randstate_t state;
    gmp_randinit_default(state);
    unsigned long random_buffer; //SEED

    do
    {
        if (getrandom(&random_buffer, sizeof(unsigned long), GRND_NONBLOCK) == -1) why2_die("getrandom fn failed!");

        //GENERATE RANDOM PRIME USING random_buffer SEED
        gmp_randseed_ui(state, random_buffer);
        mpz_urandomb(x, state, WHY2_CHAT_KEY_BITS);
        mpz_nextprime(x, x);
    } while (mpz_probab_prime_p(x, WHY2_CHAT_PRIME_ITERS) == 0); //CHECK FOR PRIME PROBABILITY

    //DEALLOCATION
    gmp_randclear(state);
}

void read_file(FILE *file, char **output)
{
    //VARIABLES
    int buffer_size;
    char *buffer;

    //GET LENGTH
    fseek(file, 0, SEEK_END);
    buffer_size = ftell(file);
    rewind(file);

    //READ
    buffer = why2_calloc(buffer_size + 1, sizeof(char));
    if (fread(buffer, buffer_size, 1, file) != 1) why2_die("Reading keyfile failed!");
    buffer[buffer_size] = '\0';

    //ASSIGN OUTPUT
    *output = buffer;
}

char *exp_mod(char *to_exp, char *exponent)
{
    //VARIABLES
    char *output;
    mpz_t m, c, n, e;
    mpz_init(c);

    //GET ALL STUFF
    mpz_init_set_str(m, to_exp, 10);
    mpz_init_set_str(n, why2_get_chat_modulus(), WHY2_CHAT_KEY_BASE);
    mpz_init_set_str(e, exponent, WHY2_CHAT_KEY_BASE);

    //ENCRYPT MESSAGE
    mpz_powm(c, m, e, n);

    output = why2_malloc(mpz_sizeinbase(c, 10) + 2); //ALLOCATE OUTPUT
    mpz_get_str(output, 10, c); //GET OUTPUT

    //DEALLOCATION
    mpz_clears(m, c, n, e, NULL);

    return output;
}

//GLOBAL
void why2_chat_init_keys(void)
{
    //KEY FILES
    FILE *public; //TECHNICALLY, PUBLIC KEY CONTAINS ONLY THE MODULUS AND PRIVATE CONTAINS ONLY THE d
    FILE *private;

    //GET PATH TO KEY DIR
    char *path = why2_replace(WHY2_CHAT_KEY_LOCATION, "{USER}", getenv("USER"));

    //ALLOCATE THE KEY PATHS
    char *public_path = why2_malloc(strlen(path) + strlen(WHY2_CHAT_PUB_KEY) + 3);
    char *private_path = why2_malloc(strlen(path) + strlen(WHY2_CHAT_PRI_KEY) + 3);

    //GET THE ACTUAL KEY PATHS
    sprintf(public_path, "%s/%s%c", path, WHY2_CHAT_PUB_KEY, '\0');
    sprintf(private_path, "%s/%s%c", path, WHY2_CHAT_PRI_KEY, '\0');

    //CHECK IF KEYS EXIST
    if (access(path, R_OK) != 0)
    {
        mkdir(path, 0700);

        //SOME USER OUTPUT
        printf("You are probably running WHY2-Chat for the first time now.\nGenerating RSA keys...\n");

        //VARIABLES
        mpz_t p, q, e, d, n, phi_n, buffer_1, buffer_2;
        mpz_inits(p, q, e, d, n, phi_n, buffer_1, buffer_2, NULL);

        //GENERATE PRIMES
        generate_prime(p);
        generate_prime(q);

        //SET e
        mpz_set_str(e, WHY2_CHAT_RSA_EXPONENT, 10);

        //GET n
        mpz_mul(n, p, q);

        //GET phi
        mpz_sub_ui(buffer_1, p, 1);
        mpz_sub_ui(buffer_2, q, 1);
        mpz_mul(phi_n, buffer_1, buffer_2);

        //COUNT d
        mpz_invert(d, e, phi_n);

        printf("Saving keys...\n");

        //WRITE THE KEYS INTO KEY-FILES
        public = why2_fopen(public_path, "w+");
        private = why2_fopen(private_path, "w+");

        mpz_out_str(public, WHY2_CHAT_KEY_BASE, n);
        mpz_out_str(private, WHY2_CHAT_KEY_BASE, d);

        //KEYGEN DEALLOCATION
        mpz_clears(p, q, e, d, n, phi_n, buffer_1, buffer_2, NULL);
    } else
    {
        //OPEN FILES
        public = why2_fopen(public_path, "r");
        private = why2_fopen(private_path, "r");

        //READ THE KEYS
        read_file(public, &rsa_modulus);
        read_file(private, &rsa_d);
    }

    //DEALLOCATION
    why2_deallocate(path);
    why2_deallocate(public_path);
    why2_deallocate(private_path);
    why2_deallocate(public);
    why2_deallocate(private);
}

void why2_chat_deallocate_keys(void)
{
    why2_deallocate(rsa_modulus);
    why2_deallocate(rsa_d);
}

char *why2_get_chat_modulus(void)
{
    return rsa_modulus;
}

char *why2_get_chat_d(void)
{
    return rsa_d;
}

char *why2_chat_rsa_pub_encrypt(char *to_encrypt)
{
    return exp_mod(to_encrypt, WHY2_CHAT_RSA_EXPONENT);
}

char *why2_chat_rsa_pri_decrypt(char *to_decrypt)
{
    return exp_mod(to_decrypt, why2_get_chat_d());
}

char *why2_sha256(char *input)
{
    unsigned char *output = why2_malloc(SHA256_DIGEST_LENGTH + 1);
    char *formatted_output = why2_malloc(SHA256_DIGEST_LENGTH * 2 + 2);

    SHA256((unsigned char*) input, strlen(input), output);

    //SAVE AS STRING IN HEX
    for (int i = 0; i < SHA256_DIGEST_LENGTH; i++)
    {
        sprintf(formatted_output + (i * 2), "%02x", output[i]);
    }
    formatted_output[SHA256_DIGEST_LENGTH * 2] = '\0';

    //DEALLOCATION
    why2_deallocate(output);

    return formatted_output;
}