/*
 * scdtools - Tools for Scdaemon and OpenPGP smartcards
 * Copyright (C) 2014,2015 Damien Goutte-Gattat
 *
 * This program 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.
 *
 * This program 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 this program. If not, see <http://www.gnu.org/licenses/>.
 */

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

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <ctype.h>
#include <time.h>
#include <getopt.h>
#include <err.h>

#include <gcrypt.h>

#include "gpg-util.h"

#define DEFAULT_TIME_STEP   30
#define DEFAULT_DIGITS      6
#define MAX_KEY_SIZE        64

static void
usage(int status)
{
    puts("Usage: scdtotp [options]\n\
Generate a time-based one-time password (TOTP) from a seed\n\
found in a OpenPGP smartcard.\n\
");

    puts("Options:\n\
  -h, --help            Display this help message.\n\
  -v, --version         Display the version message.\n\
");

    printf("\
  -t, --time SECONDS    Generate OTP for the specified time\n\
                        (in seconds) instead of current time.\n\
  -s, --step N          Specify the time step factor\n\
                        (defualt is %d).\n\
  -d, --digits N        Specify number of digits to output\n\
                        (default is %d).\n\
\n", DEFAULT_TIME_STEP, DEFAULT_DIGITS);

    puts("\
  -m, --mac-algo ALG    Use the specified HMAC algorithm.\n\
                        ALG can be 'sha1' (default)',\n\
                        'sha256' or 'sha512'.\n\
");

    printf("\
  -n, --private-do N    Read key from private data object N\n\
                        (default is %d).\n\
\n", DEFAULT_PRIVATE_DO);

    printf("Report bugs to <%s>.\n", PACKAGE_BUGREPORT);

    exit(status);
}

static void
info(void)
{
    printf("\
scdtotp (scdtools %s)\n\
Copyright (C) 2014,2015 Damien Goutte-Gattat\n\
\n\
This program is released under the GNU General Public License.\n\
See the COPYING file or <http://www.gnu.org/licenses/gpl.html>.\n\
", VERSION);

    exit(EXIT_SUCCESS);
}

/*
 * Convert a time value (in seconds) into a moving factor for the
 * TOTP algorithm as per RFC 6238.
 *
 * @param time   Elapsed seconds since Unix epoch.
 * @param step   Time step parameter.
 * @param buffer Byte buffer to store the formatted factor. Must be
 *               at least 8 bytes long.
 * @param len    Size of the output buffer.
 *
 * @return The number of bytes written to the buffer (always 8).
 */
static int
time2factor(time_t time, int step, unsigned char *buffer, size_t len)
{
    unsigned n_steps;
    int i, j;

    n_steps = time / step;

    for ( i = 0, j = 7; i < 8; i++, j-- )
        buffer[i] = (n_steps & (uint64_t)(0xFF << (j * 8))) >> (j * 8);

    return 8;
}

/*
 * Convert a hexstring into a bytes array.
 *
 * @param hexstring The string to convert.
 * @param buffer    The output buffer.
 * @param len       Size of the output buffer.
 *
 * @return The number of bytes written to the output buffer, or
 *         -1 if the buffer is not large enough.
 */
static int
hexstring2bytes(const char *hexstring, unsigned char *buffer, size_t len)
{
    int high_nibble, n;
    unsigned val;

    high_nibble = 1;
    n = 0;

    while ( *hexstring && n < len ) {
        switch ( *hexstring ) {
        case '0':
        case '1':
        case '2':
        case '3':
        case '4':
        case '5':
        case '6':
        case '7':
        case '8':
        case '9':
            val = *hexstring - '0';
            break;

        case 'a':
        case 'b':
        case 'c':
        case 'd':
        case 'e':
        case 'f':
            val = *hexstring - 'a' + 10;
            break;

        case 'A':
        case 'B':
        case 'C':
        case 'D':
        case 'E':
        case 'F':
            val = *hexstring - 'A' + 10;
            break;

        default:
            hexstring += 1;
            continue;
        }

        if ( high_nibble ) {
            buffer[n] = val << 4;
            high_nibble = 0;
        }
        else {
            buffer[n] |= (val & 0xF);
            high_nibble = 1;
            n += 1;
        }

        hexstring += 1;
    }

    if ( *hexstring )
        return -1;  /* Buffer too small. */

    return n;
}

