/*
 * Fast multithreaded pseudorandom number generator
 *
 * Copyright (c) 2012 Michael Buesch <m@bues.ch>
 * Licensed under the GNU/GPL version 2 or (at your option) any later version.
 */

#include "util.h"

#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <stddef.h>
#include <getopt.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <pthread.h>
#include <signal.h>
#include <assert.h>
#include <sys/sysinfo.h>
#include <gsl/gsl_rng.h>
#include <openssl/sha.h>


/* RNG state is re-seeded from the entropy source on each round */
struct rng_state {
	uint64_t hashsalt;
	uint32_t rngseed;
	uint32_t bytes_left;
};

#define MAX_RESEED_INTERVAL	0x1FFFFu /* Must be a power of two minus one */
#define CHUNK_SIZE		8192u
#define OUT_CHUNK_SIZE		4000u

struct rng_context {
	int source_fd;
	gsl_rng *gslrng;
	unsigned long rngmin, rngmax;
	void (*rng_make_chunk)(struct rng_context *ctx,
			       uint8_t *chunk, size_t size);
	size_t (*rng_hash_chunk)(struct rng_context *ctx,
				 uint8_t *chunk, size_t size);
	struct rng_state state;
};

struct one_thread {
	pthread_t tid;
	sig_atomic_t stop;
	struct rng_context rng;
};

static struct threads_context {
	struct one_thread *t;
	unsigned int count;
} threads;

enum hash_type {
	HASH_NONE,
	HASH_SHA512,
};

static struct commandline_options {
	unsigned int nr_threads;
	const gsl_rng_type *rngtype;
	enum hash_type hashtype;
	bool nofold;
	const char *sourcefile;
} cmdopts;

static sig_atomic_t kill_signal_received;


#define DEFINE_RNG_MAKE_CHUNK(_nr_bytes)			\
static void rng_make_chunk_##_nr_bytes(struct rng_context *ctx,	\
				       uint8_t *chunk,		\
				       size_t size)		\
{								\
	unsigned long i, data;					\
								\
	while (size) {						\
		data = gsl_rng_get(ctx->gslrng);		\
		data -= ctx->rngmin;				\
		for (i = 0; size && i < _nr_bytes; i++) {	\
			*chunk = data;				\
			data >>= 8;				\
			chunk++;				\
			size--;					\
		}						\
	}							\
}
DEFINE_RNG_MAKE_CHUNK(1);
DEFINE_RNG_MAKE_CHUNK(2);
DEFINE_RNG_MAKE_CHUNK(3);
DEFINE_RNG_MAKE_CHUNK(4);

static size_t rng_hash_chunk_none(struct rng_context *ctx,
				  uint8_t *chunk, size_t size)
{
	return size;
}

#define DEFINE_RNG_HASH_CHUNK(_hashfunc, _hashctx, _digestbits)			\
static size_t rng_hash_chunk_##_hashfunc(struct rng_context *ctx,		\
					 uint8_t *chunk, size_t size)		\
{										\
	const size_t diglen = (_digestbits) / 8;				\
	uint8_t dig[diglen];							\
	size_t i, j, len, outidx = 0;						\
	_hashctx h;								\
										\
	for (i = 0; i < size; i += len) {					\
		len = min(diglen, size - i);					\
		_hashfunc##_Init(&h);						\
		_hashfunc##_Update(&h, &ctx->state.hashsalt,			\
				   sizeof(ctx->state.hashsalt));		\
		_hashfunc##_Update(&h, chunk + i, len);				\
		_hashfunc##_Final(dig, &h);					\
		if (cmdopts.nofold) {						\
			memcpy(chunk + outidx, dig, len);			\
			outidx += len;						\
		} else {							\
			for (j = 0; j < len / 2; j++) {				\
				chunk[outidx + j] =				\
					dig[j] ^ dig[diglen / 2 + j];		\
			}							\
			outidx += len / 2;					\
		}								\
	}									\
										\
	return outidx;								\
}
DEFINE_RNG_HASH_CHUNK(SHA512, SHA512_CTX, 512);

static int read_entropy_source(int fd, void *_buf, size_t size)
{
	uint8_t *buf = _buf;
	size_t r = 0;
	ssize_t res;

	while (r < size) {
		if (kill_signal_received)
			return -EINTR;
		res = read(fd, buf + r, size - r);
		if (res < 0) {
			if (errno == EAGAIN)
				continue;
			if (errno == EWOULDBLOCK)
				continue;
			if (errno == EINTR)
				continue;
			fprintf(stderr, "Failed to read %s: %s\n",
				cmdopts.sourcefile, strerror(errno));
			return errno;
		}
		r += res;
	}

	return 0;
}

