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

#include <openssl/sha.h>
#include <argon2.h>

#include "bhash.h"
#include "phogen.h"

#include "passgeny.h"

/*
 * =========================================================================
 * Passgeny default parameters and constants
 * =========================================================================
 */

/*
 * The default pattern used for password generation
 */
#define PASSGENY_DEFAULT_PATTERN "^6p^6p^6ps2p100d"

/*
 * List of special characters (order is important)
 */
#define PASSGENY_SPECIAL "-./=_ "

/*
 * Argon2 parameters
 */
#define PASSGENY_ARGON2_SALT            "Passgeny v1"
#define PASSGENY_ARGON2_MEMORY_COST     8192
#define PASSGENY_ARGON2_TIME_COST       6
#define PASSGENY_ARGON2_PARALLEL        4
#define PASSGENY_ARGON2_HASH_LEN        128

/* Ensure that pg_master_hash is big enough to hold a SHA256 hash */
_Static_assert(
        sizeof(((passgeny_t *)NULL)->pg_master_hash) == SHA256_DIGEST_LENGTH,
        "pg_master_hash doesn't match SHA256_DIGEST_LENGTH");

static bool passgeny_patmatch(
        const char **pstr,
        const char *pattern,
        char *out,
        size_t out_len);

static void passgeny_gen_str(char *out, size_t out_sz, const char *pool, bhash_t *bh);

/*
 * =========================================================================
 * Public API
 * =========================================================================
 */
bool passgeny_init(passgeny_t *passgeny, const char *master_password)
{
    SHA256_CTX sha256;

    memset(passgeny, 0, sizeof(*passgeny));

    SHA256_Init(&sha256);
    SHA256_Update(&sha256, master_password, strlen(master_password));
    SHA256_Final(passgeny->pg_master_hash, &sha256);

    passgeny->pg_pattern = strdup(PASSGENY_DEFAULT_PATTERN);

    return true;
}

void passgeny_fini(passgeny_t *passgeny)
{
    /* Just nuke the master hash from memory */
    explicit_bzero(passgeny->pg_master_hash, sizeof(passgeny->pg_master_hash));
    free(passgeny->pg_pattern);
}

void passgeny_set_pattern(passgeny_t *passgeny, const char *pattern)
{
    free(passgeny->pg_pattern);
    passgeny->pg_pattern = strdup(pattern);
}