/*
 * Generate a time-based one-time password as per RFC 6238.
 *
 * @param algo   Libgcrypt identifier of the HMAC algorithm to use.
 * @param key    Hexstring representation of the key.
 * @param time   Time in seconds.
 * @param step   Time step parameter.
 * @param digits Number of digits of the OTP to generate
 *               (must be 6, 7, or 8).
 * @param otp    Pointer to an unsigned integer to store the
 *               generated OTP value.
 *
 * @return 0 if OTP generation was successful, or a gpg_error_t
 *         error code.
 */
static gpg_error_t
generate_totp(int           algo,
              const char   *key,
              time_t        time,
              unsigned      step,
              int           digits,
              unsigned     *otp)
{
    gcry_mac_hd_t mac;
    gpg_error_t e;
    unsigned char buffer[MAX_KEY_SIZE];
    size_t len;
    int offset;

    static int modulo[] = { 0, 0, 0, 0, 0, 0, 1000000, 10000000, 100000000 };

    if ( (len = hexstring2bytes(key, buffer, sizeof(buffer))) == -1 )
        return gcry_error(GPG_ERR_INV_KEYLEN);

    if ( ! (e = gcry_mac_open(&mac, algo, 0, NULL)) ) {

        if ( ! (e = gcry_mac_setkey(mac, buffer, len)) ) {

            (void)time2factor(time, step, buffer, sizeof(buffer));
            if ( ! (e = gcry_mac_write(mac, buffer, 8)) ) {

                len = sizeof(buffer);
                if ( ! (e = gcry_mac_read(mac, buffer, &len)) ) {

                    offset = buffer[len - 1] & 0xF;
                    *otp =
                        ((buffer[offset]     & 0x7F) << 24) |
                        ((buffer[offset + 1] & 0xFF) << 16) |
                        ((buffer[offset + 2] & 0xFF) <<  8) |
                         (buffer[offset + 3] & 0xFF);
                    *otp = *otp % modulo[digits];

                    e = gcry_error(GPG_ERR_NO_ERROR);
                }
            }
        }

        gcry_mac_close(mac);
    }

    return e;
}

struct key_data_t {
    unsigned char  *buffer;
    size_t          len;
    int             found;
};

/*
 * Callback for the below function.
 */
static gpg_error_t
get_otp_key_cb(void *arg, const char *line)
{
    struct key_data_t *key = (struct key_data_t *)arg;

    if ( ! strncmp("PRIVATE-DO-", line, 11) ) {
        const char *cursor;
        size_t n;

        cursor = line + 12;
        n = 0;

        while ( *cursor != '\n' && n < key->len - 1 ) {
            if ( isxdigit(*cursor) )
                key->buffer[n++] = *cursor;
            cursor += 1;
        }

        key->buffer[n] = '\0';
        key->found = 1;
    }

    return 0;
}

/*
 * Retrieve an OTP key from an inserted OpenPGP smartcard.
 *
 * @param buffer String buffer to store the retrieved key.
 * @param len    Size of the output buffer.
 *
 * @return 0 if successful, or a gpg_error_t error code.
 */
static gpg_error_t
get_otp_key(unsigned privatedo, unsigned char *buffer, size_t len)
{
    assuan_context_t ctx;
    gpg_error_t e;
    struct key_data_t key;
    char command[32];

    key.buffer = buffer;
    key.len = len;
    key.found = 0;

    if ( privatedo < 1 || privatedo > 4 )
        return gcry_error(GPG_ERR_INV_ARG);

    if ( ! (e = connect_to_scdaemon(&ctx)) ) {

        snprintf(command, sizeof(command), "GETATTR PRIVATE-DO-%d", privatedo);

        if ( ! (e = assuan_transact(ctx, command, NULL, NULL,
                        NULL, NULL, get_otp_key_cb, &key)) ) {
            if ( key.found )
                e = gcry_error(GPG_ERR_NO_ERROR);
            else
                e = gcry_error(GPG_ERR_NO_DATA);
        }

        assuan_release(ctx);
    }

    return e;
}

static gpg_error_t
get_serial_cb(void *arg, const char *line)
{
    char *serial = (char *)arg;

    if ( ! strncmp("SERIALNO ", line, 9) ) {
        strcpy(serial, line + 9);
        return GPG_ERR_NO_ERROR;
    }
    else
        return GPG_ERR_NO_DATA;
}

