/* timelock -- timed release encryption utility

   Copyright (C) 2017 Arun I <arunisaac@systemreboot.net>

   This file is part of timelock.

   timelock 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.

   timelock 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 timelock.  If not, see <http://www.gnu.org/licenses/>.  */

// TODO: Implement in-place
// TODO: Strengthen randomness

#include <stdio.h>
#include <math.h>
#include <gcrypt.h>
#include <inttypes.h>
#include <time.h>
#include <argp.h>
#include <error.h>
#include <boost/preprocessor/seq.hpp>
#include "config.h"

#define KEY_SIZE 32
#define NONCE_SIZE 8
#define READ_BUFFER_SIZE 1024

#define PRIME_BITS 1024
#define KEY_BITS (KEY_SIZE * 8)
#define SQUARINGS_PER_ITERATION 1023
#define TEST_ITERATIONS 10000

#define DEFAULT_TIME 600

// Magic bytes
#define MAGIC_BYTES "timelock"
// Output format version number
#define OUTPUT_VERSION 0
// Chunk types in output
enum chunk_type {END_OF_CHUNKS, CHUNK_N, CHUNK_A, CHUNK_TIMES, CHUNK_CKEY};
void print_parameter (enum chunk_type, gcry_mpi_t);
// Cipher algorithms
enum cipher_algo {SALSA20};

#define _ALLOCATE_MPI(r, _, var) gcry_mpi_t var = gcry_mpi_new(1);
#define ALLOCATE_MPI(mpi) BOOST_PP_SEQ_FOR_EACH(_ALLOCATE_MPI, _, mpi)
#define _ALLOCATE_SECURE_MPI(r, _, var) gcry_mpi_t var = gcry_mpi_snew(1);
#define ALLOCATE_SECURE_MPI(mpi) BOOST_PP_SEQ_FOR_EACH(_ALLOCATE_SECURE_MPI, _, mpi)
#define _CLEANUP(r, _, var) _Generic((var), gcry_mpi_t: gcry_mpi_release, \
				     gcry_cipher_hd_t: gcry_cipher_close, \
				     default: free)(var);
#define CLEANUP(vars) BOOST_PP_SEQ_FOR_EACH(_CLEANUP, _, vars)
#define WITH_MPI(mpi, body) { ALLOCATE_MPI(mpi); body; CLEANUP(mpi); }

#define TRY_CATCH_GCRYPT(try, message, handle) if (handle_gcrypt_error(try, message)) { handle; exit(1); }
#define SAFE_FREAD(buffer, len, handle)		\
  if (fread(buffer, 1, len, stdin) < len)		\
    { handle; error(1, 0, "Incomplete input\n"); }
#define SAFE_FWRITE(buffer, len, handle)			\
  if (fwrite(buffer, 1, len, stdout) < len)			\
    { handle; error(1, 0, "Error writing to standard output"); }
#define SAFE_MALLOC(ptr, size, handle) if (!(ptr = malloc(size))) { handle; perror(__func__); exit(1); }
#define SAFE_REALLOC(ptr, size, handle) if (!(ptr = realloc(ptr, size))) { handle; perror(__func__); exit(1); }

// gcrypt helper functions
gcry_mpi_t mpi_set_ulli (gcry_mpi_t, unsigned long long);
unsigned long long mpi_to_ulli (gcry_mpi_t w);
gcry_error_t generate_prime (gcry_mpi_t*);
int handle_gcrypt_error (gcry_error_t, const char*);
gcry_error_t open_cipher (gcry_cipher_hd_t* handle, const void *key, const unsigned char *nonce);

size_t read_stdin (unsigned char**);
void read_parameters (gcry_mpi_t* n, gcry_mpi_t* a, unsigned long long* times, gcry_mpi_t* ckey);
void lock (unsigned long);
void unlock (int);
void info ();
int create_puzzle (gcry_mpi_t, unsigned long long, gcry_mpi_t, gcry_mpi_t, gcry_mpi_t);
void solve_puzzle (gcry_mpi_t, gcry_mpi_t, gcry_mpi_t, unsigned long long, gcry_mpi_t, int);
double squarings_per_second ();