bool passgeny_generate(
        passgeny_t *passgeny,
        char *out,
        size_t out_sz,
        const char *domain,
        const char *username,
        const char * const tokens[],
        int ntokens)
{
    char ahash[PASSGENY_ARGON2_HASH_LEN];
    char popt[32];
    bhash_t bh;
    int ti;
    int rc;

    const char *ppat = passgeny->pg_pattern;

    /*
     * Generate the Argon2 hash
     */

    /* Caculate the input length */
    size_t in_len = 0;
    in_len += strlen(domain) + 1;
    in_len += strlen(username) + 1;
    in_len += sizeof(passgeny->pg_master_hash);

    for (ti = 0; ti < ntokens; ti++)
    {
        in_len += strlen(tokens[ti]);
    }

    char in[in_len];
    char *pin = in;

    strcpy(pin, domain);
    pin += strlen(domain) + 1;

    strcpy(pin, username);
    pin += strlen(username) + 1;

    for (ti = 0; ti < ntokens; ti++)
    {
        strcpy(pin, tokens[ti]);
        pin += strlen(tokens[ti]) + 1;
    }

    memcpy(pin, passgeny->pg_master_hash, sizeof(passgeny->pg_master_hash));
    pin += sizeof(passgeny->pg_master_hash);

    rc = argon2id_hash_raw(
            PASSGENY_ARGON2_TIME_COST,
            PASSGENY_ARGON2_MEMORY_COST,
            PASSGENY_ARGON2_PARALLEL,
            in,
            pin - in,
            PASSGENY_ARGON2_SALT,
            strlen(PASSGENY_ARGON2_SALT),
            ahash,
            sizeof(ahash));
    if (rc != 0)
    {
        fprintf(stderr, "Error generating argon2 hash.");
        return false;
    }

    bhash_init(&bh, ahash, sizeof(ahash));

    /*
     * Generate the password according to the currently set pattern
     */
    char cword[256];
    char cflags = '\0';
    while (*ppat != 0)
    {
        if (passgeny_patmatch(&ppat, "\\^", popt, sizeof(popt)))
        {
            cflags = '^';
            continue;
        }
        else if (passgeny_patmatch(&ppat, "([0-9]*)p", popt, sizeof(popt)))
        {
            size_t clen = atoi(popt);
            if (clen >= sizeof(cword))
            {
                fprintf(stderr, "Invalid phogen length: %zd\n", clen);
                return false;
            }
            else if (clen == 0)
            {
                clen = 1;
            }

            phogen_encode(cword, clen + 1, &bh);
        }
        else if (passgeny_patmatch(&ppat, "([0-9]*)s", popt, sizeof(popt)))
        {
            size_t clen = atoi(popt);
            if (clen >= sizeof(cword))
            {
                fprintf(stderr, "Invalid special string length: %zd\n", clen);
                return false;
            }
            else if (clen == 0)
            {
                clen = 1;
            }

            passgeny_gen_str(cword, clen + 1, PASSGENY_SPECIAL, &bh);
        }
        else if (passgeny_patmatch(&ppat, "([0-9]*)d", popt, sizeof(popt)))
        {
            int id = atoi(popt);
            if (id > INT_MAX)
            {
                fprintf(stderr, "Invalid decimal length: %d\n", id);
                return false;
            }
            else if (id == 0)
            {
                id = 10;
            }

            snprintf(cword, sizeof(cword), "%d", bhash_mod32(&bh, id));
        }
        else if (passgeny_patmatch(&ppat, "([0-9]*)x", popt, sizeof(popt)))
        {
            size_t clen = atoi(popt);
            if (clen >= sizeof(cword))
            {
                fprintf(stderr, "Invalid hex string length: %zd\n", clen);
                return false;
            }
            else if (clen == 0)
            {
                clen = 1;
            }

            passgeny_gen_str(cword, clen + 1, "0123456789abcdef", &bh);
        }
        else
        {
            printf("Error parsing pattern at: %s\n", ppat);
            return false;
        }

        if (cflags == '^')
        {
            cword[0] = toupper(cword[0]);
        }

        /* Reset flags */
        cflags = '\0';

        /* Store cword to out */
        size_t clen = strlen(cword);
        if (clen >= out_sz)
        {
            fprintf(stderr, "passgeny: Output string buffer is too short.\n");
            return false;
        }

        memcpy(out, cword, clen);
        out += clen;
        out_sz -= clen;
    }

    *out = '\0';

    passgeny->pg_last_bits_total = bhash_bits_avail(&bh);
    passgeny->pg_last_bits_used = bhash_bits_used(&bh);

    return true;
}

void passgeny_dump_verbose(const passgeny_t *pg)
{
    fprintf(stderr, "Total hash size is %2.0f bits.\n", pg->pg_last_bits_total);
    fprintf(stderr, "Used bits when generating the password: %0.2f.\n", pg->pg_last_bits_used);
}

/*
 * =========================================================================
 * Private functions
 * =========================================================================
 */

/*
 * Perform a pattern match on string `pstr`.
 *
 * If there's no match, this function returns false.
 *
 * If there's a match, pstr will be moved forward (past the pattern) and `opt`
 * will be filled with the first group match in regex.
 *
 * The value stored in `opt` will be truncated if its too long.
 */
static bool passgeny_patmatch(
        const char **pstr,
        const char *pattern,
        char *opt,
        size_t opt_sz)
{
    regmatch_t rm[2];
    regex_t re;

    bool retval = false;

    if (regcomp(&re, pattern, REG_EXTENDED) != 0)
    {
        return false;
    }

    if (regexec(&re, *pstr, sizeof(rm) / sizeof(rm[0]), rm, 0) == REG_NOMATCH)
    {
        goto exit;
    }

    /* All matches must start at the beginning of the string */
    if (rm[0].rm_so != 0) goto exit;

    if (opt != NULL && rm[1].rm_so >= 0 && rm[1].rm_eo >= 0)
    {
        size_t ropt_sz = rm[1].rm_eo - rm[1].rm_so + 1;
        if (ropt_sz > opt_sz) ropt_sz = opt_sz;

        memcpy(opt, *pstr + rm[1].rm_so, ropt_sz - 1);
        opt[ropt_sz - 1] = '\0';
    }

    *pstr += rm[0].rm_eo;
    retval = true;

exit:
    regfree(&re);
    return retval;
}

/*
 * Generate a string by using the bhash and a pool of characters
 */
static void passgeny_gen_str(char *out, size_t out_sz, const char *pool, bhash_t *bh)
{
    int ii;
    size_t psz;

    if (out_sz == 0) return;

    psz = strlen(pool);

    for (ii = 0; ii < (out_sz - 1); ii++)
    {
        out[ii] = pool[bhash_mod32(bh, psz)];
    }
    out[ii] = '\0';
}
