#include <openssl/bn.h>
#include <openssl/evp.h>

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

#define MAX_BUF 1024
#define ENTROPY_ITERATIONS 100000
#define ENTROPY_KEY_LENGTH 32

#define CHAR_SUBSET_LOWER "abcdefghijklmnopqrstuvwxyz"
#define CHAR_SUBSET_UPPER "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
#define CHAR_SUBSET_DIGITS "0123456789"
#define CHAR_SUBSET_SYMBOLS "!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"

const char *CHAR_SET = (CHAR_SUBSET_LOWER
                        CHAR_SUBSET_UPPER
                        CHAR_SUBSET_DIGITS
                        CHAR_SUBSET_SYMBOLS);

static BIGNUM *
calc_entropy(const char *site,
             const char *login,
             uint64_t counter,
             const char *master_pass,
             int passlen)
{
    char salt[MAX_BUF + 1];
    memset(salt, 0, sizeof(salt));
    int saltlen = snprintf(salt, MAX_BUF, "%s%s%lx", site, login, counter);
    if (saltlen > MAX_BUF) {
        return 0;
    }

    unsigned char key[ENTROPY_KEY_LENGTH];
    memset(key, 0, sizeof(key));
    int status = PKCS5_PBKDF2_HMAC(master_pass, passlen,
                                   (const unsigned char *) salt, saltlen,
                                   ENTROPY_ITERATIONS,
                                   EVP_sha256(),
                                   ENTROPY_KEY_LENGTH,
                                   key);
    if (status == 0) {
        return 0;
    }

    /* NULL as last arg to allocate a new BIGNUM */
    return BN_bin2bn(key, ENTROPY_KEY_LENGTH, NULL);
}

static int
consume_entropy(char *generated_pass, BIGNUM *entropy, const char *charset, int maxlen)
{
    int retval = 1;
    BN_CTX *ctx = BN_CTX_new();
    BIGNUM *bn_charsetlen = BN_new();
    BIGNUM *bn_remainder = BN_new();
    int charsetlen = strlen(charset);
    BN_set_word(bn_charsetlen, charsetlen);
    for (int passlen = strlen(generated_pass);
         passlen < maxlen;
         passlen++) {
        BN_div(entropy, bn_remainder, entropy, bn_charsetlen, ctx);
        int remainder = BN_get_word(bn_remainder);
        if (remainder >= charsetlen) {
            retval = 0;
            goto consume_entropy_cleanup;
        }
        generated_pass[passlen] = charset[remainder];
    }
consume_entropy_cleanup:
    BN_free(bn_remainder);
    BN_free(bn_charsetlen);
    BN_CTX_free(ctx);
    return retval;
}

static int
insert_str_pseudo_randomly(char *generated_pass, BIGNUM *entropy, const char *s)
{
    int retval = 1;
    char buf[MAX_BUF + 1];
    uint64_t passlen = (uint64_t) strlen(generated_pass);
    BN_CTX *ctx = BN_CTX_new();
    BIGNUM *bn_passlen = BN_new();
    BIGNUM *bn_remainder = BN_new();
    for (char c = *s; *s != '\0'; c = *(++s)) {
        if (BN_set_word(bn_passlen, passlen) == 0) {
            retval = 0;
            goto insert_str_pseudo_randomly_cleanup;
        }
        BN_div(entropy, bn_remainder, entropy, bn_passlen, ctx);
        uint64_t remainder = BN_get_word(bn_remainder);
        /*
         * Idea here is to add the char `c` at position `remainder` in `generated_pass`.
         * 1. Copy the part that would need to be shifted into `buf`.
         */
        memset(buf, 0, sizeof(buf));
        strncpy(buf, &generated_pass[remainder], passlen - remainder);
        /*
         * 2. Add new character, then copy `buf` back into `generated_pass`.
         */
        generated_pass[remainder] = c;
        strncpy(&generated_pass[remainder + 1], buf, passlen - remainder);
        passlen++;
    }
insert_str_pseudo_randomly_cleanup:
    BN_free(bn_remainder);
    BN_free(bn_passlen);
    BN_CTX_free(ctx);
    return retval;
}

static void
render_pass(BIGNUM *entropy, char *out, int length)
{
    /* XXX: FIXME: implement charset restrictions */
    int num_charsets = 4;
    consume_entropy(out, entropy, CHAR_SET, length - num_charsets);

    /*
     * After generating the initial password, add one character of
     * each charset to ensure at least one is there.
     */
    char str_to_add[num_charsets + 1];
    memset(str_to_add, 0, sizeof(str_to_add));
    consume_entropy(&str_to_add[0], entropy, CHAR_SUBSET_LOWER, 1);
    consume_entropy(&str_to_add[1], entropy, CHAR_SUBSET_UPPER, 1);
    consume_entropy(&str_to_add[2], entropy, CHAR_SUBSET_DIGITS, 1);
    consume_entropy(&str_to_add[3], entropy, CHAR_SUBSET_SYMBOLS, 1);
    insert_str_pseudo_randomly(out, entropy, str_to_add);
}

int
main(int argc, char *argv[])
{
    if (argc != 5) {
        return 1;
    }

    const char *master_pass = argv[1];
    int passlen = strlen(master_pass);
    const char *site = argv[2];
    const char *login = argv[3];
    int length = 0;
    if (sscanf(argv[4], "%d", &length) == 0) {
        return 1;
    }

    uint64_t counter = 1;
    BIGNUM *entropy = calc_entropy(site, login, counter, master_pass, passlen);
    if (entropy == NULL) {
        return 1;
    }

    char generated_pass[length + 1];
    memset(generated_pass, 0, sizeof(generated_pass));
    render_pass(entropy, generated_pass, length);
    printf("%s\n", generated_pass);
    BN_free(entropy);
    return 0;
}