/*
 * The SHA3 implementation is based on rhash:
 * 2013 by Aleksey Kravchenko <rhash.admin@gmail.com>
 */

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

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

#include <rhash.h>

#include <igloo/digest.h>
#include <igloo/error.h>
#include <igloo/ro.h>
#include <igloo/sp.h>

#define DIGEST_ALGO_SHA3_224 "SHA3-224"
#define DIGEST_ALGO_SHA3_256 "SHA3-256"
#define DIGEST_ALGO_SHA3_384 "SHA3-384"
#define DIGEST_ALGO_SHA3_512 "SHA3-512"

struct igloo_digest_tag {
    igloo_ro_full_t __parent;

    /* metadata */
    const char *algo;

    /* state */
    int done;
    union {
        struct {
            unsigned int hash_id;
            rhash rhash;
        } rhash;
    } state;
};

static void __digest_free(igloo_ro_t self);

igloo_RO_PUBLIC_TYPE(igloo_digest_t, igloo_ro_full_t,
        igloo_RO_TYPEDECL_FREE(__digest_free)
        );

struct igloo_hmac_tag {
    igloo_ro_full_t __parent;

    const char *algo;
    char *key;
    size_t key_len;
    igloo_digest_t * inner;
};

static void __hmac_free(igloo_ro_t self)
{
    igloo_hmac_t *hmac = igloo_ro_to_type(self, igloo_hmac_t);
    free(hmac->key);
    igloo_ro_unref(&(hmac->inner));
}

igloo_RO_PUBLIC_TYPE(igloo_hmac_t, igloo_ro_full_t,
        igloo_RO_TYPEDECL_FREE(__hmac_free)
        );

static igloo_error_t igloo_digest_name_to_rhash_id (unsigned int *id, const char *algo)
{
    if (strcmp(algo, DIGEST_ALGO_SHA3_224) == 0) {
        *id = RHASH_SHA3_224;
    } else if (strcmp(algo, DIGEST_ALGO_SHA3_256) == 0) {
        *id = RHASH_SHA3_256;
    } else if (strcmp(algo, DIGEST_ALGO_SHA3_384) == 0) {
        *id = RHASH_SHA3_384;
    } else if (strcmp(algo, DIGEST_ALGO_SHA3_512) == 0) {
        *id = RHASH_SHA3_512;
    } else {
        return igloo_ERROR_NOSYS;
    }

    return igloo_ERROR_NONE;
}

static void __digest_free(igloo_ro_t self)
{
    igloo_error_t err;

    igloo_digest_t *digest = igloo_ro_to_type(self, igloo_digest_t);

    err = igloo_sp_unref(&(digest->algo), self);
    (void)err; /* make compiler happy */

    if (digest->state.rhash.rhash) {
        rhash_free(digest->state.rhash.rhash);
        digest->state.rhash.rhash = NULL;
    }
}

igloo_error_t igloo_digest_new(igloo_digest_t **self, igloo_ro_t group, const char *algo)
{
    igloo_error_t err;
    unsigned int rhash_id;

    err = igloo_ro_new_raw(self, igloo_digest_t, group);
    if (err != igloo_ERROR_NONE) {
        return err;
    }

    /* strcmp is not defined for NULL */
    if (!algo) {
        igloo_ro_unref(self);
        return igloo_ERROR_FAULT;
    }

    err = igloo_digest_name_to_rhash_id(&rhash_id, algo);
    if (err != igloo_ERROR_NONE) {
        igloo_ro_unref(self);
        return err;
    }

    (*self)->state.rhash.rhash = rhash_init(rhash_id);
    (*self)->state.rhash.hash_id = rhash_id;

    if ((*self)->state.rhash.rhash == NULL) {
        igloo_ro_unref(self);
        return igloo_ERROR_GENERIC;
    }

    err = igloo_sp_replace(algo, &((*self)->algo), group);
    if (err != igloo_ERROR_NONE) {
        igloo_ro_unref(self);
        return err;
    }

    return igloo_ERROR_NONE;
}

igloo_error_t igloo_digest_copy(igloo_digest_t **new, igloo_ro_t group, igloo_digest_t *old)
{
    (void)new, (void)group;

    if (!old || !*new)
        return igloo_ERROR_FAULT;

    return igloo_ERROR_NOSYS;
}

ssize_t igloo_digest_write(igloo_digest_t *digest, const void *data, size_t len)
{
    if (!digest || digest->done || !data)
        return -1;

    rhash_update(digest->state.rhash.rhash, data, len);
    return len;
}

ssize_t igloo_digest_read(igloo_digest_t *digest, void *buf, size_t len)
{
    int digestlen;

    if (!digest || digest->done || !buf)
        return -1;

    digestlen = rhash_get_digest_size(digest->state.rhash.hash_id);
    if ((ssize_t)len < digestlen)
        return -1;
    len = digestlen;

    digest->done = 1;

    rhash_final(digest->state.rhash.rhash, buf);

    return len;
}

static ssize_t __digest_algo_blocklength(const char *algo)
{
    /* strcmp is not defined for NULL */
    if (!algo) {
        return -1;
    }

    // 1600bit - 2 * digestlen[bit]:
    if (strcmp(algo, DIGEST_ALGO_SHA3_224) == 0) return 144;
    if (strcmp(algo, DIGEST_ALGO_SHA3_256) == 0) return 136;
    if (strcmp(algo, DIGEST_ALGO_SHA3_384) == 0) return 104;
    if (strcmp(algo, DIGEST_ALGO_SHA3_512) == 0) return 72;

    return -1; // invalid algorithm
}