static int write_output(int fd, const void *_buf, size_t size)
{
	size_t written;
	ssize_t res;
	const uint8_t *buf = _buf;

	for (written = 0; written < size; ) {
		if (kill_signal_received)
			return -EINTR;
		res = write(STDOUT_FILENO,
			    buf + written,
			    min(OUT_CHUNK_SIZE, size - written));
		if (res < 0) {
			if (errno == EAGAIN)
				continue;
			if (errno == EWOULDBLOCK)
				continue;
			if (errno == EINTR)
				continue;
			return errno;
		}
		written += res;
	}
	assert(written == size);

	return 0;
}

#define RANDOM_EXTRACT(to, from_ptr)			\
	do {						\
		(to) = *((typeof(to) *)(from_ptr));	\
		(from_ptr) += sizeof(to);		\
	} while (0)

static void random_gen(struct rng_context *rng)
{
	int err;
	uint8_t chunk[CHUNK_SIZE];
	unsigned int chunksize;

	/* Re-seed? */
	while (!rng->state.bytes_left) {
		uint8_t random[sizeof(rng->state)];
		uint8_t *random_p = random;
		size_t len;

		len = sizeof(rng->state.rngseed) +
		      sizeof(rng->state.bytes_left);
		if (rng->rng_hash_chunk != rng_hash_chunk_none)
			len += sizeof(rng->state.hashsalt);
		err = read_entropy_source(rng->source_fd, random, len);
		if (err)
			return;
		RANDOM_EXTRACT(rng->state.rngseed, random_p);
		RANDOM_EXTRACT(rng->state.bytes_left, random_p);
		if (rng->rng_hash_chunk != rng_hash_chunk_none)
			RANDOM_EXTRACT(rng->state.hashsalt, random_p);

		if (rng->state.rngseed == 0) /* seed 0 is special. Don't use it. */
			continue;
		gsl_rng_set(rng->gslrng, rng->state.rngseed);
		rng->state.bytes_left &= MAX_RESEED_INTERVAL;

#if 0
		fprintf(stderr, "Seed values: "
			"rngseed=%08X, bytes_left=%08X, hashsalt=%08X%08X\n",
			rng->state.rngseed, rng->state.bytes_left,
			(unsigned int)((rng->state.hashsalt >> 32) & 0xFFFFFFFF),
			(unsigned int)(rng->state.hashsalt & 0xFFFFFFFF));
#endif
	}

	/* Determine current chunk size */
	chunksize = min(sizeof(chunk), rng->state.bytes_left);
	rng->state.bytes_left -= chunksize;

	/* Generate the random bits */
	rng->rng_make_chunk(rng, chunk, chunksize);

	/* Hash the random bits */
	chunksize = rng->rng_hash_chunk(rng, chunk, chunksize);

	/* Write the final data to the output fd. */
	err = write_output(STDOUT_FILENO, chunk, chunksize);
	if (err > 0) {
		fprintf(stderr, "Failed to dump chunk\n");
		exit(1);
	}
}

static void * thread_function(void *data)
{
	struct one_thread *t = data;

	while (!t->stop)
		random_gen(&t->rng);

	return NULL;
}

static void stop_threads(void)
{
	struct one_thread *t;
	unsigned int i;

	for (i = 0; i < threads.count; i++) {
		t = &threads.t[i];

		if (!t->tid)
			continue;

		t->stop = 1;
		pthread_join(t->tid, NULL);

		gsl_rng_free(t->rng.gslrng);
		if (t->rng.source_fd != STDIN_FILENO)
			close(t->rng.source_fd);

		memset(t, 0, sizeof(*t));
	}
	free(threads.t);
	memset(&threads, 0, sizeof(threads));
}

