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

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

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

#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);

enum CharSet {
    CHAR_SET_LOWER   = 1 << 0,
    CHAR_SET_UPPER   = 1 << 1,
    CHAR_SET_DIGITS  = 1 << 2,
    CHAR_SET_SYMBOLS = 1 << 3,
};

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 NULL;
    }

    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 NULL;
    }

    /* 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);
}

static void
err(const char *msg)
{
    fprintf(stderr, "lpass: %s\n", msg);
    exit(EXIT_FAILURE);
}

static void
usage()
{
    fputs("usage: lpass SITE LOGIN [OPTIONS]\n", stderr);
    fputs("\nlpass is a C implementation of LessPass, a stateless password manager.\n", stderr);
    fputs("\nPositional arguments:\n", stderr);
    fputs("  SITE    The domain name or URL used in password generation.\n", stderr);
    fputs("  LOGIN   The username/login used in that domain or URL.\n", stderr);
    fputs("\nOptions:\n", stderr);
    fputs("  -L LENGTH, --length LENGTH\n", stderr);
    fputs("                      password length (default: 16)\n", stderr);
    fputs("  --no-symbols        no symbols in password\n", stderr);
    fputs("\nThe environment variable LESSPASS_MASTER_PASSWORD must be set with the master\n", stderr);
    fputs("password for producing the desired password.\n", stderr);
}

int
main(int argc, char *argv[])
{
    const char *master_pass = getenv("LESSPASS_MASTER_PASSWORD");
    if (master_pass == NULL) {
        err("environment variable LESSPASS_MASTER_PASSWORD is not set");
    }

    int c;
    int length = DEFAULT_LENGTH;
    uint8_t charsets = 0;
    for (;;) {
        int option_index = 0;
        static struct option long_options[] = {
            {"length",     required_argument, 0, 'L'},
            {"no-symbols", no_argument,       0,  0 },
            {0,            0,                 0,  0 }
        };

        c = getopt_long(argc, argv, "L:", long_options, &option_index);
        if (c == -1) {
            break;
        }

        switch (c) {
        case 0:
            if (option_index == 1) {
                /* TODO: no-symbols */
            }
            break;
        case 'L':
            if (!optarg || sscanf(optarg, "%d", &length) == 0 || length < 1) {
                usage();
                exit(EXIT_FAILURE);
            }
            break;
        default:
            usage();
            exit(EXIT_FAILURE);
        }
    }

    /*
     * We are still looking for 2 positional arguments, abort if they
     * are not there.
     */
    if (optind > argc - 2) {
        usage();
        exit(EXIT_FAILURE);
    }

    const char *site = argv[optind++];
    const char *login = argv[optind++];
    int passlen = strlen(master_pass);
    uint64_t counter = 1;
    BIGNUM *entropy = calc_entropy(site, login, counter, master_pass, passlen);
    if (entropy == NULL) {
        err("Failed to calculate entropy");
    }

    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 EXIT_SUCCESS;
}