const char* argp_program_version = PACKAGE_STRING;
const char* argp_program_bug_address = PACKAGE_BUGREPORT;
static char doc[] = "timelock -- encrypt files which can be opened only in the future";
static char args_doc[] = "[lock|unlock]";

static struct argp_option options[] = {
  {"time", 't', "SECONDS", 0, "Time (in seconds) for which the input should be locked", 0},
  {"quiet", 'q', 0, 0, "Don't produce any output", 0},
  {"help", 'h', 0, 0, "Give this help list", 0},
  {"usage", 1, 0, OPTION_HIDDEN, "Give a short usage message", 0},
  { 0 }
};

struct arguments {
  enum {LOCK, UNLOCK, INFO} command;
  unsigned long seconds;
  int quiet;
  char* input_file;
};

static error_t parse_opt (int key, char *arg, struct argp_state *state)
{
  struct arguments *arguments = state->input;
  switch (key) {
  case 't':
    sscanf(arg, "%lu", &arguments->seconds);
    break;
  case 'i':
    arguments->input_file = arg;
    break;
  case 'q':
    arguments->quiet = 1;
    break;
  case 'h':
    argp_state_help(state, stderr, ARGP_HELP_STD_HELP);
    break;
  case 1:
    argp_state_help(state, stderr, ARGP_HELP_STD_USAGE);
    break;
  case ARGP_KEY_ARG:
    if (state->arg_num >= 1)
      argp_usage(state);
    if (!strcmp(arg, "lock"))
      arguments->command = LOCK;
    else if (!strcmp(arg, "unlock"))
      arguments->command = UNLOCK;
    else {
      fprintf(stderr, "%s: Invalid command\n", arg);
      argp_usage(state);
      exit(1);
    }
    break;
  case ARGP_KEY_END:
    if (((arguments->command == UNLOCK) || (arguments->command == INFO)) && (arguments->seconds != 0)) {
      fprintf(stderr, "-t|--times: Option only valid while locking\n");
      argp_usage(state);
      exit(1);
    }
    break;
  default:
    return ARGP_ERR_UNKNOWN;
  }
  return 0;
}

static struct argp argp = {options, parse_opt, args_doc, doc, NULL, NULL, NULL};

int main (int argc, char *argv[])
{
  struct arguments arguments;
  arguments.command = INFO;
  arguments.seconds = 0;
  arguments.input_file = "-";
  arguments.quiet = 0;
  argp_parse(&argp, argc, argv, ARGP_NO_HELP, 0, &arguments);

  if (!gcry_check_version (GCRYPT_VERSION))
    error(2, 0, "libgcrypt version mismatch\n");
  gcry_control(GCRYCTL_SUSPEND_SECMEM_WARN);
  gcry_control(GCRYCTL_INIT_SECMEM, 16384, 0);
  gcry_control(GCRYCTL_RESUME_SECMEM_WARN);
  gcry_control(GCRYCTL_INITIALIZATION_FINISHED, 0);

  if (arguments.command == LOCK)
    lock(arguments.seconds == 0 ? DEFAULT_TIME : arguments.seconds);
  else if (arguments.command == UNLOCK) unlock(arguments.quiet != 1);
  else if (arguments.command == INFO) info();
  return 0;
}

