/*
 * lpass.c
 * lpass is a LessPass clone implemented in C.
 *
 * This file is part of lpass.
 *
 * lpass is free software: you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * lpass 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
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with lpass. If not, see <http://www.gnu.org/licenses/>.
 */

#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 "!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"

enum CharSet {
    CHARSET_LOWER   = 1 << 0,
    CHARSET_UPPER   = 1 << 1,
    CHARSET_DIGITS  = 1 << 2,
    CHARSET_SYMBOLS = 1 << 3,
};

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();
    if (ctx == NULL || bn_charsetlen == NULL || bn_remainder == NULL) {
        retval = 0;
        goto consume_entropy_cleanup;
    }

    int charsetlen = strlen(charset);
    if (BN_set_word(bn_charsetlen, charsetlen) == 0) {
        retval = 0;
        goto consume_entropy_cleanup;
    }
    for (int passlen = strlen(generated_pass);
         passlen < maxlen;
         passlen++) {
        BN_div(entropy, bn_remainder, entropy, bn_charsetlen, ctx);

        /*
         * If `remainder` cannot store the value in `bn_remainder`, it
         * will contain a very large number. Abort by checking if the
         * remainder is too larger according to `passlen`, since that
         * is also a failure case anyway.
         */
        uint64_t 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();
    if (ctx == NULL || bn_passlen == NULL || bn_remainder == NULL) {
        retval = 0;
        goto insert_str_pseudo_randomly_cleanup;
    }
    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);

        /*
         * If `remainder` cannot store the value in `bn_remainder`, it
         * will contain a very large number. Abort by checking if the
         * remainder is too larger according to `passlen`, since that
         * is also a failure case anyway.
         */
        uint64_t remainder = BN_get_word(bn_remainder);
        if (remainder >= passlen) {
            retval = 0;
            goto insert_str_pseudo_randomly_cleanup;
        }

        /*
         * 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 int
charsets_has_set(uint8_t charsets, enum CharSet set)
{
    return (charsets & set) != 0;
}

static int
build_charset(char *charset, uint8_t allowed_charsets)
{
    int c = 0;
    charset[0] = '\0';
    if (charsets_has_set(allowed_charsets, CHARSET_LOWER)) {
        strcat(charset, CHAR_SUBSET_LOWER);
        c++;
    }
    if (charsets_has_set(allowed_charsets, CHARSET_UPPER)) {
        strcat(charset, CHAR_SUBSET_UPPER);
        c++;
    }
    if (charsets_has_set(allowed_charsets, CHARSET_DIGITS)) {
        strcat(charset, CHAR_SUBSET_DIGITS);
        c++;
    }
    if (charsets_has_set(allowed_charsets, CHARSET_SYMBOLS)) {
        strcat(charset, CHAR_SUBSET_SYMBOLS);
        c++;
    }
    return c;
}

int
render_pass(BIGNUM *entropy, uint8_t allowed_charsets, char *out, int length)
{
    char charset[MAX_BUF + 1];
    int num_charsets = build_charset(charset, allowed_charsets);
    if (consume_entropy(out, entropy, charset, length - num_charsets) == 0) {
        return 0;
    }

    /*
     * After generating the initial password, add one character of
     * each charset to ensure at least one from that charset is there.
     */
    char str_to_add[num_charsets + 1];
    memset(str_to_add, 0, sizeof(str_to_add));
    int c = 0;
    if (charsets_has_set(allowed_charsets, CHARSET_LOWER)
        && consume_entropy(&str_to_add[c++], entropy, CHAR_SUBSET_LOWER, 1) == 0) {
        return 0;
    }
    if (charsets_has_set(allowed_charsets, CHARSET_UPPER)
        && consume_entropy(&str_to_add[c++], entropy, CHAR_SUBSET_UPPER, 1) == 0) {
        return 0;
    }
    if (charsets_has_set(allowed_charsets, CHARSET_DIGITS)
        && consume_entropy(&str_to_add[c++], entropy, CHAR_SUBSET_DIGITS, 1) == 0) {
        return 0;
    }
    if (charsets_has_set(allowed_charsets, CHARSET_SYMBOLS)
        && consume_entropy(&str_to_add[c++], entropy, CHAR_SUBSET_SYMBOLS, 1) == 0) {
        return 0;
    }
    return 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-lower          no lowercase letters in password\n", stderr);
    fputs("  --no-upper          no uppercase letters in password\n", stderr);
    fputs("  --no-digits         no digits in password\n", stderr);
    fputs("  --no-symbols        no special symbols (" CHAR_SUBSET_SYMBOLS ")\n", stderr);
    fputs("  -h, --help          output this message and exit\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[])
{
    int c;
    int length = DEFAULT_LENGTH;
    uint8_t allowed_charsets = CHARSET_LOWER
        | CHARSET_UPPER
        | CHARSET_DIGITS
        | CHARSET_SYMBOLS;
    for (;;) {
        int option_index = 0;
        static struct option long_options[] = {
            {"help",       no_argument,       0, 'h'},
            {"length",     required_argument, 0, 'L'},
            {"no-lower",   no_argument,       0,  0 },
            {"no-upper",   no_argument,       0,  0 },
            {"no-digits",  no_argument,       0,  0 },
            {"no-symbols", no_argument,       0,  0 },
            {0,            0,                 0,  0 }
        };

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

        switch (c) {
        case 0:
            if (strcmp("no-lower", long_options[option_index].name) == 0) {
                allowed_charsets &= ~CHARSET_LOWER;
            } else if (strcmp("no-upper", long_options[option_index].name) == 0) {
                allowed_charsets &= ~CHARSET_UPPER;
            } else if (strcmp("no-digits", long_options[option_index].name) == 0) {
                allowed_charsets &= ~CHARSET_DIGITS;
            } else if (strcmp("no-symbols", long_options[option_index].name) == 0) {
                allowed_charsets &= ~CHARSET_SYMBOLS;
            }
            break;
        case 'h':
            usage();
            exit(EXIT_SUCCESS);
        case 'L':
            if (!optarg
                || sscanf(optarg, "%d", &length) == 0
                || length < 1
                || length >= MAX_BUF) {
                usage();
                exit(EXIT_FAILURE);
            }
            break;
        default:
            usage();
            exit(EXIT_FAILURE);
        }
    }

    const char *master_pass = getenv("LESSPASS_MASTER_PASSWORD");
    if (master_pass == NULL) {
        err("environment variable LESSPASS_MASTER_PASSWORD is not set");
    }

    /*
     * 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));
    if (render_pass(entropy, allowed_charsets, generated_pass, length) == 0) {
        BN_free(entropy);
        err("Failed to generate password");
    }
    printf("%s\n", generated_pass);
    BN_free(entropy);
    return EXIT_SUCCESS;
}