static int start_threads(unsigned int count)
{
	struct one_thread *t;
	int err;
	unsigned int i;
	unsigned long rngrange;

	memset(&threads, 0, sizeof(threads));
	threads.t = calloc(count, sizeof(*threads.t));
	if (!threads.t) {
		fprintf(stderr, "Out of memory\n");
		return -ENOMEM;
	}
	threads.count = count;

	for (i = 0; i < threads.count; i++) {
		t = &threads.t[i];

		if (strcmp(cmdopts.sourcefile, "-") == 0) {
			t->rng.source_fd = STDIN_FILENO;
		} else {
			t->rng.source_fd = open(cmdopts.sourcefile, O_RDONLY);
			if (t->rng.source_fd < 0) {
				fprintf(stderr, "Failed to open %s\n", cmdopts.sourcefile);
				goto error;
			}
		}
		t->rng.gslrng = gsl_rng_alloc(cmdopts.rngtype);
		if (!t->rng.gslrng) {
			fprintf(stderr, "Failed to init GSL RNG\n");
			t->tid = 0;
			if (t->rng.source_fd != STDIN_FILENO)
				close(t->rng.source_fd);
			goto error;
		}
		t->rng.rngmin = gsl_rng_min(t->rng.gslrng);
		t->rng.rngmax = gsl_rng_max(t->rng.gslrng);
		rngrange = t->rng.rngmax - t->rng.rngmin;
		if (rngrange >= 0xFFFFFFFF) {
			t->rng.rng_make_chunk = rng_make_chunk_4;
		} else if (rngrange >= 0xFFFFFF) {
			t->rng.rng_make_chunk = rng_make_chunk_3;
		} else if (rngrange >= 0xFFFF) {
			t->rng.rng_make_chunk = rng_make_chunk_2;
		} else if (rngrange >= 0xFF) {
			t->rng.rng_make_chunk = rng_make_chunk_1;
		} else {
			fprintf(stderr, "Weird RNG! Could not determine chunk function\n");
			t->tid = 0;
			gsl_rng_free(t->rng.gslrng);
			if (t->rng.source_fd != STDIN_FILENO)
				close(t->rng.source_fd);
			goto error;
		}
		switch (cmdopts.hashtype) {
		case HASH_SHA512:
			t->rng.rng_hash_chunk = rng_hash_chunk_SHA512;
			break;
		case HASH_NONE:
			t->rng.rng_hash_chunk = rng_hash_chunk_none;
			break;
		}

		err = pthread_create(&t->tid, NULL, thread_function, t);
		if (err) {
			t->tid = 0;
			gsl_rng_free(t->rng.gslrng);
			if (t->rng.source_fd != STDIN_FILENO)
				close(t->rng.source_fd);
			goto error;
		}
	}

	return 0;
error:
	stop_threads();
	return -ENODEV;
}

static unsigned int get_nr_system_cpus(void)
{
	return sysconf(_SC_NPROCESSORS_ONLN);
}

static void usage(FILE *fd, int argc, char **argv)
{
#undef P
#define P(fmt, ...)	fprintf(fd, fmt "\n", ##__VA_ARGS__)

	P("Fast multithreaded pseudorandom number generator");
	P("");
	P("Usage: %s [OPTIONS]", argv[0]);
	P("");
	P("  -t|--threads COUNT        Set the number of threads.");
	P("                            Defaults to the number of processors in the system.");
	P("  -r|--rng NAME             Set the random number generator.");
	P("                            See \"cumulative options\" for defaults.");
	P("  -H|--hash NAME            Set the hash.");
	P("                            See \"cumulative options\" for defaults.");
	P("  -F|--nofold               Don't fold the hash in half.");
	P("                            In --secure and --insane mode, the hash is folded in half.");
	P("                            Folding halves speed, but may also improve security.");
	P("  -s|--source FILE          Entropy source for seeding the RNG.");
	P("                            Default: /dev/urandom (unless --insane).");
	P("                            Use - for stdin.");
	P("  -h|--help                 Print this help text.");
	P("");
	P("Cumulative options:");
	P("  -3|--insane               Same as --secure, but uses /dev/random instead");
	P("                            of /dev/urandom as --source.");
	P("  -2|--secure     (default) Generate secure random bitstream for use in cryptographic");
	P("                            applications.");
	P("                            Sets: -r taus2 -H SHA512");
	P("  -1|--fair                 Generate random bitstream good enough for use in most");
	P("                            cryptographic applications.");
	P("                            Sets: -r taus2 -H SHA512 -F");
	P("  -0|--fast                 Generate random bitstream _NOT_ for use in cryptographic");
	P("                            applications.");
	P("                            Sets: -r taus2 -H none");
	P("");
	P("Valid RNGs are:");
	P("  taus2, mt19937, ranlux389, cmrg, mrg, gfsr4");
	P("");
	P("Valid hashes are:");
	P("  none, SHA512");

#undef P
}

static void set_rng_level(int level)
{
	switch (level) {
	case 3: /* --insane */
		cmdopts.sourcefile = "/dev/random";
		cmdopts.hashtype = HASH_SHA512;
		cmdopts.nofold = 0;
		cmdopts.rngtype = gsl_rng_taus2;
		break;
	case 2: /* --secure */
		cmdopts.sourcefile = "/dev/urandom";
		cmdopts.hashtype = HASH_SHA512;
		cmdopts.nofold = 0;
		cmdopts.rngtype = gsl_rng_taus2;
		break;
	case 1: /* --fair */
		cmdopts.sourcefile = "/dev/urandom";
		cmdopts.hashtype = HASH_SHA512;
		cmdopts.nofold = 1;
		cmdopts.rngtype = gsl_rng_taus2;
		break;
	case 0: /* --fast */
		cmdopts.sourcefile = "/dev/urandom";
		cmdopts.hashtype = HASH_NONE;
		cmdopts.nofold = 1;
		cmdopts.rngtype = gsl_rng_taus2;
		break;
	default:
		fprintf(stderr, "Internal error: Invalid rng level %d\n",
			level);
		exit(1);
	}
}