size_t read_stdin (unsigned char **buffer)
{
  // Allocate memory for read buffer
  SAFE_MALLOC(*buffer, READ_BUFFER_SIZE, return 0);
  size_t bytes_read = 0;

  // Read READ_BUFFER_SIZE bytes at a time
  for (size_t size=READ_BUFFER_SIZE; ; size += READ_BUFFER_SIZE) {
    size_t len = fread(*buffer + size - READ_BUFFER_SIZE, 1, READ_BUFFER_SIZE, stdin);
    bytes_read += len;
    if (len < READ_BUFFER_SIZE)
      if (ferror(stdin)) {
	perror(__func__);
	CLEANUP((*buffer));
	return 0;
      }
      else return bytes_read;
    else SAFE_REALLOC(*buffer, size + READ_BUFFER_SIZE, return 0);
  }
}

gcry_error_t generate_prime (gcry_mpi_t* p)
{
  return gcry_prime_generate(p, PRIME_BITS, 0, NULL, NULL, NULL, GCRY_WEAK_RANDOM, 0);
}

void print_parameter (enum chunk_type type, gcry_mpi_t param)
{
  unsigned char* buffer;
  size_t nwritten;
  gcry_mpi_aprint(GCRYMPI_FMT_PGP, &buffer, &nwritten, param);
  SAFE_FWRITE(&type, 1, );
  SAFE_FWRITE(buffer, nwritten, );
  CLEANUP((buffer));
}

gcry_mpi_t mpi_set_ulli (gcry_mpi_t w, unsigned long long u)
{
  unsigned long long u_be = htobe64(u);
  gcry_mpi_scan(&w, GCRYMPI_FMT_USG, &u_be, sizeof(u_be), NULL);
  return w;
}

unsigned long long mpi_to_ulli (gcry_mpi_t w)
{
  unsigned long long u;
  size_t nwritten;
  gcry_mpi_print(GCRYMPI_FMT_USG, (unsigned char*)&u, sizeof(u), &nwritten, w);
  u = be64toh(u) >> (sizeof(u) - nwritten)*8;
  return u;
}

void lock (const unsigned long seconds)
{
  unsigned char key_bytes[KEY_SIZE];
  WITH_MPI
    ((n)(a)(times_mpi)(ckey)(key_mpi),
     unsigned long long times = seconds * squarings_per_second();

     // Create and output puzzle
     if (create_puzzle(key_mpi, times, n, a, ckey)) {
       CLEANUP((n)(a)(ckey)(key_mpi));
       exit(1);
     }
     printf("%s%c", MAGIC_BYTES, OUTPUT_VERSION);
     print_parameter(CHUNK_N, n);
     print_parameter(CHUNK_A, a);
     times_mpi = mpi_set_ulli(times_mpi, times);
     print_parameter(CHUNK_TIMES, times_mpi);
     print_parameter(CHUNK_CKEY, ckey);
     printf("%c", END_OF_CHUNKS);
     TRY_CATCH_GCRYPT(gcry_mpi_print(GCRYMPI_FMT_USG, key_bytes, KEY_SIZE, NULL, key_mpi),
		      "Convert MPI key to bytes", CLEANUP((n)(a)(times_mpi)(ckey)(key_mpi))););

  // Read plaintext of message
  unsigned char* buffer;
  size_t len = read_stdin(&buffer);
  if (!len) exit(1);

  // Generate and print nonce
  unsigned char nonce[NONCE_SIZE];
  gcry_create_nonce(nonce, NONCE_SIZE);
  printf("%c", SALSA20);
  SAFE_FWRITE(nonce, NONCE_SIZE, CLEANUP((buffer)));

  // Write ciphertext of message
  gcry_cipher_hd_t handle;
  TRY_CATCH_GCRYPT(open_cipher(&handle, key_bytes, nonce),
		   "Open cipher handle", CLEANUP((handle)(buffer)));
  TRY_CATCH_GCRYPT(gcry_cipher_encrypt(handle, buffer, len, NULL, 0),
		   "Encrypt message", CLEANUP((handle)(buffer)));
  SAFE_FWRITE(buffer, len, CLEANUP((handle)(buffer)));

  CLEANUP((handle)(buffer));
}

