/*
 * Command line tool for setting up PKCS #11.  Mostly this means
 * things like setting PINs.
 */

#define _POSIX_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <termios.h>
#include <getopt.h>

#include <hal.h>

#include "sql_common.h"

/*
 * Apparently the cool kids don't use getpassword() anymore, and there
 * is no fully portable replacement, so the advice is just to roll
 * one's own using termios.  Yuck, but gives us the opportunity to
 * provide a bit of visual feedback, what the heck.
 *
 * This version requires /dev/tty.  Use something else if you're
 * reading from a pipe.
 */

static int getpin_tty(const char *prompt,
		      char *pinbuf,
		      const size_t pinbuf_len)
{
  struct termios oflags, nflags;
  int c = '\0', ok = 0, fixtty = 0;
  char *p = pinbuf;
  FILE *f = NULL;

  assert(prompt != NULL && pinbuf != NULL);

  if ((f = fopen("/dev/tty", "r+")) == NULL ||
      setvbuf(f, NULL, _IONBF, 0) == EOF)
    goto fail;

  tcgetattr(fileno(f), &oflags);
  nflags = oflags;
  nflags.c_lflag &= ~ (ECHO | ICANON);

  if (tcsetattr(fileno(f), TCSANOW, &nflags) < 0)
    goto fail;

  fixtty = 1;
  fputs(prompt, f);

  while (p < pinbuf + pinbuf_len - 1 && (c = fgetc(f)) != EOF && c != '\n') {
    fputc('*', f);
    *p++ = (char) c;
  }

  assert(p <= pinbuf + pinbuf_len - 1);
  *p = '\0';
  fputc('\n', f);
  ok = (c == '\n');

 fail:

  if (!ok)
    perror("getpin_tty() failed");

  if (fixtty && tcsetattr(fileno(f), TCSANOW, &oflags) < 0)
    perror("getpin_tty() couldn't restore termios settings");

  if (f != NULL)
    fclose(f);

  return ok;
}



#define OPTIONS									\
  OPT_FLG('s', "set-so-pin",     "set Security Officer PIN")			\
  OPT_FLG('u', "set-user-pin",	 "set \"user\" PIN")				\
  OPT_ARG('i', "set-iterations", "set PBKDF2 iteration count")			\
  OPT_FLG('p', "pin-from-stdin", "read PIN from stdin instead of /dev/tty")	\
  OPT_END

#define OPT_END

static void usage(const int code, const char *jane)
{
  assert(jane != NULL);
  FILE *f = code ? stderr : stdout;

  fprintf(f, "usage: %s [options]\noptions:\n", jane);

#define OPT_FLG(_short_, _long_, _help_)  fprintf(f, "  -%c      --%-32s%s\n", _short_, _long_, _help_);
#define OPT_ARG(_short_, _long_, _help_)  fprintf(f, "  -%c ARG  --%-32s%s\n", _short_, _long_ " ARG", _help_);
  OPTIONS;
#undef OPT_ARG
#undef OPT_FLG

  exit(code);
}

static void parse_args(int argc, char *argv[],
		       int *do_set_so_pin,
		       int *do_set_user_pin,
		       int *do_set_iterations,
		       int *read_from_stdin,
		       unsigned long *iterations)
{
  char *endptr;
  int c;

#define OPT_FLG(_short_, _long_, _help_) _short_,
#define OPT_ARG(_short_, _long_, _help_) _short_, ':',
  static const char short_opts[] = { OPTIONS '\0' };
#undef OPT_ARG
#undef OPT_FLG

#define OPT_FLG(_short_, _long_, _help_) { _long_, no_argument,       NULL, _short_ },
#define OPT_ARG(_short_, _long_, _help_) { _long_, required_argument, NULL, _short_ },
  struct option long_opts[] = { OPTIONS { NULL } };
#undef OPT_ARG
#undef OPT_FLG

  assert(argv != 0 &&
	 do_set_so_pin != 0 && do_set_user_pin != 0 && do_set_iterations != NULL &&
	 read_from_stdin != NULL && iterations != NULL);
  opterr = 0;

  if (argc == 1)
    usage(0, argv[0]);

  while ((c = getopt_long(argc, argv, short_opts, long_opts, NULL)) > 0) {
    switch (c) {

    case 'h':
      usage(0, argv[0]);

    case 'i':
      *do_set_iterations = 1;
      *iterations = strtoul(optarg, &endptr, 0);
      if (*optarg == '\0' || *endptr != '\0')
	usage(1, argv[0]);
      continue;

    case 'p':
      *read_from_stdin = 1;
      continue;

    case 's':
      *do_set_so_pin = 1;
      continue;

    case 'u':
      *do_set_user_pin = 1;
      continue;

    default:
      usage(1, argv[0]);
    }
  }

  if (optind != argc)
    usage(1, argv[0]);
}