static gpg_error_t
verify_pin(int admin)
{
    assuan_context_t ctx;
    gpg_error_t e;
    char command[64], serial[64];
    char *line;
    size_t len;

    if ( ! (e = connect_to_agent(&ctx, 1)) ) {

        if ( ! (e = assuan_transact(ctx, "SCD GETATTR SERIALNO",
                        NULL, NULL, NULL, NULL, get_serial_cb, serial)) ) {

            snprintf(command, sizeof(command), "SCD CHECKPIN %s%s",
                    serial, admin ? "[CHV3]" : "");
            e = assuan_transact(ctx, command, NULL, NULL, NULL, NULL,
                    NULL, NULL);
        }

        assuan_release(ctx);
    }

    return e;
}


static void
print_otp(unsigned otp, unsigned digits)
{
    char fmt[] = "%06d\n";

    fmt[2] = digits + '0';
    printf(fmt, otp);
}

static unsigned long
get_uinteger_or_die(const char *arg)
{
    unsigned long val;
    char *endptr;

    errno = 0;
    val = strtoul(arg, &endptr, 10);
    if ( errno != 0 || endptr == arg )
        errx(EXIT_FAILURE, "Invalid argument, unsigned integer expected: %s", arg);

    return val;
}

int
main(int argc, char **argv)
{
    int c, algo;
    unsigned char key[MAX_KEY_SIZE * 2 + 1];
    unsigned otp, step, digits, privatedo;
    time_t secs;
    gcry_error_t e;

    struct option options[] = {
        { "help",       0, NULL, 'h' },
        { "version",    0, NULL, 'v' },
        { "time",       1, NULL, 't' },
        { "step",       1, NULL, 's' },
        { "digits",     1, NULL, 'd' },
        { "mac-algo",   1, NULL, 'm' },
        { "private-do", 1, NULL, 'n' },
        { NULL,         0, NULL, 0 }
    };

    setprogname(argv[0]);
    secs = time(NULL);
    step = DEFAULT_TIME_STEP;
    digits = DEFAULT_DIGITS;
    algo = GCRY_MAC_HMAC_SHA1;
    privatedo = DEFAULT_PRIVATE_DO;

    while ( (c = getopt_long(argc, argv, "hvt:s:d:m:n:",
                             options, NULL)) != -1 ) {
        switch ( c ) {
        case 'h':
            usage(EXIT_SUCCESS);
            break;

        case '?':
            usage(EXIT_FAILURE);
            break;

        case 'v':
            info();
            break;

        case 't':
            secs = get_uinteger_or_die(optarg);
            break;

        case 's':
            step = get_uinteger_or_die(optarg);
            break;

        case 'd':
            digits = get_uinteger_or_die(optarg);
            if ( digits < 6 || digits > 8 )
                errx(EXIT_FAILURE, "digits must be either 6, 7, or 8");
            break;

        case 'm':
            if ( ! strcmp("sha1", optarg) )
                algo = GCRY_MAC_HMAC_SHA1;
            else if ( ! strcmp("sha256", optarg) )
                algo = GCRY_MAC_HMAC_SHA256;
            else if ( ! strcmp("sha512", optarg) )
                algo = GCRY_MAC_HMAC_SHA512;
            else
                errx(EXIT_FAILURE, "unsupported HMAC algorithm: %s", optarg);
            break;

        case 'n':
            privatedo = get_uinteger_or_die(optarg);
            if ( privatedo < 1 || privatedo > 4 )
                errx(EXIT_FAILURE, "DO number must be 1, 2, 3, or 4");
            break;
        }
    }

    if ( ! gcry_check_version(GCRYPT_VERSION) )
        errx(EXIT_FAILURE, "libgcrypt version mismatch");

    gcry_control(GCRYCTL_DISABLE_SECMEM, 0);
    gcry_control(GCRYCTL_INITIALIZATION_FINISHED, 0);

    if ( privatedo > 2 && (e = verify_pin(privatedo == 4)) )
        errx(EXIT_FAILURE, "Cannot get key from token: %s", gcry_strerror(e));

    if ( (e = get_otp_key(privatedo, key, sizeof(key))) )
        errx(EXIT_FAILURE, "Cannot get key from token: %s", gcry_strerror(e));

    if ( (e = generate_totp(algo, key, secs, step, digits, &otp)) )
        errx(EXIT_FAILURE, "Cannot generate OTP: %s", gcry_strerror(e));

    print_otp(otp, digits);

    return EXIT_SUCCESS;
}