double squarings_per_second ()
{
  double squarings = 0;
  WITH_MPI
    ((n)(a)(ckey)(key),
     if (create_puzzle(key, TEST_ITERATIONS, n, a, ckey)) {
       CLEANUP((n)(a)(ckey)(key));
       exit(1);
     }
     clock_t start = clock();
     solve_puzzle(key, n, a, TEST_ITERATIONS, ckey, 0);
     clock_t stop = clock();
     squarings = CLOCKS_PER_SEC * TEST_ITERATIONS / ((double)(stop - start)););
  return squarings;
}

void unlock (int progress)
{
  unsigned char key_bytes[KEY_SIZE];
  unsigned long long times;
  WITH_MPI
    ((n)(a)(ckey)(key),
     read_parameters(&n, &a, &times, &ckey);
     solve_puzzle(key, n, a, times, ckey, progress);
     TRY_CATCH_GCRYPT(gcry_mpi_print(GCRYMPI_FMT_USG, key_bytes, KEY_SIZE, NULL, key),
		      "Convert MPI key to bytes", CLEANUP((n)(a)(ckey)(key))););

  // Read nonce
  enum cipher_algo cipher;
  SAFE_FREAD(&cipher, 1, );
  if ((uint8_t)cipher != SALSA20) {
    error(1, 0, "Cipher algorithm not supported\n");
    exit(1);
  }
  unsigned char nonce[NONCE_SIZE];
  SAFE_FREAD(nonce, NONCE_SIZE, );

  // Open cipher handle
  gcry_cipher_hd_t handle;
  TRY_CATCH_GCRYPT(open_cipher(&handle, key_bytes, nonce),
		   "Open cipher handle", CLEANUP((handle)));

  // Read stdin
  unsigned char* buffer;
  size_t len = read_stdin(&buffer);
  if (!len) CLEANUP((handle));

  // Decrypt
  TRY_CATCH_GCRYPT(gcry_cipher_decrypt(handle, buffer, len, NULL, 0),
		   "Decrypt message", CLEANUP((handle)(buffer)));
  SAFE_FWRITE(buffer, len, );
  CLEANUP((handle)(buffer));
}

int create_puzzle (const gcry_mpi_t key, const unsigned long long _times, gcry_mpi_t n, gcry_mpi_t a, gcry_mpi_t ckey)
{
  ALLOCATE_SECURE_MPI((p)(q)(phi)(two)(times));
  times = mpi_set_ulli(times, _times);
  two = mpi_set_ulli(two, 2);
  int rv = 1;

  // Primes p, q; n, phi and a
  TRY_CATCH_GCRYPT((generate_prime(&p) || generate_prime(&q)), "Generate primes", goto cleanup);
  gcry_mpi_mul(n, p, q);
  gcry_mpi_sub_ui(p, p, 1);
  gcry_mpi_sub_ui(q, q, 1);
  gcry_mpi_mul(phi, p, q);
  while (!((gcry_mpi_cmp_ui(a, 1) > 0) && (gcry_mpi_cmp(a, n) < 0)))
    gcry_mpi_randomize(a, gcry_mpi_get_nbits(n), GCRY_WEAK_RANDOM);

  // ciphertext of key
  gcry_mpi_powm(ckey, two, times, phi);
  gcry_mpi_powm(ckey, a, ckey, n);
  while (!((gcry_mpi_cmp_ui(key, 1) > 0) && (gcry_mpi_cmp(key, n) < 0)))
    gcry_mpi_randomize(key, KEY_BITS, GCRY_WEAK_RANDOM);
  gcry_mpi_addm(ckey, key, ckey, n);
  rv = 0;

 cleanup:
  CLEANUP((p)(q)(phi)(two)(times));
  return rv;
}

