/* Copyright (C) 2020-2022  Philipp "ph3-der-loewe" Schafft <lion@lion.leolix.org>
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library 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
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <stdio.h>
#include <string.h>
#include <math.h>
#include <pthread.h>

#ifdef HAVE_UNAME
#include <sys/utsname.h>
#endif

#ifdef HAVE_UNISTD_H
#include <unistd.h>
#endif

#ifdef HAVE_SYS_TYPES_H
#include <sys/types.h>
#endif
#ifdef HAVE_SYS_STAT_H
#include <sys/stat.h>
#endif
#ifdef HAVE_FCNTL_H
#include <fcntl.h>
#endif

#ifdef HAVE_SYS_RANDOM_H
#include <sys/random.h>
#endif

#include <igloo/prng.h>
#include <igloo/error.h>
#include <igloo/digest.h>
#include <igloo/ro.h>
#include <igloo/time.h>

#include "private.h"

#define HASHFUNC            "SHA3-512" // must be a 512 bit function!
#define BLOCK_LENGTH        (512/8)

#define MAX_BITS_PER_HASH   192
#define MAX_BITS_PER_STATE  (MAX_BITS_PER_HASH*3)

#define STATE_INITED        0x01U
#define STATE_AUTO_SEED_0   0x02U
#define STATE_MAN_SEED_0    0x04U

typedef struct {
    igloo_ctime_t ctime;
    pthread_t thread;
    const void *stack;
    const void *instance;
    const void *buffer;
    size_t bufferlen;
} igloo_prng_buffer_extra_t;

static inline void buffer_extra_init(igloo_prng_buffer_extra_t *extra, const void *instance)
{
    igloo_error_t error;

    memset(extra, 0, sizeof(*extra));
    extra->thread = pthread_self();
    extra->stack = &error; // any object on the current stack.
    extra->instance = instance;

    error = igloo_ctime_from_now(&(extra->ctime), igloo_CLOCK_MONOTONIC);
    if (error != igloo_ERROR_NONE)
        igloo_ctime_from_now(&(extra->ctime), igloo_CLOCK_REALTIME);
}

igloo_error_t   igloo_instance_prng_init(igloo_prng_state_t *state)
{
    if (pthread_mutex_init(&(state->lock), NULL) != 0)
        return igloo_ERROR_GENERIC;

    state->overcommitment = 32;

    return igloo_ERROR_NONE;
}

void            igloo_instance_prng_destroy(igloo_prng_state_t *state)
{
    pthread_mutex_destroy(&(state->lock));
}

igloo_error_t   igloo_instance_prng_stringify(igloo_prng_state_t *state, char *buf, size_t len)
{
    (void)state;


    pthread_mutex_lock(&(state->lock));
    snprintf(buf, len, "{bits=%zu}", state->bits);
    pthread_mutex_unlock(&(state->lock));

    return igloo_ERROR_NONE;
}

static igloo_error_t get_inited_state(igloo_prng_state_t **state, size_t *instancelen, igloo_ro_t instance)
{
    igloo_prng_state_t *self;
    size_t len;

    *state = self = igloo_instance_get_prng_state(instance, &len);
    if (!*state)
        return igloo_ERROR_BADSTATE;

    if (instancelen)
        *instancelen = len;

    pthread_mutex_lock(&(self->lock));

    if (!(self->state & STATE_INITED)) {
        igloo_digest_t *digest = NULL;
        igloo_error_t error;
        const void *instanceraw = igloo_ro_to_type(instance, igloo_ro_stub_t);
        struct {
            igloo_prng_buffer_extra_t extra;

            int debian;
#ifdef HAVE_UNAME
            struct utsname utsname;
#endif
#ifdef HAVE_GETUID
            uid_t uid;
#endif
#ifdef HAVE_GETPID
            pid_t pid;
#endif
#ifdef HAVE_GETPPID
            pid_t ppid;
#endif
        } seed;

        buffer_extra_init(&(seed.extra), instanceraw);
        seed.extra.buffer = instanceraw;
        seed.extra.bufferlen = len;

        memset(&seed, 0, sizeof(seed));

        seed.debian = 4;
#ifdef HAVE_UNAME
        uname(&seed.utsname);
#endif
#ifdef HAVE_GETUID
        seed.uid = getuid();
#endif
#ifdef HAVE_GETPID
        seed.pid = getpid();
#endif
#ifdef HAVE_GETPPID
        seed.ppid = getppid();
#endif


        error = igloo_digest_new(&digest, instance, HASHFUNC);
        if (error != igloo_ERROR_NONE)
            return error;

        if (igloo_digest_write(digest, instanceraw, len) != (ssize_t)len) {
            igloo_ro_unref(&digest);
            return igloo_ERROR_GENERIC;
        }

        if (igloo_digest_write(digest, &seed, sizeof(seed)) != (ssize_t)sizeof(seed)) {
            igloo_ro_unref(&digest);
            return igloo_ERROR_GENERIC;
        }


        if (igloo_digest_read(digest, self->initial, sizeof(self->initial)) != sizeof(self->initial)) {
            igloo_ro_unref(&digest);
            return igloo_ERROR_GENERIC;
        }

        error = igloo_ro_unref(&digest);
        if (error != igloo_ERROR_NONE)
            return error;

        self->bits = 5;

        self->state |= STATE_INITED;

#if defined(PATH_BOOT_ID) || defined(PATH_MACHINE_ID)
        // this is safe as we are inited now. In worst case we read initial seeds twice.
        pthread_mutex_unlock(&(self->lock));
#ifdef PATH_BOOT_ID
        igloo_prng_read_file(instance, PATH_BOOT_ID, -1, -1, igloo_PRNG_FLAG_NONE);
#endif
#ifdef PATH_MACHINE_ID
        igloo_prng_read_file(instance, PATH_MACHINE_ID, -1, 0, igloo_PRNG_FLAG_NONE);
#endif
        pthread_mutex_lock(&(self->lock));
#endif
    }

    return igloo_ERROR_NONE;
}

static ssize_t igloo_prng_estimate_bits(const void *buffer, size_t len)
{
    size_t counts[256];
    size_t i;
    double entpb = 0;

    if (!buffer)
        return -1;

    if (len < 2)
        return 0;

    memset(counts, 0, sizeof(counts));

    for (i = 0; i < len; i++)
        counts[((const unsigned char *)buffer)[i]]++;

    for (i = 0; i < 256; i++) {
        double probability;

        if (!counts[i])
            continue;

        probability = counts[i]/(double)len;
        entpb += probability * log2(1. / probability);
    }

    return entpb * len;
}

igloo_error_t igloo_prng_configure(igloo_ro_t instance, igloo_prng_flags_t addflags, igloo_prng_flags_t removeflags, ssize_t overcommitment)
{
    igloo_prng_state_t *self;

    self = igloo_instance_get_prng_state(instance, NULL);
    if (!self)
        return igloo_ERROR_BADSTATE;

    pthread_mutex_lock(&(self->lock));
    self->flags |= addflags|removeflags;
    self->flags -= removeflags;
    if (overcommitment > 0)
        self->overcommitment = overcommitment;
    pthread_mutex_unlock(&(self->lock));

    return igloo_ERROR_NONE;
}

static igloo_error_t digest_instance_extra_once(const void *instanceraw, size_t instancelen, igloo_ro_t instance, igloo_prng_buffer_extra_t *extra, const void *mid, size_t midlen, igloo_digest_512_t out)
{
    igloo_digest_t *digest = NULL;
    igloo_error_t error;

    error = igloo_digest_new(&digest, instance, HASHFUNC);
    if (error != igloo_ERROR_NONE)
        return error;

    if (igloo_digest_write(digest, instanceraw, instancelen) != (ssize_t)instancelen) {
        igloo_ro_unref(&digest);
        return igloo_ERROR_GENERIC;
    }

    if (mid && midlen) {
        if (igloo_digest_write(digest, mid, midlen) != (ssize_t)midlen) {
            igloo_ro_unref(&digest);
            return igloo_ERROR_GENERIC;
        }
    }

    if (igloo_digest_write(digest, extra, sizeof(*extra)) != (ssize_t)sizeof(*extra)) {
        igloo_ro_unref(&digest);
        return igloo_ERROR_GENERIC;
    }

    if (igloo_digest_read(digest, out, sizeof(igloo_digest_512_t)) != sizeof(igloo_digest_512_t)) {
        igloo_ro_unref(&digest);
        return igloo_ERROR_GENERIC;
    }

    error = igloo_ro_unref(&digest);
    if (error != igloo_ERROR_NONE)
        return error;

    return igloo_ERROR_NONE;
}

static inline size_t igloo_prng_auto_reseed_unlocked_urandom_buffer(void *buffer, size_t bufferlen)
{
#if defined(PATH_URANDOM) || defined(HAVE_GETRANDOM)
    ssize_t ret;
#endif
#ifdef PATH_URANDOM
    int fh;
#endif

#ifdef HAVE_GETRANDOM
    ret = getrandom(buffer, bufferlen, 0);
    if (ret > 0)
        return ret;
#endif

#ifdef HAVE_GETENTROPY
    if (getentropy(buffer, bufferlen) == 0) {
        return bufferlen;
    }
#endif

#ifdef PATH_URANDOM
#ifdef O_CLOEXEC
    fh = open(PATH_URANDOM, O_RDONLY|O_CLOEXEC, 0);
#else
    fh = open(PATH_URANDOM, O_RDONLY, 0);
#endif
    if (fh >= 0) {
        ret = read(fh, buffer, bufferlen);
        close(fh);

        if (ret > 0)
            return ret;
    }
#endif

    return 0;
}

igloo_error_t igloo_prng_auto_reseed_unlocked(igloo_prng_state_t *self, const void *instanceraw, size_t instancelen, igloo_ro_t instance)
{
    igloo_prng_buffer_extra_t extra;
    igloo_error_t error;
    char buffer[32];
    size_t extrabytes;

    buffer_extra_init(&extra, instanceraw);
    self->callcount++;

    extrabytes = igloo_prng_auto_reseed_unlocked_urandom_buffer(buffer, sizeof(buffer));

    if (self->state & STATE_AUTO_SEED_0) {
        error = digest_instance_extra_once(instanceraw, instancelen, instance, &extra, buffer, extrabytes, self->auto_seed0);
    } else {
        error = digest_instance_extra_once(instanceraw, instancelen, instance, &extra, buffer, extrabytes, self->auto_seed1);
    }
    if (error != igloo_ERROR_NONE) {
        pthread_mutex_unlock(&(self->lock));
        return error;
    }
    self->state ^= STATE_AUTO_SEED_0;

    self->bits += 3 + (extrabytes * 8 / 2);

    return igloo_ERROR_NONE;
}

igloo_error_t igloo_prng_auto_reseed(igloo_ro_t instance, igloo_prng_flags_t flags)
{
    igloo_prng_state_t *self;
    size_t instancelen;
    igloo_error_t error = get_inited_state(&self, &instancelen, instance);
    const void *instanceraw = igloo_ro_to_type(instance, igloo_ro_stub_t);

    (void)flags;

    if (error != igloo_ERROR_NONE)
        return error;

    error = igloo_prng_auto_reseed_unlocked(self, instanceraw, instancelen, instance);

    pthread_mutex_unlock(&(self->lock));

    return error;
}

igloo_error_t igloo_prng_write(igloo_ro_t instance, const void *buffer, size_t len, ssize_t bits, igloo_prng_flags_t flags)
{
    igloo_prng_state_t *self;
    size_t instancelen;
    igloo_error_t error;
    const void *instanceraw;
    igloo_prng_buffer_extra_t extra;

    (void)flags;

    if (bits > (ssize_t)(len*8))
        return igloo_ERROR_INVAL;

    if (bits < 0) {
        ssize_t ent = igloo_prng_estimate_bits(buffer, len);

        bits = (len * 8) / 32;
        if (ent >= 0 && bits > ent)
            bits = ent;
    }

    if (bits > MAX_BITS_PER_HASH)
        bits = MAX_BITS_PER_HASH;

    error = get_inited_state(&self, &instancelen, instance);
    if (error != igloo_ERROR_NONE)
        return error;

    instanceraw = igloo_ro_to_type(instance, igloo_ro_stub_t);

    buffer_extra_init(&extra, instanceraw);

    self->callcount++;

    if (self->state & STATE_MAN_SEED_0) {
        error = digest_instance_extra_once(instanceraw, instancelen, instance, &extra, buffer, len, self->manual_seed0);
    } else {
        error = digest_instance_extra_once(instanceraw, instancelen, instance, &extra, buffer, len, self->manual_seed1);
    }
    if (error != igloo_ERROR_NONE) {
        pthread_mutex_unlock(&(self->lock));
        return error;
    }
    self->state ^= STATE_MAN_SEED_0;
    self->bits += bits;
    self->bits += 1;

    if (self->bits > MAX_BITS_PER_STATE)
        self->bits = MAX_BITS_PER_STATE;

    pthread_mutex_unlock(&(self->lock));
    return igloo_ERROR_NONE;
}

static igloo_error_t igloo_prng_read__inner(igloo_prng_state_t *self, const void *instanceraw, size_t instancelen, igloo_ro_t instance, igloo_prng_buffer_extra_t *extra)
{
    self->callcount++;

    return digest_instance_extra_once(instanceraw, instancelen, instance, extra, NULL, 0, self->last_result);
}

ssize_t igloo_prng_read(igloo_ro_t instance, void *buffer, size_t len, igloo_prng_flags_t flags)
{
    igloo_prng_state_t *self;
    size_t instancelen;
    igloo_error_t error = get_inited_state(&self, &instancelen, instance);
    const void *instanceraw = igloo_ro_to_type(instance, igloo_ro_stub_t);
    size_t done = 0;
    size_t bitswant;
    igloo_prng_buffer_extra_t extra;

    (void)flags;

    if (error != igloo_ERROR_NONE)
        return -1;

    bitswant = ((len * 8) / self->overcommitment) + 1;

    if (bitswant > self->bits) {
        size_t try;
        for (try = 0; bitswant > self->bits && try < 2; try++) {
            error = igloo_prng_auto_reseed_unlocked(self, instanceraw, instancelen, instance);
            if (error != igloo_ERROR_NONE)
                break;
        }

        if (bitswant > self->bits) {
            pthread_mutex_unlock(&(self->lock));
            return -1;
        }
    }

    self->bits -= bitswant;

    buffer_extra_init(&extra, instanceraw);
    extra.buffer = buffer;
    extra.bufferlen = len;

    while (done < len) {
        size_t todo = (len - done) < sizeof(self->last_result) ? len - done : sizeof(self->last_result);

        error = igloo_prng_read__inner(self, instanceraw, instancelen, instance, &extra);
        if (error != igloo_ERROR_NONE) {
            pthread_mutex_unlock(&(self->lock));
            if (done > 0) {
                return done;
            } else {
                return -1;
            }
        }

        memcpy(buffer + done, self->last_result, todo);
        done += todo;
    }


    pthread_mutex_unlock(&(self->lock));
    return done;
}

/* write len pseudo random bytes to a file. If len is -1 a default value is used. */
igloo_error_t igloo_prng_write_file(igloo_ro_t instance, const char *filename, ssize_t len, igloo_prng_flags_t flags)
{
    char buffer[BLOCK_LENGTH*16];
    size_t done = 0;
    FILE *file;

    if (len < 0)
        len = 1024;

    file = fopen(filename, "wb");
    if (!file)
        return igloo_ERROR_GENERIC;

    while (done < (size_t)len) {
        size_t todo = (size_t)len < sizeof(buffer) ? (size_t)len : sizeof(buffer);
        ssize_t res = igloo_prng_read(instance, buffer, todo, flags);

        if (res < 1) {
            fclose(file);
            return igloo_ERROR_GENERIC;
        }

        if (fwrite(buffer, 1, res, file) != (size_t)res) {
            fclose(file);
            return igloo_ERROR_GENERIC;
        }

        done += res;
    }

    fclose(file);

    return igloo_ERROR_NONE;
}