static int parse_args(int argc, char **argv)
{
	static struct option long_options[] = {
		{ "threads", required_argument, 0, 't', },
		{ "rng", required_argument, 0, 'r', },
		{ "hash", required_argument, 0, 'H', },
		{ "nofold", no_argument, 0, 'F', },
		{ "source", required_argument, 0, 's', },
		{ "help", no_argument, 0, 'h', },
		{ "insane", no_argument, 0, '3', },
		{ "secure", no_argument, 0, '2', },
		{ "fair", no_argument, 0, '1', },
		{ "fast", no_argument, 0, '0', },
		{ 0, },
	};
	int c, idx;

	/* Set defaults */
	cmdopts.nr_threads = max(2u, get_nr_system_cpus());
	set_rng_level(2);

	while (1) {
		c = getopt_long(argc, argv, "ht:r:H:Fs:3210",
				long_options, &idx);
		if (c == -1)
			break;
		switch (c) {
		case 't': /* --threads */
			if (sscanf(optarg, "%u", &cmdopts.nr_threads) != 1) {
				fprintf(stderr, "Failed to parse --threads argument\n");
				return -1;
			}
			if (cmdopts.nr_threads < 1 || cmdopts.nr_threads > 4096) {
				fprintf(stderr, "Invalid --threads value\n");
				return -1;
			}
			break;
		case 'r': /* --rng */
			if (strcasecmp(optarg, "taus2") == 0)
				cmdopts.rngtype = gsl_rng_taus2;
			else if (strcasecmp(optarg, "mt19937") == 0)
				cmdopts.rngtype = gsl_rng_mt19937;
			else if (strcasecmp(optarg, "ranlux389") == 0)
				cmdopts.rngtype = gsl_rng_ranlux389;
			else if (strcasecmp(optarg, "cmrg") == 0)
				cmdopts.rngtype = gsl_rng_cmrg;
			else if (strcasecmp(optarg, "mrg") == 0)
				cmdopts.rngtype = gsl_rng_mrg;
			else if (strcasecmp(optarg, "gfsr4") == 0)
				cmdopts.rngtype = gsl_rng_gfsr4;
			else {
				fprintf(stderr, "Invalid --rng name\n");
				return -1;
			}
			break;
		case 'H': /* --hash */
			if (strcasecmp(optarg, "sha512") == 0)
				cmdopts.hashtype = HASH_SHA512;
			else if (strcasecmp(optarg, "none") == 0)
				cmdopts.hashtype = HASH_NONE;
			else {
				fprintf(stderr, "Invalid --hash name\n");
				return -1;
			}
			break;
		case 'F': /* --nofold */
			cmdopts.nofold = 1;
			break;
		case 's': /* --source */
			cmdopts.sourcefile = optarg;
			break;
		case 'h': /* --help */
			usage(stdout, argc, argv);
			return 1;
		case '3': /* --insane */
			set_rng_level(3);
			break;
		case '2': /* --secure */
			set_rng_level(2);
			break;
		case '1': /* --fair */
			set_rng_level(1);
			break;
		case '0': /* --fast */
			set_rng_level(0);
			break;
		default:
			return -1;
		}
	}

	return 0;
}

static void signal_handler(int signum)
{
	switch (signum) {
	case SIGINT:
	case SIGTERM:
		fprintf(stderr, "Terminating signal received.\n");
		kill_signal_received = 1;
		break;
	default:
		fprintf(stderr, "Received unknown signal %d\n", signum);
	}
}

static void setup_sighandler(void)
{
	struct sigaction act;

	memset(&act, 0, sizeof(act));
	sigemptyset(&act.sa_mask);
	act.sa_flags = 0;
	act.sa_handler = signal_handler;

	sigaction(SIGINT, &act, NULL);
	sigaction(SIGTERM, &act, NULL);
}

int main(int argc, char **argv)
{
	int err;

	err = parse_args(argc, argv);
	if (err > 0)
		return 0;
	if (err)
		return err;

	if (isatty(STDOUT_FILENO)) {
		fprintf(stderr, "stdout is a tty. Not dumping random data to a terminal. "
			"Aborting...\n");
		return 1;
	}

	setup_sighandler();
	err = start_threads(cmdopts.nr_threads);
	if (err)
		return 1;

	while (!kill_signal_received)
		sleep(1);

	stop_threads();

	return 0;
}