void solve_puzzle (gcry_mpi_t key, const gcry_mpi_t n, const gcry_mpi_t a, const unsigned long long times, const gcry_mpi_t ckey, int progress)
{
  // Compute number of squarings
  unsigned long long full = times / SQUARINGS_PER_ITERATION;
  unsigned long long partial = times % SQUARINGS_PER_ITERATION;

  // Perform squarings
  gcry_mpi_set(key, a);
  WITH_MPI
    ((one)(exponent),
     one = mpi_set_ulli(one, 1);
     gcry_mpi_mul_2exp(exponent, one, SQUARINGS_PER_ITERATION);
     for (unsigned long long i=0; i<full; i++) {
       if (progress) fprintf(stderr, "\rProgress: %.2f%%", (100.0*i)/full);
       gcry_mpi_powm(key, key, exponent, n);
     }
     gcry_mpi_mul_2exp(exponent, one, partial);
     gcry_mpi_powm(key, key, exponent, n);
     if (progress) fprintf(stderr, "\rProgress: 100.00%%\n");
     gcry_mpi_subm(key, ckey, key, n););
}

gcry_error_t open_cipher (gcry_cipher_hd_t* handle, const void *key, const unsigned char *nonce)
{
  gcry_error_t err;
  if (gcry_cipher_open(handle, GCRY_CIPHER_SALSA20, GCRY_CIPHER_MODE_STREAM, 0)) return err;
  if (gcry_cipher_setkey(*handle, key, KEY_SIZE)) return err;
  if (gcry_cipher_setiv(*handle, nonce, NONCE_SIZE)) return err;
  return 0;
}

int handle_gcrypt_error (gcry_error_t err, const char* function)
{
  if (err) fprintf(stderr, "%s: %s/%s\n", function, gcry_strsource(err), gcry_strerror(err));
  return err ? 1 : 0;
}

void info ()
{
  unsigned long long times;
  double sps = squarings_per_second();
  if (sps == 0) exit(1);
  WITH_MPI
    ((n)(a)(ckey),
     read_parameters(&n, &a, &times, &ckey);
     printf("n: %u bits\n", gcry_mpi_get_nbits(n));
     printf("a: %u bits\n", gcry_mpi_get_nbits(a));
     printf("squarings: %llu [ETA: %.0f second(s)]\n", times, round(times/sps));
     printf("ckey: %u bits\n", gcry_mpi_get_nbits(ckey)););
}

void read_parameters (gcry_mpi_t* n, gcry_mpi_t* a, unsigned long long* times, gcry_mpi_t* ckey)
{
  unsigned char* buffer;
  SAFE_MALLOC(buffer, strlen(MAGIC_BYTES) + 1, exit(1));
  ALLOCATE_MPI((times_mpi));

  // Check for magic bytes
  buffer[strlen(MAGIC_BYTES)] = 0;
  SAFE_FREAD(buffer, strlen(MAGIC_BYTES), );
  if (!strstr((char*)buffer, MAGIC_BYTES))
    error(1, 0, "Magic bytes missing. Input data is invalid.\n");

  // Check for output version
  uint8_t version = fgetc(stdin);
  if (version > OUTPUT_VERSION)
    error(1, 0, "Input data format version %d not supported. Please upgrade %s.\n", version, PACKAGE);

  // Read parameters
  enum chunk_type type;
  gcry_mpi_t* params[] = {NULL, n, a, &times_mpi, ckey};
  SAFE_REALLOC(buffer, 2, exit(1));
  while ((type = fgetc(stdin))) {
    SAFE_FREAD(buffer, 2, );
    size_t len = be16toh(*(uint16_t*)buffer);
    len = (len / 8) + (len % 8 ? 1 : 0);
    SAFE_REALLOC(buffer, len + 2, exit(1));
    SAFE_FREAD(buffer + 2, len, );
    TRY_CATCH_GCRYPT(gcry_mpi_scan(params[type], GCRYMPI_FMT_PGP, buffer, len + 2, NULL),
		     "Read parameter", exit(1));
  }
  *times = mpi_to_ulli(times_mpi);
  CLEANUP((times_mpi)(buffer));
}