/* read at max len bytes from the file and see the PRNG with it. if len is -1 all of the file is read. */
igloo_error_t igloo_prng_read_file(igloo_ro_t instance, const char *filename, ssize_t len, ssize_t bits, igloo_prng_flags_t flags)
{
    char buffer[BLOCK_LENGTH*16];
    size_t done = 0;
    int fh;
    igloo_error_t error = igloo_ERROR_NONE;


    if (len < 0 || len > 1048576)
        len = 1048576;

#ifdef O_CLOEXEC
    fh = open(filename, O_RDONLY|O_CLOEXEC, 0);
#else
    fh = open(filename, O_RDONLY, 0);
#endif
    if (fh < 0)
        return igloo_ERROR_GENERIC;

    while (done < (size_t)len) {
        size_t todo = (size_t)len < sizeof(buffer) ? (size_t)len : sizeof(buffer);
        size_t res = read(fh, buffer, todo);
        ssize_t bufferbits;

        if (res < 1) {
            close(fh);
            return igloo_ERROR_NONE;
        }

        if (bits < 0 || len < 1) {
            bufferbits = -1;
        } else {
            bufferbits = (bits * res) / len;
        }

        error = igloo_prng_write(instance, buffer, res, bufferbits, flags);
        if (error != igloo_ERROR_NONE)
            break;

        done += res;
    }

    close(fh);
    return error;
}