#define lose(_msg_)			\
  do {					\
    fprintf(stderr, "%s\n", _msg_);	\
    goto fail;				\
  } while (0)

static int set_iterations(unsigned long iterations)
{
  static const char update_query[] =
    " UPDATE global SET pbkdf2_iterations = ?";

  sqlite3_stmt *q = NULL;
  int ok = 0;
  
  if (!sql_check_ok(sql_prepare(&q, update_query))		||
      !sql_check_ok(sqlite3_bind_int64(q, 1, iterations))	||
      !sql_check_done(sqlite3_step(q)))
    lose("Couldn't update database");

  ok = 1;

 fail:
  sqlite3_finalize(q);
  return ok;
}

static int set_pin(const char * const pin_type, const int read_from_stdin)
{
  static const char iterations_query[] = 
    " SELECT pbkdf2_iterations FROM global";

  static const char update_format[] =
    " UPDATE global SET %s_pin = ?1, %s_pin_salt = ?2";

  /* Allow user to change these lengths? */
  uint8_t pinbuf[32], salt[16];

  char pin[P11_MAX_PIN_LENGTH + 1], *p;
  sqlite3_stmt *q = NULL;
  hal_error_t err;
  int ok = 0;

  if (read_from_stdin) {
    if (fgets(pin, sizeof(pin), stdin) == NULL) {
      perror("Couldn't read PIN");
      return 0;
    }
    if ((p = strchr(pin, '\n')) != NULL)
      *p = '\0';
  }

  else {
    char prompt[sizeof("Enter user PIN:  ")];
    snprintf(prompt, sizeof(prompt), "Enter %s PIN: ", pin_type);
    if (!getpin_tty(prompt, pin, sizeof(pin)))
      return 0;
  }

  const size_t len = strlen(pin);

  if (len < P11_MIN_PIN_LENGTH || len > P11_MAX_PIN_LENGTH) {
    fprintf(stderr, "Unacceptable length %lu for %s PIN, allowd range [%lu, %lu]\n",
	    (unsigned long) len, pin_type,
	    (unsigned long) P11_MIN_PIN_LENGTH, (unsigned long) P11_MAX_PIN_LENGTH);
    return 0;
  }

  if (!sql_check_ok(sql_prepare(&q, iterations_query))	||
      !sql_check_row(sqlite3_step(q))			||
      sqlite3_column_type(q, 0) == SQLITE_NULL)
    lose("Couldn't retrieve PBKDF2 iteration count from SQL");
    
  if ((err = hal_get_random(salt, sizeof(salt))) != HAL_OK) {
    fprintf(stderr, "Couldn't generate salt: %s\n", hal_error_string(err));
    goto fail;
  }

  if ((err = hal_pbkdf2(hal_hash_sha256, (uint8_t *) pin, len, salt, sizeof(salt),
			pinbuf, sizeof(pinbuf), sqlite3_column_int(q, 0))) != HAL_OK) {
    fprintf(stderr, "Couldn't process new PIN: %s\n", hal_error_string(err));
    goto fail;
  }

  if (!sql_check_ok(sql_finalize_and_clear(&q))					||
      !sql_check_ok(sql_prepare(&q, update_format, pin_type, pin_type))		||
      !sql_check_ok(sqlite3_bind_blob(q, 1, pinbuf, sizeof(pinbuf), NULL))	||
      !sql_check_ok(sqlite3_bind_blob(q, 2, salt, sizeof(salt), NULL))		||
      !sql_check_done(sqlite3_step(q)))
    lose("Couldn't update database");
      
  ok = 1;

 fail:
  sqlite3_finalize(q);
  memset(pin, 0, sizeof(pin));
  memset(pinbuf, 0, sizeof(pinbuf));
  memset(salt, 0, sizeof(salt));
  return ok;
}

int main(int argc, char *argv[])
{
  int do_set_so_pin = 0, do_set_user_pin = 0, do_set_iterations = 0, read_from_stdin = 0;
  unsigned long iterations;
  int ok = 0;

  parse_args(argc, argv, &do_set_so_pin, &do_set_user_pin, &do_set_iterations, &read_from_stdin, &iterations);

  if (!sql_init() || !sql_exec("BEGIN"))
    lose("Couldn't initialize SQL, giving up");

  if (do_set_iterations && !set_iterations(iterations))
    lose("Couldn't set PBKDF2 iteration count");

  if (do_set_so_pin && !set_pin("so", read_from_stdin))
    lose("Couldn't set SO PIN");

  if (do_set_user_pin && !set_pin("user", read_from_stdin))
    lose("Couldn't set user PIN");

  if (!sql_exec("COMMIT"))
    lose("Couldn't commit SQL transaction");

  ok = 1;

 fail:
  sql_fini();
  return !ok;
}