static ssize_t __digest_algo_length_bytes(const char *algo)
{
    unsigned int rhash_id;

    /* strcmp is not defined for NULL */
    if (!algo) {
        return -1;
    }

    if (igloo_digest_name_to_rhash_id(&rhash_id, algo) != igloo_ERROR_NONE) {
        return -1;
    }

    return rhash_get_digest_size(rhash_id);
}

igloo_error_t igloo_hmac_new(igloo_hmac_t **self, igloo_ro_t group, const char *algo, const void *key, size_t keylen)
{
    ssize_t blocklen;
    char *keycopy;
    igloo_error_t err;
    size_t i;
    igloo_digest_t *inner;

    if (!self || igloo_ro_is_null(group) || !algo || !key)
        return igloo_ERROR_FAULT;

    blocklen = __digest_algo_blocklength(algo);
    if (blocklen < 1) {
        return igloo_ERROR_NOSYS;
    }

    keycopy = calloc(1, blocklen);
    if (!keycopy)
        return igloo_ERROR_NOMEM;

    if (keylen > (size_t)blocklen) {
        ssize_t digestlen = __digest_algo_length_bytes(algo);
        igloo_digest_t *digest;

        if (digestlen < 1) {
            free(keycopy);
            return igloo_ERROR_NOSYS;
        }

        if ((err = igloo_digest_new(&digest, group, algo)) != igloo_ERROR_NONE) {
            free(keycopy);
            return err;
        }

        if (igloo_digest_write(digest, key, keylen) != (ssize_t)keylen) {
            igloo_ro_unref(&digest);
            free(keycopy);
            return igloo_ERROR_GENERIC;
        }

        if (igloo_digest_read(digest, keycopy, digestlen) != (ssize_t)digestlen) {
            igloo_ro_unref(&digest);
            free(keycopy);
            return igloo_ERROR_GENERIC;
        }

        igloo_ro_unref(&digest);
    } else {
        memcpy(keycopy, key, keylen);
    }

    err = igloo_ro_new_raw(self, igloo_hmac_t, group);
    if (err != igloo_ERROR_NONE) {
        return err;
    }

    (*self)->algo = algo;
    (*self)->key = keycopy;
    (*self)->key_len = blocklen;

    if ((err = igloo_digest_new(&inner, group, algo)) != igloo_ERROR_NONE)
        return err;

    (*self)->inner = inner;

    for (i = 0; i < (size_t)blocklen; i++) {
        char x = keycopy[i] ^ 0x36;

        if (igloo_digest_write((*self)->inner, &x, 1) != 1) {
            igloo_ro_unref(self);
            return igloo_ERROR_GENERIC;
        }
    }

    return igloo_ERROR_NONE;
}

igloo_error_t igloo_hmac_copy(igloo_hmac_t **new, igloo_ro_t group, igloo_hmac_t *old)
{
    igloo_error_t err;
    igloo_digest_t *inner;

    if (!old)
        return igloo_ERROR_FAULT;

    err = igloo_ro_new_raw(new, igloo_hmac_t, group);
    if (err != igloo_ERROR_NONE)
        return err;

    (*new)->algo = old->algo;
    (*new)->key = malloc(old->key_len);
    if (!(*new)->key) {
        igloo_ro_unref(new);
        return igloo_ERROR_NOMEM;
    }
    memcpy((*new)->key, old->key, old->key_len);
    (*new)->key_len = old->key_len;

    err = igloo_digest_copy(&inner, group, old->inner);
    if (err != igloo_ERROR_NONE) {
        igloo_ro_unref(new);
        return err;
    }
    (*new)->inner = inner;

    return igloo_ERROR_NONE;
}

ssize_t igloo_hmac_write(igloo_hmac_t *hmac, const void *data, size_t len)
{
    if (!hmac)
        return -1;

    return igloo_digest_write(hmac->inner, data, len);
}

ssize_t igloo_hmac_read(igloo_hmac_t *hmac, void *buf, size_t len)
{
    igloo_digest_t *digest;
    igloo_error_t err;
    ssize_t digestlen;
    char *res;
    size_t i;
    ssize_t ret;

    if (!hmac)
        return -1;

    err = igloo_digest_new(&digest, hmac, hmac->algo);
    if (err != igloo_ERROR_NONE)
        return -1;

    digestlen = __digest_algo_length_bytes(hmac->algo);
    if (digestlen < 1)
        return -1;

    res = malloc(digestlen);
    if (!res) {
        igloo_ro_unref(&digest);
        return -1;
    }

    for (i = 0; i < hmac->key_len; i++) {
        char x = hmac->key[i] ^ 0x5C;
        if (igloo_digest_write(digest, &x, 1) != 1) {
            free(res);
            igloo_ro_unref(&digest);
            return -1;
        }
    }

    if (igloo_digest_read(hmac->inner, res, digestlen) != (ssize_t)digestlen) {
        free(res);
        igloo_ro_unref(&digest);
        return -1;
    }

    if (igloo_digest_write(digest, res, digestlen) != (ssize_t)digestlen) {
        free(res);
        igloo_ro_unref(&digest);
        return -1;
    }

    ret = igloo_digest_read(digest, buf, len);

    free(res);
    igloo_ro_unref(&digest);

    return ret;
}
